Add tf_xla_null_cluster_outputs debug flag to bypass XLA cluster execution

tensorflow/compiler/jit/flags.cc

@@ -220,6 +220,7 @@ void AllocateAndParseFlags() {
   build_ops_flags->tf_xla_check_cluster_input_numerics = false;
   build_ops_flags->tf_xla_check_cluster_output_numerics = false;
   build_ops_flags->tf_xla_disable_constant_folding = false;
+  build_ops_flags->tf_xla_null_cluster_outputs = false;
   build_ops_flags->tf_xla_disable_full_embedding_pipelining = false;
   build_ops_flags->tf_xla_disable_full_embedding_pipelining_with_summaries =
       true;
@@ -310,6 +311,14 @@ void AllocateAndParseFlags() {
             &build_ops_flags->tf_xla_disable_constant_folding,
             "If true then disables constant folding on TF graph before XLA "
             "compilation."),
+        Flag("tf_xla_null_cluster_outputs",
+             &build_ops_flags->tf_xla_null_cluster_outputs,
+             "If true then XLA clusters are compiled but not executed. "
+             "Outputs are filled with null/zero values: constant outputs "
+             "return their compile-time values, DT_RESOURCE outputs pass "
+             "input resource tensors through unchanged, and all other outputs "
+             "are zero-initialized CPU tensors with the correct type and "
+             "shape. Only the CPU scenario is considered."),
        Flag("tf_xla_disable_full_embedding_pipelining",
             &build_ops_flags->tf_xla_disable_full_embedding_pipelining,
             "If true then disables full embedding pipelining and instead use "

tensorflow/compiler/jit/flags.h

@@ -273,6 +273,15 @@ struct BuildXlaOpsPassFlags {
   // guarantee that tests are run on XLA and not on TF's CPU implementation.
   bool tf_xla_disable_constant_folding;
 
+  // If true, XLA clusters are compiled but not executed. Outputs are filled
+  // with null/zero values instead of actual computation results:
+  //   - Constant outputs still return their compile-time constant values.
+  //   - DT_RESOURCE outputs pass the input resource tensor through unchanged.
+  //   - All other outputs are allocated with the correct type and shape and
+  //     zero-initialized. Only CPU (host) tensors are supported; this flag
+  //     should not be used with GPU or other accelerator devices.
+  bool tf_xla_null_cluster_outputs;
+
   // Disables full embedding pipelining when true. Instead, strict SparseCore
   // TensorCore sequencing will be used.
   bool tf_xla_disable_full_embedding_pipelining;

tensorflow/compiler/jit/kernels/xla_ops.cc

@@ -568,11 +568,19 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
             done);
         OP_REQUIRES_OK_ASYNC(ctx, LockVariables(absl::MakeSpan(variable_infos)),
                              done);
-        OP_REQUIRES_OK_ASYNC(
-            ctx,
-            RunPjRtExecutable(inputs, variable_infos, *compilation_result,
-                              pjrt_client, pjrt_executable, ctx),
-            done);
+        if (GetBuildXlaOpsPassFlags()->tf_xla_null_cluster_outputs) {
+          OP_REQUIRES_OK_ASYNC(
+              ctx,
+              PopulateNullOutputs(ctx, compilation_result,
+                                  /*missing_ctx_input_prefix=*/0),
+              done);
+        } else {
+          OP_REQUIRES_OK_ASYNC(
+              ctx,
+              RunPjRtExecutable(inputs, variable_infos, *compilation_result,
+                                pjrt_client, pjrt_executable, ctx),
+              done);
+        }
       }
       VLOG(2) << "Done executing with PJRT.";
       done();
@@ -640,19 +648,27 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
       xla::RunId run_id(0);
       run_options.set_run_id(run_id);
 
-      absl::StatusOr<xla::ExecutionOutput> execution_output = RunExecutable(
-          platform_info, launch_context, std::move(*execution_inputs),
-          run_options, executable, ctx, allocator.get());
-      OP_REQUIRES_ASYNC(ctx, execution_output.ok(), execution_output.status(),
-                        done);
+      if (GetBuildXlaOpsPassFlags()->tf_xla_null_cluster_outputs) {
+        OP_REQUIRES_OK_ASYNC(
+            ctx,
+            PopulateNullOutputs(ctx, compilation_result,
+                                /*missing_ctx_input_prefix=*/0),
+            done);
+      } else {
+        absl::StatusOr<xla::ExecutionOutput> execution_output = RunExecutable(
+            platform_info, launch_context, std::move(*execution_inputs),
+            run_options, executable, ctx, allocator.get());
+        OP_REQUIRES_ASYNC(ctx, execution_output.ok(), execution_output.status(),
+                          done);
 
-      OP_REQUIRES_OK_ASYNC(
-          ctx,
-          launch_context.PopulateOutputs(
-              ctx, compilation_result, execution_output->ConsumeResult(),
-              /*missing_ctx_input_prefix=*/0, absl::MakeSpan(variable_infos),
-              input_output_alias, resource_var_ptrs),
-          done);
+        OP_REQUIRES_OK_ASYNC(
+            ctx,
+            launch_context.PopulateOutputs(
+                ctx, compilation_result, execution_output->ConsumeResult(),
+                /*missing_ctx_input_prefix=*/0, absl::MakeSpan(variable_infos),
+                input_output_alias, resource_var_ptrs),
+            done);
+      }
       VLOG(1) << "Done";
     }
     done();
