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" 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" 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