From 2406f50d11832c044bc82d995dc395a5c91ae49e Mon Sep 17 00:00:00 2001 From: Feng Liu Date: Tue, 20 Oct 2020 11:08:21 -0700 Subject: [PATCH] Add metrics for the usage of the op composition feature PiperOrigin-RevId: 338095086 Change-Id: I575646b06937cb5a43b3c5a0aa5a438ffe4a4793 --- tensorflow/compiler/mlir/tfr/BUILD | 4 ++++ .../tfr/integration/graph_decompose_pass.cc | 10 ++++++++++ .../tfr/integration/node_expansion_pass.cc | 10 ++++++++++ .../compiler/mlir/tfr/passes/decompose.cc | 20 ++++++++++++++++++- 4 files changed, 43 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/tfr/BUILD b/tensorflow/compiler/mlir/tfr/BUILD index 8a32fc93b7f..ae02faf8be2 100644 --- a/tensorflow/compiler/mlir/tfr/BUILD +++ b/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", diff --git a/tensorflow/compiler/mlir/tfr/integration/graph_decompose_pass.cc b/tensorflow/compiler/mlir/tfr/integration/graph_decompose_pass.cc index 99890e9f621..7041545637a 100644 --- a/tensorflow/compiler/mlir/tfr/integration/graph_decompose_pass.cc +++ b/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)); diff --git a/tensorflow/compiler/mlir/tfr/integration/node_expansion_pass.cc b/tensorflow/compiler/mlir/tfr/integration/node_expansion_pass.cc index 61c4d1c8953..5bb7d235fa7 100644 --- a/tensorflow/compiler/mlir/tfr/integration/node_expansion_pass.cc +++ b/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 diff --git a/tensorflow/compiler/mlir/tfr/passes/decompose.cc b/tensorflow/compiler/mlir/tfr/passes/decompose.cc index 9265437cca9..e1a9ae8c2e6 100644 --- a/tensorflow/compiler/mlir/tfr/passes/decompose.cc +++ b/tensorflow/compiler/mlir/tfr/passes/decompose.cc @@ -19,6 +19,7 @@ limitations under the License. #include #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 { - explicit DecomposeTFOpsPass(llvm::Optional 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();