Add metrics to compilesingleop and compilefunction.

PiperOrigin-RevId: 631523608
tensorflow · May 7, 2024 · 4d3287a · 4d3287a
1 parent 8fad1c4
commit 4d3287a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
@@ -499,6 +499,7 @@ cc_library(
         "//tensorflow/compiler/jit:flags",
         "//tensorflow/compiler/jit:shape_inference",
         "//tensorflow/compiler/jit:xla_compile_util",
+        "//tensorflow/compiler/mlir/tensorflow:attribute_utils",
         "//tensorflow/compiler/mlir/tf2xla:mlir_bridge_rollout_policy",
         "//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes",
         "//tensorflow/compiler/mlir/utils:array_container_utils",

diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/shape_inference.h"
 #include "tensorflow/compiler/jit/xla_compile_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
 #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
 #include "tensorflow/compiler/mlir/utils/array_container_utils.h"
 #include "tensorflow/compiler/tf2xla/graph_compiler.h"
@@ -84,6 +85,12 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+// Name of component for error logging. This name is fixed and required to
+// enable logging.
+constexpr char kSingleOpComponent[] = "TF2XLA_XLA_COMPILER_COMPILE_SINGLE_OP";
+constexpr char kCompileFunctionComponent[] =
+    "TF2XLA_XLA_COMPILER_COMPILE_FUNCTION";
+
 // Checks that arguments `args` match types `types`.
 Status CheckSignature(const DataTypeVector& types,
                       absl::Span<const XlaCompiler::Argument> args) {
@@ -770,6 +777,9 @@ Status XlaCompiler::CompileSingleOp(
     tensorflow::metrics::IncrementPhase2XlaCompilerCounter(
         tensorflow::metrics::Phase2XlaCompilerMetric::
             kCompileSingleOpXlaBuilderFailure);
+    tsl::error_logging::Log(mlir::TF::kBridgeComponent, kSingleOpComponent,
+                            status.ToString())
+        .IgnoreError();
   }
   return status;
 }
@@ -862,12 +872,7 @@ Status XlaCompiler::CompileFunction(
 
   VLOG(1) << "====================================================";
 
-  auto state = ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_DISABLED;
-  if (options.is_entry_computation) {
-    state = GetMlirBridgeRolloutState(config_proto);
-  }
-
-  VLOG(1) << "MLIR bridge off. Using the old bridge to compile the function";
+  VLOG(1) << "CompileFunction with XlaBuilder";
   auto status =
       CompileGraph(options, function_id, std::move(graph), args, result);
   if (!status.ok()) {
@@ -880,6 +885,9 @@ Status XlaCompiler::CompileFunction(
         ". Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and "
         "--vmodule=xla_compiler=2 to obtain a dump of the compiled "
         "functions.");
+    tsl::error_logging::Log(mlir::TF::kBridgeComponent,
+                            kCompileFunctionComponent, status.ToString())
+        .IgnoreError();
     return status;
   }