@@ -902,11 +918,17 @@ void XlaRunOp::Compute(OpKernelContext* ctx) {
                              closure.num_constant_args());
       OP_REQUIRES_OK(ctx, updated_variables.status());
       OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*updated_variables)));
-      OP_REQUIRES_OK(
-          ctx, RunPjRtExecutable(closure.num_constant_args(), inputs,
-                                 variable_snapshots, *updated_variables,
-                                 *closure.compilation_result(),
-                                 closure.client(), closure.executable(), ctx));
+      if (GetBuildXlaOpsPassFlags()->tf_xla_null_cluster_outputs) {
+        OP_REQUIRES_OK(ctx,
+                       PopulateNullOutputs(ctx, closure.compilation_result(),
+                                           closure.num_constant_args()));
+      } else {
+        OP_REQUIRES_OK(
+            ctx, RunPjRtExecutable(closure.num_constant_args(), inputs,
+                                   variable_snapshots, *updated_variables,
+                                   *closure.compilation_result(),
+                                   closure.client(), closure.executable(), ctx));
+      }
     }
 
     OP_REQUIRES_OK(ctx, absl::OkStatus());
@@ -961,6 +983,13 @@ void XlaRunOp::Compute(OpKernelContext* ctx) {
       GetRecvDeviceMemoryFunction(ctx, key);
   run_options.set_recv_device_memory_function(&recv_function);
 
+  if (GetBuildXlaOpsPassFlags()->tf_xla_null_cluster_outputs) {
+    OP_REQUIRES_OK(ctx,
+                   PopulateNullOutputs(ctx, closure.compilation_result(),
+                                       closure.num_constant_args()));
+    return;
+  }
+
   absl::StatusOr<xla::ExecutionOutput> execution_output = RunExecutable(
       platform_info_, launch_context, std::move(*execution_inputs), run_options,
       closure.executable(), ctx, allocator.get());

tensorflow/compiler/jit/xla_compile_on_demand_op.cc

@@ -270,8 +270,16 @@ void XlaCompileOnDemandOp::Compute(OpKernelContext* ctx) {
     VLOG(2) << "Compiled op with PJRT: " << ctx->status();
     VLOG(2) << "result != nullptr: " << (result != nullptr);
     VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr);
-    VLOG(2) << "Executing with PJRT ...";
 
+    if (GetBuildXlaOpsPassFlags()->tf_xla_null_cluster_outputs) {
+      VLOG(2) << "Skipping PJRT execution (tf_xla_null_cluster_outputs=true).";
+      OP_REQUIRES_OK(ctx,
+                     PopulateNullOutputs(ctx, result,
+                                         /*missing_ctx_input_prefix=*/0));
+      return;
+    }
+
+    VLOG(2) << "Executing with PJRT ...";
     OP_REQUIRES_OK(ctx, RunPjRtExecutable(inputs, variables, *result,
                                           pjrt_device_compiler->client(),
                                           pjrt_executable, ctx));
@@ -301,6 +309,14 @@ void XlaCompileOnDemandOp::Compute(OpKernelContext* ctx) {
     core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
     core::ScopedUnref profiler_ref(profiler);
 
+    if (GetBuildXlaOpsPassFlags()->tf_xla_null_cluster_outputs) {
+      VLOG(2) << "Skipping XLA execution (tf_xla_null_cluster_outputs=true).";
+      OP_REQUIRES_OK(ctx,
+                     PopulateNullOutputs(ctx, result,
+                                         /*missing_ctx_input_prefix=*/0));
+      return;
+    }
+
     // Locks are acquired again when populating the `ctx` outputs.
     OP_REQUIRES_OK(
         ctx, Run(variable_args, result, xla_device_compiler, executable, ctx));

tensorflow/compiler/jit/xla_launch_util.cc

@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/jit/xla_launch_util.h"
 
 #include <cstdint>
+#include <cstring>
 #include <memory>
 #include <optional>
 #include <set>
@@ -320,6 +321,36 @@ absl::Status SetOutputForConstant(
   return absl::OkStatus();
 }
 
+absl::Status PopulateNullOutputs(
+    OpKernelContext* ctx,
+    const XlaCompiler::CompilationResult* compilation_result,
+    int missing_ctx_input_prefix) {
+  TF_RET_CHECK(ctx->num_outputs() ==
+               static_cast<int>(compilation_result->outputs.size()));
+  for (int i = 0; i < ctx->num_outputs(); ++i) {
+    const XlaOutputDescription& descr = compilation_result->outputs[i];
+    if (descr.is_constant) {
+      TF_RETURN_IF_ERROR(SetOutputForConstant(
+          ctx, /*requires_copy_to_device=*/false, compilation_result, i));
+    } else if (descr.type == DT_RESOURCE) {
+      int input_index = descr.input_index - missing_ctx_input_prefix;
+      TF_RET_CHECK(input_index >= 0 && input_index < ctx->num_inputs())
+          << "Invalid input index for null output " << i << ": " << input_index;
+      ctx->set_output(i, ctx->input(input_index));
+    } else {
+      Tensor* output_tensor;
+      TF_RETURN_IF_ERROR(ctx->allocate_output(i, descr.shape, &output_tensor));
+      // DT_STRING tensors store std::string objects, not raw bytes, so they
+      // must not be zero-initialized with memset; their default constructor
+      // already produces valid empty strings.
+      if (descr.type != DT_STRING && output_tensor->NumElements() > 0) {
+        memset(output_tensor->data(), 0, output_tensor->TotalBytes());
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
 static absl::StatusOr<Var*> GetOrCreateResourceVar(
     OpKernelContext* ctx, const ResourceHandle& handle,
     const XlaCompiler::ResourceUpdate& write) {

tensorflow/compiler/jit/xla_launch_util.h

@@ -53,6 +53,20 @@ absl::Status SetOutputForConstant(
     OpKernelContext* ctx, bool requires_copy_to_device,
     const XlaCompiler::CompilationResult* compilation_result, int output_num);
 
+// Populates ctx outputs with null/zero values, skipping actual XLA execution.
+// Used when tf_xla_null_cluster_outputs is enabled for debugging.
+//   - Constant outputs: returns the compiled constant value.
+//   - DT_RESOURCE outputs: passes the corresponding input tensor through.
+//   - All other outputs: allocates a zero-initialized CPU tensor with the
+//     correct type and shape. Only the CPU (host) scenario is supported.
+//
+// `missing_ctx_input_prefix` is the number of leading inputs that are
+// constants baked into the compiled kernel and are absent from `ctx`.
+absl::Status PopulateNullOutputs(
+    OpKernelContext* ctx,
+    const XlaCompiler::CompilationResult* compilation_result,
+    int missing_ctx_input_prefix);
+
 // Converts input tensors and variables which are parameters of the
 // XlaComputation into PjRtBuffers to be fed as input to the
 // PjRtLoadedExecutable.