Memory space related copies should not be normalized.

PiperOrigin-RevId: 673102513
openxla · Sep 17, 2024 · 40b5de1 · 40b5de1
1 parent 891d972
commit 40b5de1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/xla/service/host_offload_utils.cc b/xla/service/host_offload_utils.cc
@@ -244,5 +244,17 @@ bool IsHostAsyncStart(const HloInstruction* instruction) {
          instruction->async_execution_thread() == HloInstruction::kHostThread;
 }
 
+bool IsSynchronousCopyFromOrToHost(const HloInstruction* instruction) {
+  if (instruction->opcode() != HloOpcode::kCopy) {
+    return false;
+  }
+  return (instruction->shape().has_layout() &&
+          instruction->shape().layout().memory_space() ==
+              Layout::kHostMemorySpace) ||
+         (instruction->operand(0)->shape().has_layout() &&
+          instruction->operand(0)->shape().layout().memory_space() ==
+              Layout::kHostMemorySpace);
+}
+
 }  // namespace host_offload_utils
 }  // namespace xla
diff --git a/xla/service/host_offload_utils.h b/xla/service/host_offload_utils.h
@@ -98,6 +98,9 @@ bool IsValidDuringPureMemoryOffload(const HloInstruction* instruction);
 // Returns true if the instruction is an async-start with host thread.
 bool IsHostAsyncStart(const HloInstruction* instruction);
 
+// Returns true if the copy is from or to host memory space.
+bool IsSynchronousCopyFromOrToHost(const HloInstruction* instruction);
+
 }  // namespace host_offload_utils
 }  // namespace xla