[feat] Add support for Ascend device memory transfers (#382)

PR adds support for transferring data between device and host memory using Ascend native APIs. Currently, only Ascend-specific native APIs have been integrated, with core computing kernels to be added soon.

- `ascend_stream.cc`: Provides wrappers for ACLRT copy operations, including synchronous, asynchronous, and batched operations.
- `ascend_buffer.cc`: Offers wrappers for ACLRT memory management APIs, including device memory allocation/deallocation, host memory allocation/deallocation, and host memory mapping/unmapping.

Author: mag1c-h

ucm/shared/test/case/trans/trans_test.cc

@@ -68,6 +68,7 @@ TEST_F(UCTransUnitTest, CopyDataWithSM)
     ASSERT_EQ(device.Setup(deviceId), ok);
     auto buffer = device.MakeBuffer();
     auto stream = device.MakeSMStream();
+    if (!stream) { return; }
     auto hPtr1 = buffer->MakeHostBuffer(size * number);
     ASSERT_NE(hPtr1, nullptr);
     ASSERT_EQ(buffer->MakeDeviceBuffers(size, number), ok);

ucm/shared/trans/ascend/CMakeLists.txt

@@ -0,0 +1,15 @@
+set(ASCEND_ROOT "/usr/local/Ascend/ascend-toolkit/latest" CACHE PATH "Path to Ascend root directory")
+add_library(Ascend::ascendcl UNKNOWN IMPORTED)
+set_target_properties(Ascend::ascendcl PROPERTIES
+    INTERFACE_INCLUDE_DIRECTORIES "${ASCEND_ROOT}/include"
+    IMPORTED_LOCATION "${ASCEND_ROOT}/lib64/libascendcl.so"
+)
+add_library(trans STATIC
+    ascend_device.cc
+    ascend_buffer.cc
+    ascend_stream.cc
+)
+target_link_libraries(trans PUBLIC
+    fmt
+    Ascend::ascendcl
+)

ucm/shared/trans/ascend/ascend_buffer.cc

@@ -0,0 +1,76 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ * */
+#include "ascend_buffer.h"
+#include <acl/acl.h>
+
+namespace UC::Trans {
+
+std::shared_ptr<void> Trans::AscendBuffer::MakeDeviceBuffer(size_t size)
+{
+    void* device = nullptr;
+    auto ret = aclrtMalloc(&device, size, ACL_MEM_TYPE_HIGH_BAND_WIDTH);
+    if (ret == ACL_SUCCESS) { return std::shared_ptr<void>(device, aclrtFree); }
+    return nullptr;
+}
+
+std::shared_ptr<void> Trans::AscendBuffer::MakeHostBuffer(size_t size)
+{
+    void* host = nullptr;
+    auto ret = aclrtMallocHost(&host, size);
+    if (ret == ACL_SUCCESS) { return std::shared_ptr<void>(host, aclrtFreeHost); }
+    return nullptr;
+}
+
+Status Trans::AscendBuffer::RegisterHostBuffer(void* ptr, size_t size)
+{
+    if (registerOnHost_) {
+        aclrtHostUnregister(registerOnHost_);
+        registerOnHost_ = nullptr;
+        registerOnDevice_ = nullptr;
+    }
+    auto ret = aclrtHostRegister(ptr, size, ACL_HOST_REGISTER_MAPPED, &registerOnDevice_);
+    if (ret == ACL_SUCCESS) {
+        registerOnHost_ = ptr;
+        return Status::OK();
+    }
+    registerOnDevice_ = nullptr;
+    return Status{ret, std::to_string(ret)};
+}
+
+void Trans::AscendBuffer::UnregisterHostBuffer(void* ptr)
+{
+    aclrtHostUnregister(ptr);
+    if (registerOnHost_ == ptr) {
+        registerOnHost_ = nullptr;
+        registerOnDevice_ = nullptr;
+    }
+}
+
+void* Trans::AscendBuffer::GetHostPtrOnDevice(void* ptr)
+{
+    if (registerOnHost_ == ptr) { return registerOnDevice_; }
+    return nullptr;
+}
+
+} // namespace UC::Trans

ucm/shared/trans/ascend/ascend_buffer.h

@@ -0,0 +1,46 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ * */
+#ifndef UNIFIEDCACHE_TRANS_ASCEND_BUFFER_H
+#define UNIFIEDCACHE_TRANS_ASCEND_BUFFER_H
+
+#include "trans/detail/reserved_buffer.h"
+
+namespace UC::Trans {
+
+class AscendBuffer : public ReservedBuffer {
+    void* registerOnHost_{nullptr};
+    void* registerOnDevice_{nullptr};
+
+public:
+    std::shared_ptr<void> MakeDeviceBuffer(size_t size) override;
+    std::shared_ptr<void> MakeHostBuffer(size_t size) override;
+
+    Status RegisterHostBuffer(void* ptr, size_t size) override;
+    void UnregisterHostBuffer(void* ptr) override;
+    void* GetHostPtrOnDevice(void* ptr) override;
+};
+
+} // namespace UC::Trans
+
+#endif

ucm/shared/trans/ascend/ascend_device.cc

@@ -0,0 +1,62 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ * */
+#include <acl/acl.h>
+#include "ascend_buffer.h"
+#include "ascend_stream.h"
+#include "trans/device.h"
+
+namespace UC::Trans {
+
+Status Device::Setup(int32_t deviceId)
+{
+    if (deviceId < 0) { return Status::Error(fmt::format("invalid device id({})", deviceId)); }
+    auto ret = aclrtSetDevice(deviceId);
+    if (ret == ACL_SUCCESS) { return Status::OK(); }
+    return Status{ret, std::to_string(ret)};
+}
+
+std::unique_ptr<Stream> Device::MakeStream()
+{
+    std::unique_ptr<Stream> stream = nullptr;
+    try {
+        stream = std::make_unique<AscendStream>();
+    } catch (...) {
+        return nullptr;
+    }
+    if (stream->Setup().Success()) { return stream; }
+    return nullptr;
+}
+
+std::unique_ptr<Stream> Device::MakeSMStream() { return nullptr; }
+
+std::unique_ptr<Buffer> Device::MakeBuffer()
+{
+    try {
+        return std::make_unique<AscendBuffer>();
+    } catch (...) {
+        return nullptr;
+    }
+}
+
+} // namespace UC::Trans

ucm/shared/trans/ascend/ascend_stream.cc

@@ -0,0 +1,177 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ * */
+#include "ascend_stream.h"
+
+namespace UC::Trans {
+
+AscendStream::~AscendStream()
+{
+    if (cbThread_.joinable()) {
+        auto tid = cbThread_.native_handle();
+        (void)aclrtUnSubscribeReport(tid, stream_);
+        stop_ = true;
+        cbThread_.join();
+    }
+    if (stream_) {
+        (void)aclrtDestroyStream(stream_);
+        stream_ = nullptr;
+    }
+}
+
+Status AscendStream::Setup()
+{
+    auto ret = aclrtCreateStream(&stream_);
+    if (ret != ACL_SUCCESS) [[unlikely]] { return Status{ret, std::to_string(ret)}; }
+    cbThread_ = std::thread([this] {
+        while (!this->stop_) { (void)aclrtProcessReport(10); }
+    });
+    auto tid = cbThread_.native_handle();
+    ret = aclrtSubscribeReport(tid, stream_);
+    if (ret != ACL_SUCCESS) [[unlikely]] { return Status{ret, std::to_string(ret)}; }
+    return Status::OK();
+}
+
+Status AscendStream::DeviceToHost(void* device, void* host, size_t size)
+{
+    auto ret = aclrtMemcpy(host, size, device, size, ACL_MEMCPY_DEVICE_TO_HOST);
+    if (ret == ACL_SUCCESS) { return Status::OK(); }
+    return Status{ret, std::to_string(ret)};
+}
+
+Status AscendStream::DeviceToHost(void* device[], void* host[], size_t size, size_t number)
+{
+    auto s = DeviceToHostAsync(device, host, size, number);
+    if (s.Failure()) [[unlikely]] { return s; }
+    return Synchronized();
+}
+
+Status AscendStream::DeviceToHost(void* device[], void* host, size_t size, size_t number)
+{
+    auto s = DeviceToHostAsync(device, host, size, number);
+    if (s.Failure()) [[unlikely]] { return s; }
+    return Synchronized();
+}
+
+Status AscendStream::DeviceToHostAsync(void* device, void* host, size_t size)
+{
+    auto ret = aclrtMemcpyAsync(host, size, device, size, ACL_MEMCPY_DEVICE_TO_HOST, stream_);
+    if (ret == ACL_SUCCESS) { return Status::OK(); }
+    return Status{ret, std::to_string(ret)};
+}
+
+Status AscendStream::DeviceToHostAsync(void* device[], void* host[], size_t size, size_t number)
+{
+    for (size_t i = 0; i < number; i++) {
+        auto s = DeviceToHostAsync(device[i], host[i], size);
+        if (s.Failure()) [[unlikely]] { return s; }
+    }
+    return Status::OK();
+}
+
+Status AscendStream::DeviceToHostAsync(void* device[], void* host, size_t size, size_t number)
+{
+    for (size_t i = 0; i < number; i++) {
+        auto pHost = (void*)(((int8_t*)host) + size * i);
+        auto s = DeviceToHostAsync(device[i], pHost, size);
+        if (s.Failure()) [[unlikely]] { return s; }
+    }
+    return Status::OK();
+}
+
+Status AscendStream::HostToDevice(void* host, void* device, size_t size)
+{
+    auto ret = aclrtMemcpy(device, size, host, size, ACL_MEMCPY_HOST_TO_DEVICE);
+    if (ret == ACL_SUCCESS) { return Status::OK(); }
+    return Status{ret, std::to_string(ret)};
+}
+
+Status AscendStream::HostToDevice(void* host[], void* device[], size_t size, size_t number)
+{
+    auto s = HostToDeviceAsync(host, device, size, number);
+    if (s.Failure()) [[unlikely]] { return s; }
+    return Synchronized();
+}
+
+Status AscendStream::HostToDevice(void* host, void* device[], size_t size, size_t number)
+{
+    auto s = HostToDeviceAsync(host, device, size, number);
+    if (s.Failure()) [[unlikely]] { return s; }
+    return Synchronized();
+}
+
+Status AscendStream::HostToDeviceAsync(void* host, void* device, size_t size)
+{
+    auto ret = aclrtMemcpyAsync(device, size, host, size, ACL_MEMCPY_HOST_TO_DEVICE, stream_);
+    if (ret == ACL_SUCCESS) { return Status::OK(); }
+    return Status{ret, std::to_string(ret)};
+}
+
+Status AscendStream::HostToDeviceAsync(void* host[], void* device[], size_t size, size_t number)
+{
+    for (size_t i = 0; i < number; i++) {
+        auto s = HostToDeviceAsync(host[i], device[i], size);
+        if (s.Failure()) [[unlikely]] { return s; }
+    }
+    return Status::OK();
+}
+
+Status AscendStream::HostToDeviceAsync(void* host, void* device[], size_t size, size_t number)
+{
+    for (size_t i = 0; i < number; i++) {
+        auto pHost = (void*)(((int8_t*)host) + size * i);
+        auto s = HostToDeviceAsync(pHost, device[i], size);
+        if (s.Failure()) [[unlikely]] { return s; }
+    }
+    return Status::OK();
+}
+
+using Closure = std::function<void(bool)>;
+
+static void Trampoline(void* data)
+{
+    auto c = static_cast<Closure*>(data);
+    (*c)(true);
+    delete c;
+}
+
+Status Trans::AscendStream::AppendCallback(std::function<void(bool)> cb)
+{
+    auto c = new (std::nothrow) Closure{std::move(cb)};
+    if (!c) [[unlikely]] { return Status::Error("out of memory for appending callback"); }
+    auto ret = aclrtLaunchCallback(Trampoline, (void*)c, ACL_CALLBACK_NO_BLOCK, stream_);
+    if (ret != ACL_SUCCESS) [[unlikely]] {
+        delete c;
+        return Status{ret, std::to_string(ret)};
+    }
+    return Status::OK();
+}
+
+Status AscendStream::Synchronized()
+{
+    auto ret = aclrtSynchronizeStream(stream_);
+    if (ret == ACL_SUCCESS) { return Status::OK(); }
+    return Status{ret, std::to_string(ret)};
+}
+
+} // namespace UC::Trans