[pull] master from tensorflow:master

third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc

@@ -17,7 +17,6 @@ limitations under the License.
 #include <memory>
 #include <utility>
 
-#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "xla/tsl/platform/macros.h"
 #include "xla/tsl/platform/test.h"

third_party/xla/xla/pjrt/dump/BUILD

@@ -56,6 +56,7 @@ xla_cc_test(
         ":dump",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/pjrt:pjrt_compiler",
         "//xla/pjrt:pjrt_device_description",
         "//xla/pjrt:pjrt_executable",

third_party/xla/xla/pjrt/dump/dump_test.cc

@@ -45,6 +45,7 @@ limitations under the License.
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/test.h"
+#include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/path.h"
 

third_party/xla/xla/tests/constraint_propagator.cc

@@ -374,10 +374,26 @@ ConstraintPropagator::Run(
         get_index_known_zeroes) {
   ConstraintPropagator propagator(get_index_known_zeroes);
   auto computations = module.MakeComputationPostOrder();
+
+  // Phase 1: Seed constraints and ML patterns across all computations in the
+  // module.
   for (HloComputation* computation : computations) {
-    ABSL_RETURN_IF_ERROR(propagator.Propagate(computation));
+    ABSL_RETURN_IF_ERROR(propagator.SeedConstraints(computation));
+    ABSL_RETURN_IF_ERROR(propagator.SeedMLPatternsConstraints(computation));
+    ABSL_RETURN_IF_ERROR(propagator.PropagateSeedConstraints(computation));
   }
 
+  // Phase 2: Inter-procedural fixed-point propagation across all computations
+  // in the module. Post-order iteration ensures constraints propagate backward
+  // across fusion boundaries in each iteration until convergence.
+  absl::flat_hash_map<const HloInstruction*, ConstraintState> before;
+  do {
+    before = propagator.states_;
+    for (HloComputation* computation : computations) {
+      ABSL_RETURN_IF_ERROR(propagator.PropagateConstraints(computation));
+    }
+  } while (before != propagator.states_);
+
   // Extract only the parameters
   absl::flat_hash_map<const HloInstruction*, ConstraintState> result;
   for (const HloInstruction* param :
@@ -848,6 +864,10 @@ absl::Status ConstraintPropagator::SeedMLPatternsConstraints(
 
 absl::Status ConstraintPropagator::PropagateConstraintsExact(
     const HloInstruction* instruction) {
+  if (instruction->opcode() == HloOpcode::kFusion) {
+    return PropagateFusionBoundary(instruction);
+  }
+
   ConstraintState output_state = states_[instruction];
   ConstraintInterval output_interval = output_state.GetConstraintInterval();
   StructuralConstraints output_structural =
@@ -969,6 +989,46 @@ absl::Status ConstraintPropagator::PropagateConstraintsExact(
   return absl::OkStatus();
 }
 
+absl::Status ConstraintPropagator::PropagateFusionBoundary(
+    const HloInstruction* fusion_instruction) {
+  const HloComputation* fused_comp =
+      fusion_instruction->fused_instructions_computation();
+  if (fused_comp == nullptr) {
+    return absl::OkStatus();
+  }
+
+  // 1. Output / Root binding:
+  const HloInstruction* fused_root =
+      fusion_instruction->fused_expression_root();
+  if (fused_root != nullptr) {
+    // Backward: outer constraint on fusion result flows into inner root.
+    ConstraintState fusion_state = states_[fusion_instruction];
+    states_[fused_root].Merge(fusion_state);
+    // Forward: internal constraint computed on root flows out to fusion result.
+    ConstraintState root_state = states_[fused_root];
+    states_[fusion_instruction].Merge(root_state);
+  }
+
+  // 2. Operands / Parameters binding:
+  for (int64_t i = 0; i < fusion_instruction->operand_count(); ++i) {
+    const HloInstruction* operand = fusion_instruction->operand(i);
+    const HloInstruction* fused_param = fusion_instruction->fused_parameter(i);
+    if (fused_param == nullptr) {
+      continue;
+    }
+    // Backward: constraints accumulated on the internal parameter flow out
+    // to the caller operand.
+    ConstraintState param_state = states_[fused_param];
+    states_[operand].Merge(param_state);
+    // Forward: constraints established on the caller operand flow into the
+    // internal parameter.
+    ConstraintState operand_state = states_[operand];
+    states_[fused_param].Merge(operand_state);
+  }
+
+  return absl::OkStatus();
+}
+
 void ConstraintPropagator::PropagateAddApprox(
     const HloInstruction* instruction,
     const ConstraintInterval& output_interval) {

third_party/xla/xla/tests/constraint_propagator.h

@@ -108,6 +108,12 @@ class ConstraintPropagator {
   // formatting can simply propagate the exact constraints to their operands.
   absl::Status PropagateConstraintsExact(const HloInstruction* instruction);
 
+  // Propagates constraints bidirectionally across a kFusion instruction
+  // boundary, mapping caller operands to callee fused parameters and caller
+  // result to the fused expression root.
+  absl::Status PropagateFusionBoundary(
+      const HloInstruction* fusion_instruction);
+
   // Propagates constraints from the output of an instruction to its operands.
   // This is approximate and introduces approximations for ops like add, sub,
   // etc. These approximations can reduce the valid search space for an input

third_party/xla/xla/tests/constraint_propagator_test.cc

@@ -1311,5 +1311,59 @@ ENTRY main {
   EXPECT_DOUBLE_EQ(x_int.min, -128.0);
   EXPECT_DOUBLE_EQ(x_int.max, 127.0);
 }
+TEST_F(ConstraintPropagatorTest,
+       FusionBoundaryPropagatesConstraintsToCallerOperands) {
+  const char* hlo = R"(
+HloModule TestModule
+
+%inner_computation (inner_param: f32[8,128]) -> f32[8,128] {
+  %inner_param = f32[8,128] parameter(0)
+  ROOT %sqrt = f32[8,128] sqrt(%inner_param)
+}
+
+ENTRY main {
+  %param_0 = f32[8,128] parameter(0)
+  ROOT %fusion = f32[8,128] fusion(%param_0), kind=kLoop, calls=%inner_computation
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
+  ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module));
+
+  auto p0_int = states[module->entry_computation()->parameter_instruction(0)]
+                    .GetConstraintInterval();
+  EXPECT_FALSE(p0_int.IsEmpty());
+  EXPECT_TRUE(p0_int.IsPositive());
+  EXPECT_GE(p0_int.min, 0.0);
+}
+
+TEST_F(ConstraintPropagatorTest,
+       NestedFusionBoundaryPropagatesConstraintsAcrossMultipleLevels) {
+  const char* hlo = R"(
+HloModule TestModule
+
+%innermost_computation (p: f32[8,128]) -> f32[8,128] {
+  %p = f32[8,128] parameter(0)
+  ROOT %sqrt = f32[8,128] sqrt(%p)
+}
+
+%outer_computation (p_outer: f32[8,128]) -> f32[8,128] {
+  %p_outer = f32[8,128] parameter(0)
+  ROOT %inner_fusion = f32[8,128] fusion(%p_outer), kind=kLoop, calls=%innermost_computation
+}
+
+ENTRY main {
+  %param_0 = f32[8,128] parameter(0)
+  ROOT %outer_fusion = f32[8,128] fusion(%param_0), kind=kOutput, calls=%outer_computation
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
+  ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module));
+
+  auto p0_int = states[module->entry_computation()->parameter_instruction(0)]
+                    .GetConstraintInterval();
+  EXPECT_FALSE(p0_int.IsEmpty());
+  EXPECT_TRUE(p0_int.IsPositive());
+  EXPECT_GE(p0_int.min, 0.0);
+}
 }  // namespace
 }  // namespace xla