This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 08b32a7976 [Runtime][ROCm] Enable ROCm host memory support (#17037)
08b32a7976 is described below
commit 08b32a797642515b0b263ead292af6962fea0cf4
Author: Ruihang Lai <ruiha...@cs.cmu.edu>
AuthorDate: Thu May 30 07:28:26 2024 -0400
[Runtime][ROCm] Enable ROCm host memory support (#17037)
This PR enables the ROCMHost memory support in ROCm device API.
---
src/runtime/ndarray.cc | 3 ++-
src/runtime/rocm/rocm_device_api.cc | 40 ++++++++++++++++++++++++++++++++-----
2 files changed, 37 insertions(+), 6 deletions(-)
diff --git a/src/runtime/ndarray.cc b/src/runtime/ndarray.cc
index c2efa79c0c..c2cf5f388a 100644
--- a/src/runtime/ndarray.cc
+++ b/src/runtime/ndarray.cc
@@ -316,7 +316,8 @@ void NDArray::CopyFromTo(const DLTensor* from, DLTensor*
to, TVMStreamHandle str
ICHECK(from->device.device_type == to->device.device_type ||
from->device.device_type == kDLCPU ||
to->device.device_type == kDLCPU || from->device.device_type ==
kDLCUDAHost ||
- to->device.device_type == kDLCUDAHost)
+ to->device.device_type == kDLCUDAHost || from->device.device_type ==
kDLROCMHost ||
+ to->device.device_type == kDLROCMHost)
<< "Can not copy across different device types directly. From device
type: "
<< from->device.device_type << " to device type: " <<
to->device.device_type;
diff --git a/src/runtime/rocm/rocm_device_api.cc
b/src/runtime/rocm/rocm_device_api.cc
index f3cc46f927..e2a5048ca0 100644
--- a/src/runtime/rocm/rocm_device_api.cc
+++ b/src/runtime/rocm/rocm_device_api.cc
@@ -144,16 +144,26 @@ class ROCMDeviceAPI final : public DeviceAPI {
*rv = value;
}
void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType
type_hint) final {
- ROCM_CALL(hipSetDevice(dev.device_id));
ICHECK_EQ(256 % alignment, 0U) << "ROCM space is aligned at 256 bytes";
void* ret;
- ROCM_CALL(hipMalloc(&ret, nbytes));
+ if (dev.device_type == kDLROCMHost) {
+ VLOG(1) << "allocating " << nbytes << "bytes on host";
+ ROCM_CALL(hipHostMalloc(&ret, nbytes));
+ } else {
+ ROCM_CALL(hipSetDevice(dev.device_id));
+ VLOG(1) << "allocating " << nbytes << " bytes on device";
+ ROCM_CALL(hipMalloc(&ret, nbytes));
+ }
return ret;
}
void FreeDataSpace(Device dev, void* ptr) final {
- ROCM_CALL(hipSetDevice(dev.device_id));
- ROCM_CALL(hipFree(ptr));
+ if (dev.device_type == kDLROCMHost) {
+ ROCM_CALL(hipHostFree(ptr));
+ } else {
+ ROCM_CALL(hipSetDevice(dev.device_id));
+ ROCM_CALL(hipFree(ptr));
+ }
}
void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t
to_offset, size_t size,
@@ -162,6 +172,21 @@ class ROCMDeviceAPI final : public DeviceAPI {
hipStream_t hip_stream = static_cast<hipStream_t>(stream);
from = static_cast<const char*>(from) + from_offset;
to = static_cast<char*>(to) + to_offset;
+
+ if (dev_from.device_type == kDLROCMHost) {
+ dev_from.device_type = kDLCPU;
+ }
+
+ if (dev_to.device_type == kDLROCMHost) {
+ dev_to.device_type = kDLCPU;
+ }
+
+ // In case there is a copy from host mem to host mem */
+ if (dev_to.device_type == kDLCPU && dev_from.device_type == kDLCPU) {
+ memcpy(to, from, size);
+ return;
+ }
+
if (dev_from.device_type == kDLROCM && dev_to.device_type == kDLROCM) {
ROCM_CALL(hipSetDevice(dev_from.device_id));
if (dev_from.device_id == dev_to.device_id) {
@@ -210,7 +235,7 @@ class ROCMDeviceAPI final : public DeviceAPI {
private:
static void GPUCopy(const void* from, void* to, size_t size, hipMemcpyKind
kind,
hipStream_t stream) {
- if (stream != 0) {
+ if (stream != nullptr) {
ROCM_CALL(hipMemcpyAsync(to, from, size, kind, stream));
} else {
ROCM_CALL(hipMemcpy(to, from, size, kind));
@@ -229,6 +254,11 @@ TVM_REGISTER_GLOBAL("device_api.rocm").set_body([](TVMArgs
args, TVMRetValue* rv
*rv = static_cast<void*>(ptr);
});
+TVM_REGISTER_GLOBAL("device_api.rocm_host").set_body([](TVMArgs args,
TVMRetValue* rv) {
+ DeviceAPI* ptr = ROCMDeviceAPI::Global();
+ *rv = static_cast<void*>(ptr);
+});
+
class ROCMTimerNode : public TimerNode {
public:
virtual void Start() {
