From ba73aa33d5712108696e873bd8761ef41ff2b3b7 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 29 Aug 2026 17:34:02 -0700 Subject: [PATCH 1/3] Automated Code Change PiperOrigin-RevId: 973246125 --- third_party/xla/xla/pjrt/dump/BUILD | 1 + third_party/xla/xla/pjrt/dump/dump_test.cc | 1 + 2 files changed, 2 insertions(+) diff --git a/third_party/xla/xla/pjrt/dump/BUILD b/third_party/xla/xla/pjrt/dump/BUILD index ea8a252995ffd0..defbd319d92f0e 100644 --- a/third_party/xla/xla/pjrt/dump/BUILD +++ b/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", diff --git a/third_party/xla/xla/pjrt/dump/dump_test.cc b/third_party/xla/xla/pjrt/dump/dump_test.cc index 5f2f4ef7ca9aa5..cffeac5c587fa5 100644 --- a/third_party/xla/xla/pjrt/dump/dump_test.cc +++ b/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" From 619d48ef9b9ffa80a10324e1d6b9441b5b5f25a8 Mon Sep 17 00:00:00 2001 From: Bhatu Date: Sat, 29 Aug 2026 21:28:26 -0700 Subject: [PATCH 2/3] Propagate input constraints across nested fusion instruction boundaries. Previously, ConstraintPropagator only analyzed computations in isolation and stopped at kFusion instructions. In multi-level fusions (such as kOutput fusions containing inner kLoop fusions), constraints from internal operations failed to propagate backward to the module's entry parameters. Consequently, parameters remained at their default unconstrained ranges causing floating-point overflow to +inf and downstream NaN generation. This change: - Implements PropagateFusionBoundary in PropagateConstraintsExact to bidirectionally synchronize caller operands with callee fused parameters, and caller results with the fused expression root. - Updates ConstraintPropagator::Run to perform an inter-procedural fixed-point iteration across all computations in topological post-order until global convergence. - Adds unit tests verifying constraint propagation across both single-level and multi-level nested fusions. PiperOrigin-RevId: 973306649 --- .../xla/xla/tests/constraint_propagator.cc | 62 ++++++++++++++++++- .../xla/xla/tests/constraint_propagator.h | 6 ++ .../xla/tests/constraint_propagator_test.cc | 54 ++++++++++++++++ 3 files changed, 121 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index e47a042be3cc70..8ce31e5b30cff4 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/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 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 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) { diff --git a/third_party/xla/xla/tests/constraint_propagator.h b/third_party/xla/xla/tests/constraint_propagator.h index 9ada9ec0953f01..d70045c768df57 100644 --- a/third_party/xla/xla/tests/constraint_propagator.h +++ b/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 diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 0b937c57f66212..340d0f881b65aa 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/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 From 9794b52cf8e167fdee4570c63791fdaac2f12d38 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 29 Aug 2026 21:35:54 -0700 Subject: [PATCH 3/3] Automated Code Change PiperOrigin-RevId: 973308466 --- .../third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc index 355c9ad410c0a6..2c504201fc53d7 100644 --- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc +++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include "absl/memory/memory.h" #include "absl/status/status.h" #include "xla/tsl/platform/macros.h" #include "xla/tsl/platform/test.h"