Introduce host_offload_utils::ComputeTypeIsHost to OSS.

PiperOrigin-RevId: 686610541
openxla · Oct 16, 2024 · 02261a0 · 02261a0
1 parent d5ff7c1
commit 02261a0
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/xla/service/BUILD b/xla/service/BUILD
@@ -6628,6 +6628,7 @@ cc_library(
         ":pattern_matcher",
         "//xla:literal_util",
         "//xla:shape_util",
+        "//xla:side_effect_util",
         "//xla:status_macros",
         "//xla:util",
         "//xla/hlo/ir:hlo",

diff --git a/xla/service/host_offload_utils.cc b/xla/service/host_offload_utils.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/service/call_graph.h"
 #include "xla/service/host_memory_offload_annotations.h"
 #include "xla/shape_util.h"
+#include "xla/side_effect_util.h"
 #include "xla/util.h"
 
 namespace xla {
@@ -256,5 +257,14 @@ bool IsSynchronousCopyFromOrToHost(const HloInstruction* instruction) {
               Layout::kHostMemorySpace);
 }
 
+bool ComputeTypeIsHost(const HloInstruction* hlo_instruction) {
+  const auto& frontend_attributes_map =
+      hlo_instruction->frontend_attributes().map();
+  return (frontend_attributes_map.find(kXlaComputeTypeAttr) !=
+              frontend_attributes_map.end() &&
+          frontend_attributes_map.find(kXlaComputeTypeAttr)->second ==
+              kXlaComputeTypeHost);
+}
+
 }  // namespace host_offload_utils
 }  // namespace xla
diff --git a/xla/service/host_offload_utils.h b/xla/service/host_offload_utils.h
@@ -101,6 +101,8 @@ bool IsHostAsyncStart(const HloInstruction* instruction);
 // Returns true if the copy is from or to host memory space.
 bool IsSynchronousCopyFromOrToHost(const HloInstruction* instruction);
 
+bool ComputeTypeIsHost(const HloInstruction* hlo_instruction);
+
 }  // namespace host_offload_utils
 }  // namespace xla