Move GpuDriver::GetPointerMemorySpace to the appropriate Executor classes.

PiperOrigin-RevId: 685292532

Author: klucke

third_party/xla/xla/stream_executor/cuda/cuda_driver.cc

@@ -944,22 +944,6 @@ int GpuDriver::GetDeviceCount() {
   return device_count;
 }
 
-absl::StatusOr<MemoryType> GpuDriver::GetPointerMemorySpace(
-    CUdeviceptr pointer) {
-  unsigned int value;
-  TF_RETURN_IF_ERROR(cuda::ToStatus(cuPointerGetAttribute(
-      &value, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, pointer)));
-  switch (value) {
-    case CU_MEMORYTYPE_DEVICE:
-      return MemoryType::kDevice;
-    case CU_MEMORYTYPE_HOST:
-      return MemoryType::kHost;
-    default:
-      return absl::InternalError(
-          absl::StrCat("unknown memory space provided by CUDA API: ", value));
-  }
-}
-
 absl::Status GpuDriver::GetPointerAddressRange(CUdeviceptr dptr,
                                                CUdeviceptr* base,
                                                size_t* size) {

third_party/xla/xla/stream_executor/cuda/cuda_executor.cc

@@ -1171,5 +1171,22 @@ CudaExecutor::CreateDeviceDescription(int device_ordinal) {
   return std::make_unique<DeviceDescription>(std::move(desc));
 }
 
+absl::StatusOr<MemoryType> CudaExecutor::GetPointerMemorySpace(
+    const void* ptr) {
+  CUdeviceptr pointer = reinterpret_cast<CUdeviceptr>(const_cast<void*>(ptr));
+  unsigned int value;
+  TF_RETURN_IF_ERROR(cuda::ToStatus(cuPointerGetAttribute(
+      &value, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, pointer)));
+  switch (value) {
+    case CU_MEMORYTYPE_DEVICE:
+      return MemoryType::kDevice;
+    case CU_MEMORYTYPE_HOST:
+      return MemoryType::kHost;
+    default:
+      return absl::InternalError(
+          absl::StrCat("unknown memory space provided by CUDA API: ", value));
+  }
+}
+
 }  // namespace gpu
 }  // namespace stream_executor

third_party/xla/xla/stream_executor/cuda/cuda_executor.h

@@ -139,10 +139,7 @@ class CudaExecutor : public GpuExecutor {
   bool HostMemoryRegister(void* location, uint64_t size) override;
   bool HostMemoryUnregister(void* location) override;
 
-  absl::StatusOr<MemoryType> GetPointerMemorySpace(const void* ptr) override {
-    return GpuDriver::GetPointerMemorySpace(
-        reinterpret_cast<GpuDevicePtr>(const_cast<void*>(ptr)));
-  }
+  absl::StatusOr<MemoryType> GetPointerMemorySpace(const void* ptr) override;
 
   Stream* FindAllocatedStream(void* gpu_stream) override {
     absl::MutexLock lock(&alive_gpu_streams_mu_);

third_party/xla/xla/stream_executor/gpu/gpu_driver.h

@@ -421,9 +421,6 @@ class GpuDriver {
 
   // -- Pointer-specific calls.
 
-  // Returns the memory space addressed by pointer.
-  static absl::StatusOr<MemoryType> GetPointerMemorySpace(GpuDevicePtr pointer);
-
   // Returns the base address and size of the device pointer dptr.
   static absl::Status GetPointerAddressRange(GpuDevicePtr dptr,
                                              GpuDevicePtr* base, size_t* size);

third_party/xla/xla/stream_executor/rocm/rocm_driver.cc

@@ -766,27 +766,6 @@ absl::Status GpuDriver::GetPointerAddressRange(hipDeviceptr_t dptr,
                       reinterpret_cast<void*>(dptr), ToString(result).c_str()));
 }
 
-absl::StatusOr<MemoryType> GpuDriver::GetPointerMemorySpace(
-    hipDeviceptr_t pointer) {
-  unsigned int value;
-  hipError_t result = wrap::hipPointerGetAttribute(
-      &value, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE, pointer);
-  if (result == hipSuccess) {
-    switch (value) {
-      case hipMemoryTypeDevice:
-        return MemoryType::kDevice;
-      case hipMemoryTypeHost:
-        return MemoryType::kHost;
-      default:
-        return absl::InternalError(
-            absl::StrCat("unknown memory space provided by ROCM API: ", value));
-    }
-  }
-
-  return absl::InternalError(absl::StrCat(
-      "failed to query device pointer for memory space: ", ToString(result)));
-}
-
 absl::StatusOr<int32_t> GpuDriver::GetDriverVersion() {
   int32_t version;
   TF_RETURN_IF_ERROR(ToStatus(wrap::hipDriverGetVersion(&version),

third_party/xla/xla/stream_executor/rocm/rocm_executor.cc

@@ -37,6 +37,7 @@ limitations under the License.
 #include "absl/synchronization/mutex.h"
 #include "absl/synchronization/notification.h"
 #include "absl/types/span.h"
+#include "rocm/include/hip/driver_types.h"
 #include "rocm/include/hip/hip_runtime.h"
 #include "rocm/include/hip/hip_version.h"
 #include "rocm/rocm_config.h"
@@ -975,6 +976,29 @@ RocmExecutor::CreateDeviceDescription(int device_ordinal) {
   return std::make_unique<DeviceDescription>(std::move(desc));
 }
 
+absl::StatusOr<MemoryType> RocmExecutor::GetPointerMemorySpace(
+    const void* ptr) {
+  hipDeviceptr_t pointer =
+      reinterpret_cast<hipDeviceptr_t>(const_cast<void*>(ptr));
+  unsigned int value;
+  hipError_t result = wrap::hipPointerGetAttribute(
+      &value, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE, pointer);
+  if (result == hipSuccess) {
+    switch (value) {
+      case hipMemoryTypeDevice:
+        return MemoryType::kDevice;
+      case hipMemoryTypeHost:
+        return MemoryType::kHost;
+      default:
+        return absl::InternalError(
+            absl::StrCat("unknown memory space provided by ROCM API: ", value));
+    }
+  }
+
+  return absl::InternalError(absl::StrCat(
+      "failed to query device pointer for memory space: ", ToString(result)));
+}
+
 }  // namespace gpu
 
 }  // namespace stream_executor

third_party/xla/xla/stream_executor/rocm/rocm_executor.h

@@ -128,10 +128,7 @@ class RocmExecutor : public GpuExecutor {
     return GpuDriver::HostDeallocate(gpu_context(), location);
   }
 
-  absl::StatusOr<MemoryType> GetPointerMemorySpace(const void* ptr) override {
-    return GpuDriver::GetPointerMemorySpace(
-        reinterpret_cast<GpuDevicePtr>(const_cast<void*>(ptr)));
-  }
+  absl::StatusOr<MemoryType> GetPointerMemorySpace(const void* ptr) override;
 
   Stream* FindAllocatedStream(void* gpu_stream) override {
     absl::MutexLock lock(&alive_gpu_streams_mu_);