From ea3993e196a05254e5084022ae9f9835485f4585 Mon Sep 17 00:00:00 2001 From: Kyle Lucke Date: Sat, 12 Oct 2024 18:12:10 -0700 Subject: [PATCH] Move GpuDriver::GetPointerMemorySpace to the appropriate Executor classes. PiperOrigin-RevId: 685292532 --- xla/stream_executor/cuda/cuda_driver.cc | 16 --------------- xla/stream_executor/cuda/cuda_executor.cc | 17 ++++++++++++++++ xla/stream_executor/cuda/cuda_executor.h | 5 +---- xla/stream_executor/gpu/gpu_driver.h | 3 --- xla/stream_executor/rocm/rocm_driver.cc | 21 -------------------- xla/stream_executor/rocm/rocm_executor.cc | 24 +++++++++++++++++++++++ xla/stream_executor/rocm/rocm_executor.h | 5 +---- 7 files changed, 43 insertions(+), 48 deletions(-) diff --git a/xla/stream_executor/cuda/cuda_driver.cc b/xla/stream_executor/cuda/cuda_driver.cc index f84ce537732c8..e53b066c80eb4 100644 --- a/xla/stream_executor/cuda/cuda_driver.cc +++ b/xla/stream_executor/cuda/cuda_driver.cc @@ -944,22 +944,6 @@ int GpuDriver::GetDeviceCount() { return device_count; } -absl::StatusOr 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) { diff --git a/xla/stream_executor/cuda/cuda_executor.cc b/xla/stream_executor/cuda/cuda_executor.cc index 12e2792644d9e..8c7b408dddfbb 100644 --- a/xla/stream_executor/cuda/cuda_executor.cc +++ b/xla/stream_executor/cuda/cuda_executor.cc @@ -1171,5 +1171,22 @@ CudaExecutor::CreateDeviceDescription(int device_ordinal) { return std::make_unique(std::move(desc)); } +absl::StatusOr CudaExecutor::GetPointerMemorySpace( + const void* ptr) { + CUdeviceptr pointer = reinterpret_cast(const_cast(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 diff --git a/xla/stream_executor/cuda/cuda_executor.h b/xla/stream_executor/cuda/cuda_executor.h index af68c22019ce8..985583844e2b5 100644 --- a/xla/stream_executor/cuda/cuda_executor.h +++ b/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 GetPointerMemorySpace(const void* ptr) override { - return GpuDriver::GetPointerMemorySpace( - reinterpret_cast(const_cast(ptr))); - } + absl::StatusOr GetPointerMemorySpace(const void* ptr) override; Stream* FindAllocatedStream(void* gpu_stream) override { absl::MutexLock lock(&alive_gpu_streams_mu_); diff --git a/xla/stream_executor/gpu/gpu_driver.h b/xla/stream_executor/gpu/gpu_driver.h index 67cca9cf4686f..889db461fe65f 100644 --- a/xla/stream_executor/gpu/gpu_driver.h +++ b/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 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); diff --git a/xla/stream_executor/rocm/rocm_driver.cc b/xla/stream_executor/rocm/rocm_driver.cc index b041314774e4d..9fc5a9653e0cd 100644 --- a/xla/stream_executor/rocm/rocm_driver.cc +++ b/xla/stream_executor/rocm/rocm_driver.cc @@ -766,27 +766,6 @@ absl::Status GpuDriver::GetPointerAddressRange(hipDeviceptr_t dptr, reinterpret_cast(dptr), ToString(result).c_str())); } -absl::StatusOr 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 GpuDriver::GetDriverVersion() { int32_t version; TF_RETURN_IF_ERROR(ToStatus(wrap::hipDriverGetVersion(&version), diff --git a/xla/stream_executor/rocm/rocm_executor.cc b/xla/stream_executor/rocm/rocm_executor.cc index b2320ced736b6..010b5306f2f8a 100644 --- a/xla/stream_executor/rocm/rocm_executor.cc +++ b/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(std::move(desc)); } +absl::StatusOr RocmExecutor::GetPointerMemorySpace( + const void* ptr) { + hipDeviceptr_t pointer = + reinterpret_cast(const_cast(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 diff --git a/xla/stream_executor/rocm/rocm_executor.h b/xla/stream_executor/rocm/rocm_executor.h index 0bfa60f8279ff..5dc070ebf5d6d 100644 --- a/xla/stream_executor/rocm/rocm_executor.h +++ b/xla/stream_executor/rocm/rocm_executor.h @@ -128,10 +128,7 @@ class RocmExecutor : public GpuExecutor { return GpuDriver::HostDeallocate(gpu_context(), location); } - absl::StatusOr GetPointerMemorySpace(const void* ptr) override { - return GpuDriver::GetPointerMemorySpace( - reinterpret_cast(const_cast(ptr))); - } + absl::StatusOr GetPointerMemorySpace(const void* ptr) override; Stream* FindAllocatedStream(void* gpu_stream) override { absl::MutexLock lock(&alive_gpu_streams_mu_);