Add metrics for the usage of the op composition feature

PiperOrigin-RevId: 338095086
Change-Id: I575646b06937cb5a43b3c5a0aa5a438ffe4a4793

tensorflow/compiler/mlir/tfr/BUILD

@@ -123,7 +123,9 @@ cc_library(
         ":tfr",
         ":utils",
         "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:lib",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
@@ -228,6 +230,7 @@ cc_library(
     deps = [
         ":tfr_decompose_ctx",
         "//tensorflow/compiler/mlir:mlir_graph_optimization_pass",
+        "//tensorflow/core:lib",
         "//tensorflow/stream_executor/lib",
         "@llvm-project//mlir:IR",
     ],
@@ -257,6 +260,7 @@ cc_library(
     hdrs = ["integration/node_expansion_pass.h"],
     deps = [
         ":tfr_decompose_ctx",
+        "//tensorflow/core:lib",
         "//tensorflow/core/common_runtime/eager:core",
         "//tensorflow/core/common_runtime/eager:eager_op_rewrite_registry",
         "//tensorflow/stream_executor/lib",

tensorflow/compiler/mlir/tfr/integration/graph_decompose_pass.cc

@@ -16,9 +16,17 @@ limitations under the License.
 
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 namespace tensorflow {
+namespace {
+
+auto* tf_core_op_expansion_graph_counter =
+    monitoring::Counter<0>::New("/tensorflow/core/op_expansion/graph_counter",
+                                "The number of graphs being op expanded.");
+}  // namespace
+
 namespace tfr {
 
 bool GraphDecomposePass::IsEnabled(const ConfigProto& config_proto) const {
@@ -34,6 +42,8 @@ Status GraphDecomposePass::Run(const ConfigProto& config_proto,
     return Status::OK();
   }
 
+  tf_core_op_expansion_graph_counter->GetCell()->IncrementBy(1);
+
   LOG_FIRST_N(INFO, 1) << "Run Graph Decomposition Passes";
 
   TF_RETURN_IF_ERROR(DecomposeGraph(module));

tensorflow/compiler/mlir/tfr/integration/node_expansion_pass.cc

@@ -18,9 +18,17 @@ limitations under the License.
 
 #include "absl/strings/str_cat.h"
 #include "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 namespace tensorflow {
+namespace {
+
+auto* tf_core_op_expansion_node_counter =
+    monitoring::Counter<0>::New("/tensorflow/core/op_expansion/node_counter",
+                                "The number of nodes being op expanded.");
+}  // namespace
+
 namespace tfr {
 
 Status CompositeOpExpansion::Run(EagerOperation* orig_op,
@@ -28,6 +36,8 @@ Status CompositeOpExpansion::Run(EagerOperation* orig_op,
   if (!IsEnabled()) return Status::OK();
   if (orig_op->Device() != kVariantDeviceNull) return Status::OK();
 
+  tf_core_op_expansion_node_counter->GetCell()->IncrementBy(1);
+
   LOG_FIRST_N(INFO, 1) << "Run Node Expansion Passes";
 
   // Get the FunctionDef and insert that into the context

tensorflow/compiler/mlir/tfr/passes/decompose.cc

@@ -19,6 +19,7 @@ limitations under the License.
 #include <string>
 
 #include "absl/memory/memory.h"
+#include "absl/strings/string_view.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
@@ -50,6 +51,21 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tfr/ir/tfr_types.h"
 #include "tensorflow/compiler/mlir/tfr/passes/passes.h"
 #include "tensorflow/compiler/mlir/tfr/utils/utils.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+
+namespace tensorflow {
+namespace {
+
+auto* tf_core_op_expansion_op_counter =
+    monitoring::Counter<1>::New("/tensorflow/core/op_expansion/op_counter",
+                                "The number of composite op expanded.", "name");
+}
+
+void IncreaseOpExpansionExecuteCounterByOne(const std::string& op_name) {
+  tf_core_op_expansion_op_counter->GetCell(op_name)->IncrementBy(1);
+}
+
+}  // namespace tensorflow
 
 //===----------------------------------------------------------------------===//
 // The pass to decompose unregistered TF ops with the TFR compose function.
@@ -62,7 +78,6 @@ namespace {
 // Decompose the TF ops with the registered composition library.
 struct DecomposeTFOpsPass
     : public PassWrapper<DecomposeTFOpsPass, FunctionPass> {
-
   explicit DecomposeTFOpsPass(llvm::Optional<ModuleOp> external_tfr_module)
       : external_tfr_module(external_tfr_module) {}
 
@@ -118,6 +133,9 @@ LogicalResult DecomposeTFOpsPass::RewriteUnregisteredTFOps() {
       return;
     }
 
+    tensorflow::IncreaseOpExpansionExecuteCounterByOne(
+        op->getName().getStringRef().str());
+
     auto compose_func_type = compose_func.getType();
     builder.setInsertionPoint(op);
     TFRTensorType unconstrainted_tensor_type = builder.getType<TFRTensorType>();