diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index 6e7d16de16a4f6..eb82e85e83fc2c 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -1154,6 +1154,7 @@ absl::Status EncapsulateSubgraphsPass::Run( // TODO(b/195757077): Remove this once there is a better way to disable // GraphOptimizationPasses that are not needed due to MLIR bridge. + bool has_xla_clusters = false; for (Node* n : (*options.graph)->nodes()) { // Skip the pass if we found TPUExecute or TPUExecuteAndUpdateVariables ops // in the graph, which indicates the graph is produced by TPU TF-XLA bridge @@ -1162,6 +1163,15 @@ absl::Status EncapsulateSubgraphsPass::Run( n->type_string() == "TPUExecuteAndUpdateVariables") { return absl::OkStatus(); } + if (n->attrs().Find(kXlaClusterAttr)) { + has_xla_clusters = true; + } + } + // If no nodes have been marked for XLA compilation, skip the pass to avoid + // the overhead of creating CPU devices and FunctionLibraryRuntime objects. + if (!has_xla_clusters) { + VLOG(1) << "No XLA clusters found, skipping EncapsulateSubgraphsPass"; + return absl::OkStatus(); } std::unique_ptr graph_out; diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc index 94b136a02b99cf..0f925e21fd82f0 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/jit/extract_outside_compilation_pass.h" +#include "tensorflow/compiler/jit/mark_for_compilation_pass.h" #include "tensorflow/compiler/jit/test_util.h" #include "tensorflow/compiler/tf2xla/side_effect_util.h" #include "tensorflow/core/common_runtime/device_factory.h" @@ -2707,6 +2708,10 @@ void CreateSubgraphTouchingRefVar(const Scope& s) { } TEST(EncapsulateSubgraphsTest, RefVariablesMarked) { + // When no XLA clusters are present, EncapsulateSubgraphsPass exits early and + // does not set kXlaHasReferenceVarsAttr, even for graphs with reference + // variables. This is correct since the attribute is only consumed by + // XlaLaunch/XlaCompile nodes which don't exist without clusters. Scope root = Scope::NewRootScope().ExitOnError(); CreateSubgraphTouchingRefVar(root); @@ -2720,12 +2725,14 @@ TEST(EncapsulateSubgraphsTest, RefVariablesMarked) { EncapsulateSubgraphsPass pass; TF_ASSERT_OK(pass.Run(options)); + // With no XLA clusters, the pass exits early and does not set + // kXlaHasReferenceVarsAttr on any nodes. for (const Node* node : graph->nodes()) { - bool has_ref_var; - TF_ASSERT_OK( - GetNodeAttr(node->attrs(), kXlaHasReferenceVarsAttr, &has_ref_var)); - EXPECT_TRUE(node->IsSink() || node->IsSource() || has_ref_var) - << "All nodes apart from source and sink can access reference variable"; + bool has_ref_var = false; + EXPECT_FALSE( + GetNodeAttr(node->attrs(), kXlaHasReferenceVarsAttr, &has_ref_var).ok()) + << "kXlaHasReferenceVarsAttr should not be set when no XLA clusters " + "exist"; } } @@ -2743,6 +2750,14 @@ TEST(EncapsulateSubgraphsTest, NoRefVarsNoAttr) { auto graph = std::make_unique(OpRegistry::Global()); TF_ASSERT_OK(root.ToGraph(graph.get())); + // Add XLA cluster attributes so the pass runs its full path including ref + // variable analysis. Without any kXlaClusterAttr nodes, the pass exits early + // and skips the expensive device/FLR setup and ref var analysis (since those + // attributes are only needed by XlaLaunch/XlaCompile nodes). + for (Node* node : graph->op_nodes()) { + node->AddAttr(kXlaClusterAttr, "cluster_0"); + } + GraphOptimizationPassWrapper wrapper; GraphOptimizationPassOptions options = wrapper.CreateGraphOptimizationPassOptions(&graph); @@ -2758,5 +2773,34 @@ TEST(EncapsulateSubgraphsTest, NoRefVarsNoAttr) { } } +TEST(EncapsulateSubgraphsTest, NoXlaClustersEarlyExit) { + // Verify that EncapsulateSubgraphsPass exits early without doing expensive + // setup work (device creation, FLR creation) when no nodes have been marked + // with kXlaClusterAttr by the MarkForCompilationPass. + Scope root = Scope::NewRootScope().ExitOnError(); + Output constant = + ops::Const(root.WithOpName("constant"), Input::Initializer(1.0)); + Output neg = ops::Negate(root.WithOpName("negate"), constant); + + auto graph = std::make_unique(OpRegistry::Global()); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + GraphOptimizationPassWrapper wrapper; + GraphOptimizationPassOptions options = + wrapper.CreateGraphOptimizationPassOptions(&graph); + + EncapsulateSubgraphsPass pass; + TF_ASSERT_OK(pass.Run(options)); + + // The pass should have exited early without setting kXlaHasReferenceVarsAttr. + for (const Node* node : graph->nodes()) { + bool has_ref_var = false; + EXPECT_FALSE( + GetNodeAttr(node->attrs(), kXlaHasReferenceVarsAttr, &has_ref_var).ok()) + << "kXlaHasReferenceVarsAttr should not be set when no XLA clusters " + "are present (early exit optimization)"; + } +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 340cdbe8032c63..072841770b712a 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -1919,7 +1919,7 @@ absl::Status MarkForCompilation( } } - return MarkForCompilationPassImpl{ + TF_RETURN_IF_ERROR(MarkForCompilationPassImpl{ debug_options, graph, flib_def, @@ -1934,7 +1934,15 @@ absl::Status MarkForCompilation( .session_metadata() .name() : ""} - .Run(); + .Run()); + + // Mark the source node to indicate this graph has been fully processed by + // MarkForCompilation. This allows subsequent invocations (e.g., when + // PartitionedCall re-runs the optimization pipeline) to return early in O(1) + // time by checking the source node, instead of scanning all graph nodes. + graph->source_node()->AddAttr(kXlaAlreadyClustered, true); + + return absl::OkStatus(); } std::atomic* GetPointerToFuel(int64_t initial_value) {