Add kCpu property tag.

PiperOrigin-RevId: 719007295
openxla · Jan 23, 2025 · 53c999d · 53c999d
1 parent b3f7ab6
commit 53c999d
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/xla/service/hlo_runner.cc b/xla/service/hlo_runner.cc
@@ -718,6 +718,9 @@ bool HloRunner::HasProperty(const HloRunnerPropertyTag::Type tag) const {
     return std::holds_alternative<stream_executor::RocmComputeCapability>(
         device_description.gpu_compute_capability());
   }
+  if (tag == HloRunnerPropertyTag::kCpu) {
+    return backend().platform()->Name() == "Host";
+  }
   return false;
 }
 

diff --git a/xla/service/hlo_runner_interface.h b/xla/service/hlo_runner_interface.h
@@ -75,6 +75,8 @@ class HloRunnerPropertyTag final {
   static constexpr Type kDefault = 0;
   // Indicates that the runner is using ROCm.
   static constexpr Type kUsingGpuRocm = 1;
+  // Indicates that this runner is a CPU runner.
+  static constexpr Type kCpu = 2;
 
  private:
   HloRunnerPropertyTag() = default;

diff --git a/xla/service/hlo_runner_pjrt.cc b/xla/service/hlo_runner_pjrt.cc
@@ -662,6 +662,9 @@ bool HloRunnerPjRt::HasProperty(const HloRunnerPropertyTag::Type tag) const {
   if (tag == HloRunnerPropertyTag::kUsingGpuRocm) {
     return pjrt_client_->platform_name() == xla::RocmName();
   }
+  if (tag == HloRunnerPropertyTag::kCpu) {
+    return pjrt_client_->platform_name() == xla::CpuName();
+  }
   return false;
 }