Skip to content

Commit

Permalink
[XLA:HOST_OFFLOAD] Add Copy to the list of movable opcodes between ho…
Browse files Browse the repository at this point in the history
…st to

device and dynamic slice.

PiperOrigin-RevId: 610968466
  • Loading branch information
blakehechtman authored and tensorflower-gardener committed May 6, 2024
1 parent 8c1ea34 commit 0719f3e
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion third_party/xla/xla/service/host_offloader.cc
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,8 @@ absl::StatusOr<std::vector<HloInstruction*>> GetBufferUsersOfType(
// And the buffer is passed through the first operand.
// This is used to trace the graph between an annotation and its relevant slice.
bool CanTraverseOpBetweenAnnotation(HloInstruction* hlo) {
if (hlo->opcode() == HloOpcode::kBitcast) {
if (hlo->opcode() == HloOpcode::kBitcast ||
hlo->opcode() == HloOpcode::kCopy) {
return true;
} else if (hlo->opcode() == HloOpcode::kReshape) {
return ShapeUtil::ReshapeIsBitcast(hlo->operand(0)->shape(), hlo->shape());
Expand Down

0 comments on commit 0719f3e

Please sign in to comment.