Update graph analysis (via GetMlirBridgeRolloutPolicy) to take in a flag indicating which version of the Phase 1 pass was called and add missing logging for Session API/V1 Compat Pass.

PiperOrigin-RevId: 467271406

Author: paynecl

tensorflow/compiler/mlir/mlir_bridge_rollout_policy.cc

@@ -23,7 +23,8 @@ MlirBridgeRolloutPolicy GetMlirBridgeRolloutPolicy(
     const tensorflow::Graph& graph,
     const FunctionLibraryDefinition* function_library,
     std::optional<ConfigProto> config_proto,
-    bool uses_uninitialized_resource_args, bool record_stats) {
+    bool uses_uninitialized_resource_args, bool is_v1_compat,
+    bool record_stats) {
   switch (GetMlirBridgeRolloutState(config_proto)) {
     case ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED:
       return MlirBridgeRolloutPolicy::kEnabledByUser;
@@ -36,4 +37,10 @@ MlirBridgeRolloutPolicy GetMlirBridgeRolloutPolicy(
   }
 }
 
+void LogGraphFeatures(const Graph& graph,
+                      const FunctionLibraryDefinition* function_library,
+                      std::optional<ConfigProto> config_proto,
+                      bool uses_uninitialized_resource_args,
+                      bool is_v1_compat) {}
+
 }  // namespace tensorflow

tensorflow/compiler/mlir/mlir_bridge_rollout_policy.h

@@ -56,13 +56,25 @@ MlirBridgeRolloutPolicy GetMlirBridgeRolloutPolicy(
     const tensorflow::Graph& graph,
     const FunctionLibraryDefinition* function_library,
     std::optional<tensorflow::ConfigProto> config_proto,
-    bool uses_uninitialized_resource_args, bool record_stats = false);
+    bool uses_uninitialized_resource_args, bool is_v1_compat,
+    bool record_stats);
 
 static inline MlirBridgeRolloutPolicy GetMlirBridge2ndPhaseRolloutPolicy(
     mlir::ModuleOp module) {
   return MlirBridgeRolloutPolicy::kDisabledAfterGraphAnalysis;
 }
 
+// Explicit Interface for when we want to log features vs test the validity of
+// the graph for MLIR bridge processing.  Note that right now the logging
+// which is done in the logic used by GraphHasFeaturesUnsupportedByMlirBridge
+// has diverged and logs supported features as well.  Parameters are the same
+// as for GetMlirBridgeRolloutPolicy with the exception of
+// record_stats, which isn't needed because this interface will always record.
+void LogGraphFeatures(const Graph& graph,
+                      const FunctionLibraryDefinition* function_library,
+                      std::optional<ConfigProto> config_proto,
+                      bool uses_uninitialized_resource_args, bool is_v1_compat);
+
 }  // namespace tensorflow
 
 #endif  // THIRD_PARTY_TENSORFLOW_COMPILER_MLIR_MLIR_BRIDGE_ROLLOUT_POLICY_H_

tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <memory>
 #include <string>
 
+#include "tensorflow/compiler/mlir/mlir_bridge_rollout_policy.h"
 #include "absl/container/flat_hash_set.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -178,12 +179,7 @@ Status MlirFunctionOptimizationPass::Run(
   tensorflow::metrics::ScopedCounter<2> timings(
       tensorflow::metrics::GetGraphOptimizationCounter(),
       {kTfMlirCategory, "graph_analysis"});
-  // Capture stats on graph properties analyzed before running the MLIR bridge.
-  // We set `uses_uninitialized_resource_args` to false here because function
-  // optimization is not affected by uninitialized resource args.
-  GetMlirBridgeRolloutPolicy(**graph, flib_def, config_proto,
-                             /*uses_uninitialized_resource_args=*/false,
-                             /*record_stats=*/true);
+
   timings.ReportAndStop();
 
   if (overall_state == MlirOptimizationPassState::Disabled) {
@@ -192,6 +188,13 @@ Status MlirFunctionOptimizationPass::Run(
           << "None of the MLIR Optimization Passes are enabled "
           << "(registered " << registry_->passes().size() << ")";
     }
+    // Capture stats on graph properties analyzed before running the MLIR
+    // bridge. We set `uses_uninitialized_resource_args` to false here because
+    // function optimization is not affected by uninitialized resource args.
+    // TODO(b/241853328): Remove LogGraphFeatures when fixed
+    LogGraphFeatures(**graph, flib_def, config_proto,
+                     /*uses_uninitialized_resource_args=*/false,
+                     /*is_v1_compat=*/false);
     return OkStatus();
   }
 
