From 3a0851b9213f0280e6c2e68987177d5be6958c43 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 7 Jan 2025 15:18:13 -0800
Subject: [PATCH] Add counter for graph conversion in V1 compat pipeline.

PiperOrigin-RevId: 713058419
---
 .../compiler/mlir/mlir_graph_optimization_pass.cc     | 11 +++++++++++
 .../mlir/mlir_graph_optimization_pass_test.cc         |  7 +++++++
 2 files changed, 18 insertions(+)

diff --git a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
index 55dc00975ad9a2..0f463a5996cb51 100644
--- a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
+++ b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
@@ -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";
@@ -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();
diff --git a/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc b/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
index 6ed719000a6494..c00eba34a93ce7 100644
--- a/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
+++ b/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
@@ -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"
@@ -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 {
@@ -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) {