Skip to content

Commit

Permalink
Add counter for graph conversion in V1 compat pipeline.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 713058419
  • Loading branch information
tensorflower-gardener committed Jan 8, 2025
1 parent c9dec40 commit 3a0851b
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
11 changes: 11 additions & 0 deletions tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,14 @@ auto* mlir_function_pass_graph_conversion_count = monitoring::Counter<1>::New(
"optimization pass",
/* metric field */ "status");

auto* mlir_v1_compat_graph_conversion_count = monitoring::Counter<1>::New(
/* metric name */
"/tensorflow/core/mlir_v1_compat_graph_conversion_count",
/* metric description */
"Track success/failure of Graph to MLIR conversions in MLIR V1 compat "
"optimization pass",
/* metric field */ "status");

// The status metric field is used to record success/failure of mlir
// function/graph optimization passes.
constexpr char kSuccess[] = "kSuccess";
Expand Down Expand Up @@ -434,6 +442,9 @@ absl::Status MlirV1CompatGraphOptimizationPass::Run(
**options.graph, debug_info, *options.flib_def, import_config, &context,
/*tf_name_to_mlir_name*/ nullptr, options.session_options->config,
tensorflow::TF2XLABridgeVersion::kV1Compat);
mlir_v1_compat_graph_conversion_count
->GetCell(absl::StatusCodeToString(module_ref_status.status().code()))
->IncrementBy(1);
if (!module_ref_status.ok()) {
if (pass_state == MlirOptimizationPassState::Enabled) {
return module_ref_status.status();
Expand Down
7 changes: 7 additions & 0 deletions tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ limitations under the License.
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
Expand Down Expand Up @@ -385,6 +386,7 @@ class MlirGraphOptimizationV1PassTest : public Test {
pass_result_expected_[MlirOptimizationPassState::FallbackEnabled]
[false]);
EXPECT_EQ(mlir_function_pass_graph_conversion_count_.Read(kOk), 0);
EXPECT_EQ(mlir_v1_compat_graph_conversion_count_.Read(kOk), 1);
}

void TearDown() override {
Expand Down Expand Up @@ -417,6 +419,11 @@ class MlirGraphOptimizationV1PassTest : public Test {
monitoring::testing::CellReader<int64_t>(
/* metric name */
"/tensorflow/core/mlir_function_pass_graph_conversion_count");
monitoring::testing::CellReader<int64_t>
mlir_v1_compat_graph_conversion_count_ =
monitoring::testing::CellReader<int64_t>(
/* metric name */
"/tensorflow/core/mlir_v1_compat_graph_conversion_count");
};

TEST_F(MlirGraphOptimizationV1PassTest, OptimizationPassDoesNotFailFallback) {
Expand Down

0 comments on commit 3a0851b

Please sign in to comment.