@@ -244,6 +247,14 @@ Status MlirFunctionOptimizationPass::Run(
       std::move(module_ref_status.ValueOrDie());
   AddDevicesToOp(*module_ref, &device_set);
 
+  // Capture stats on graph properties analyzed before running the MLIR
+  // bridge. We set `uses_uninitialized_resource_args` to false here because
+  // function optimization is not affected by uninitialized resource args.
+  // TODO (b/241853328) Remove LogGraphFeatures when fixed
+  LogGraphFeatures(**graph, flib_def, config_proto,
+                   /*uses_uninitialized_resource_args=*/false,
+                   /*is_v1_compat=*/false);
+
   int per_pass_state_index = 0;
   for (auto& pass_registration : registry_->passes()) {
     llvm::StringRef name = pass_registration.pass->name();
@@ -381,11 +392,18 @@ Status MlirV1CompatGraphOptimizationPass::Run(
 
   llvm::StringRef name = pass->name();
   VLOG(2) << "Run MLIR V1 graph optimization pass: " << StringRefToView(name);
+  // If we ever have more than one MlirV1CompatOptimization pass we need to
+  // ensure the logging only happens once per graph to avoid redundant logging
+  // (see how it is used in the MLIRFunctionOptimizationPass as an example)
+  // TODO(b/241853328): Remove LogGraphFeatures when fixed
+  LogGraphFeatures(**options.graph, options.flib_def,
+                   options.session_options->config,
+                   /*uses_uninitialized_resource_args=*/false,
+                   /*is_v1_compat=*/true);
 
   if (VLOG_IS_ON(1)) {
     DumpModule(*module_ref, llvm::formatv("mlir_{0}_before_", name));
   }
-
   Status pass_status = pass->Run(options, *module_ref);
 
   if (!pass_status.ok()) {

tensorflow/compiler/tf2xla/mlir_bridge_pass.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <string>
 
+#include "tensorflow/compiler/mlir/mlir_bridge_rollout_policy.h"
 #include "absl/base/call_once.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"
 #include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h"
@@ -175,9 +176,15 @@ MlirOptimizationPassState MlirBridgePass::GetPassState(
 
   // We set `uses_uninitialized_resource_args` to false here because the first
   // phase of the bridge is not affected by uninitialized resource args.
-  MlirBridgeRolloutPolicy policy =
-      GetMlirBridgeRolloutPolicy(graph, &function_library, config_proto,
-                                 /*uses_uninitialized_resource_args=*/false);
+  // Note we are recording the stats using LogGraphFeatures in the pass
+  // that calls this one to avoid duplicate logging due to
+  // GetMlirBridgeRolloutPolicy being called multiple times for the same graph.
+  // TODO(b/241853328): Add caching of pass state and call logging/metrics
+  // related to graph analysis from here.
+  MlirBridgeRolloutPolicy policy = GetMlirBridgeRolloutPolicy(
+      graph, &function_library, config_proto,
+      /*uses_uninitialized_resource_args=*/false,
+      /*is_v1_compat=*/false, /*record_stats=*/false);
   switch (policy) {
     case MlirBridgeRolloutPolicy::kEnabledByUser:
       return MlirOptimizationPassState::Enabled;
@@ -230,6 +237,8 @@ Status MlirBridgePass::Run(const ConfigProto& config_proto,
 
   // Set device_set to nullptr here as the device specific checks are performed
   // based on the devices in the module.
+  // TODO(b/241853328): Add caching of pass state and call logging/metrics
+  // related to graph analysis from here.
   auto pass_state = GetPassState(/*device_set=*/nullptr, config_proto, graph,
                                  function_library);
 
@@ -265,9 +274,14 @@ MlirOptimizationPassState MlirBridgeV1CompatPass::GetPassState(
   // only run if it's enabled by the user explicitly.
   // We set `uses_uninitialized_resource_args` to false here because the first
   // phase of the bridge is not affected by uninitialized resource args.
+  // Note we are recording the stats using LogGraphFeatures in the pass
+  // that calls this one.
+  // TODO(b/241853328): Add caching of pass state and call logging/metrics
+  // related to graph analysis from here.
   MlirBridgeRolloutPolicy policy = GetMlirBridgeRolloutPolicy(
       graph, /*function_library=*/&function_library, config_proto,
-      /*uses_uninitialized_resource_args=*/false);
+      /*uses_uninitialized_resource_args=*/false, /*is_v1_compat=*/true,
+      /*record_stats=*/false);
   switch (policy) {
     case MlirBridgeRolloutPolicy::kEnabledByUser:
       return MlirOptimizationPassState::Enabled;