[XLA] Small refactor in sharding propagation.

mkuperst · tensorflower-gardener · commit 180f95f9f23e · 2025-06-02T11:25:02.000-07:00
Since we now propagate to both while bodies and while conditionals, while doesn't need to be handled in a special way.

PiperOrigin-RevId: 766257287
diff --git a/third_party/xla/xla/service/sharding_propagation.cc b/third_party/xla/xla/service/sharding_propagation.cc
@@ -3265,13 +3265,8 @@ absl::StatusOr<bool> ShardingPropagation::Run(
             inst->copy_sharding(sharded_inst);
           }
         }
-        if (instruction->opcode() == HloOpcode::kWhile) {
-          computation_map[instruction->while_body()] = instruction;
-          computation_map[instruction->while_condition()] = instruction;
-        } else {
-          for (HloComputation* c : instruction->called_computations()) {
-            computation_map[c] = instruction;
-          }
+        for (HloComputation* c : instruction->called_computations()) {
+          computation_map[c] = instruction;
         }
       }
     }

Original file line number	Diff line number	Diff line change
`@@ -3265,13 +3265,8 @@ absl::StatusOr<bool> ShardingPropagation::Run(`
`3265`	`3265`	`inst->copy_sharding(sharded_inst);`
`3266`	`3266`	`}`
`3267`	`3267`	`}`
`3268`		`- if (instruction->opcode() == HloOpcode::kWhile) {`
`3269`		`- computation_map[instruction->while_body()] = instruction;`
`3270`		`- computation_map[instruction->while_condition()] = instruction;`
`3271`		`- } else {`
`3272`		`- for (HloComputation* c : instruction->called_computations()) {`
`3273`		`- computation_map[c] = instruction;`
`3274`		`- }`
	`3268`	`+ for (HloComputation* c : instruction->called_computations()) {`
	`3269`	`+ computation_map[c] = instruction;`
`3275`	`3270`	`}`
`3276`	`3271`	`}`
`3277`	`3272`	`}`