Skip to content

Commit 043ea30

Browse files
fhoushmandGoogle-ML-Automation
authored andcommitted
[XLA:MSA] Allow dynamic-slice when performing post allocation transformation in MSA.
Also, bypass users of constants when checking for in-place users. PiperOrigin-RevId: 729737268
1 parent 5733cf0 commit 043ea30

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

xla/service/memory_space_assignment/algorithm.cc

+4
Original file line numberDiff line numberDiff line change
@@ -2050,6 +2050,10 @@ absl::StatusOr<HeapSimulator::Result<HloValue>> MsaAlgorithm::Finish() {
20502050
// If any of the operands of the instruction has an in-place user, we
20512051
// don't run the post-allocation transformation.
20522052
for (HloInstruction* operand : instr->operands()) {
2053+
// We don't care about users of constants.
2054+
if (operand->opcode() == HloOpcode::kConstant) {
2055+
continue;
2056+
}
20532057
for (HloInstruction* user : operand->users()) {
20542058
if (HloDataflowAnalysis::IsInPlaceOperation(user->opcode())) {
20552059
continue;

0 commit comments

Comments
 (0)