[Opt] Replace acl with abstract runtime in kv client

Co-authored-by: zhaoxiaoyan-0826 <87650636+zhaoxiaoyan-0826@users.noreply.github.com>

Author: Infinite666

CMakeLists.txt

@@ -11,6 +11,7 @@ option(BUILD_UCM_STORE "build ucm store module." ON)
 option(BUILD_UCM_DRAMPOOL "build DramPool daemon." ON)
 option(BUILD_UCM_SPARSE "build ucm sparse module." OFF)
 option(BUILD_UCM_ASU "build ucm ASU transport module." OFF)
+option(BUILD_UCM_DELEGATOR "build ucm delegator module." OFF)
 option(BUILD_UCM_MINDIE "build ucm MindIE integration module." OFF)
 option(BUILD_UNIT_TESTS "build all unit test suits." OFF)
 option(BUILD_NUMA "build numactl library." OFF)

ucm/shared/trans/simu/simu_buffer.cc

@@ -56,6 +56,13 @@ std::shared_ptr<void> SimuBuffer::MakeDeviceMappedHostBuffer(size_t size)
     return MakeHostBuffer(size);
 }
 
+std::shared_ptr<void> SimuBuffer::MakeHostMappedDeviceBuffer(size_t size, void** pDevice)
+{
+    auto host = MakeHostBuffer(size);
+    if (pDevice) { *pDevice = host.get(); }
+    return host;
+}
+
 std::shared_ptr<void> SimuBuffer::MakeDeviceBuffer(size_t size)
 {
     constexpr int8_t deviceInitVal = 0xd;

ucm/shared/trans/simu/simu_buffer.h

@@ -32,7 +32,10 @@ class SimuBuffer : public ReservedBuffer {
 public:
     // Simu has no real device, so device-mapped-host memory is plain host memory.
     bool SupportsDeviceMappedHostBuffer() const override { return true; }
+    bool SupportsHostMappedDeviceBuffer() const override { return true; }
     std::shared_ptr<void> MakeDeviceMappedHostBuffer(size_t size) override;
+    std::shared_ptr<void> MakeHostMappedDeviceBuffer(size_t size,
+                                                     void** pDevice = nullptr) override;
     std::shared_ptr<void> MakeDeviceBuffer(size_t size) override;
     std::shared_ptr<void> MakeHostBuffer(size_t size) override;
 };

ucm/store/CMakeLists.txt

@@ -13,6 +13,8 @@ if(RUNTIME_ENVIRONMENT STREQUAL "ascend" OR RUNTIME_ENVIRONMENT STREQUAL "simu")
 endif()
 if(BUILD_UCM_ASU)
     add_subdirectory(asu)
+endif()
+if(BUILD_UCM_DELEGATOR)
     add_subdirectory(delegator)
 endif()
 add_subdirectory(empty)

ucm/store/asu/CMakeLists.txt

@@ -4,7 +4,7 @@ target_include_directories(asustore
     PUBLIC
         ${CMAKE_CURRENT_SOURCE_DIR}/cc
 )
-target_link_libraries(asustore PUBLIC storeintf asu_client infra_logger PRIVATE asu_ascend_deps trans)
+target_link_libraries(asustore PUBLIC storeintf asu_client infra_logger PRIVATE trans)
 set_target_properties(asustore PROPERTIES
     BUILD_WITH_INSTALL_RPATH TRUE
     BUILD_RPATH "$ORIGIN;$ORIGIN/../../transport/kv/asu"

ucm/store/asu/cc/asu_store.cc

@@ -313,7 +313,7 @@ class AsuStore final : public StoreV1 {
         for (std::size_t index = 0; index < count; ++index) {
             if (registrations[index].addr == 0 || registrations[index].size == 0) { continue; }
             UC::ASU::MemoryRegion region;
-            region.memoryType = UC::ASU::MemoryType::ASCEND_DEVICE;
+            region.memoryType = UC::ASU::MemoryType::DEVICE;
             region.addr = static_cast<std::uint64_t>(registrations[index].addr);
             region.size = static_cast<std::uint64_t>(registrations[index].size);
             region.deviceId = config_.deviceId;
@@ -727,7 +727,7 @@ class AsuStore final : public StoreV1 {
             for (std::size_t tensorIndex = 0; tensorIndex < shard.addrs.size(); ++tensorIndex) {
                 UC::ASU::KVBuffer entry;
                 entry.key = MakeAsuKey(shard.owner);
-                entry.buffer.region.memoryType = UC::ASU::MemoryType::ASCEND_DEVICE;
+                entry.buffer.region.memoryType = UC::ASU::MemoryType::DEVICE;
                 entry.buffer.region.addr =
                     reinterpret_cast<std::uint64_t>(shard.addrs[tensorIndex]);
                 entry.buffer.region.size = config_.tensorSizes[tensorIndex];

ucm/store/test/CMakeLists.txt

@@ -3,6 +3,8 @@ if(BUILD_UNIT_TESTS)
     file(GLOB_RECURSE UCMSTORE_TEST_SOURCE_FILES "./case/*.cc")
     if(NOT BUILD_UCM_ASU)
         list(FILTER UCMSTORE_TEST_SOURCE_FILES EXCLUDE REGEX ".*[/\\\\]asu[/\\\\].*")
+    endif()
+    if(NOT BUILD_UCM_DELEGATOR)
         list(FILTER UCMSTORE_TEST_SOURCE_FILES EXCLUDE REGEX ".*[/\\\\]delegator[/\\\\].*")
     endif()
     if(RUNTIME_ENVIRONMENT STREQUAL "ascend" OR RUNTIME_ENVIRONMENT STREQUAL "simu")
@@ -48,12 +50,16 @@ if(BUILD_UNIT_TESTS)
     if(BUILD_UCM_ASU)
         target_link_libraries(ucmstore.test PRIVATE
             asu_client
-            delegator
-            delegator_store
             buffer_pool
-            asu_ascend_deps
+            trans
         )
         target_compile_definitions(ucmstore.test PRIVATE ASU_BUILD_TESTS)
     endif()
+    if(BUILD_UCM_DELEGATOR)
+        target_link_libraries(ucmstore.test PRIVATE
+            delegator
+            delegator_store
+        )
+    endif()
     gtest_discover_tests(ucmstore.test)
 endif()

ucm/transport/kv/asu/CMakeLists.txt

@@ -1,54 +1,12 @@
-if(NOT RUNTIME_ENVIRONMENT STREQUAL "ascend")
-    message(FATAL_ERROR "BUILD_UCM_ASU requires RUNTIME_ENVIRONMENT=ascend. Current value: ${RUNTIME_ENVIRONMENT}")
+if(NOT (RUNTIME_ENVIRONMENT STREQUAL "ascend" OR RUNTIME_ENVIRONMENT STREQUAL "simu"))
+    message(FATAL_ERROR
+        "BUILD_UCM_ASU requires RUNTIME_ENVIRONMENT=ascend or simu. Current value: ${RUNTIME_ENVIRONMENT}")
 endif()
 
-if(NOT DEFINED ASCEND_ROOT)
-    if(DEFINED ENV{ASCEND_HOME_PATH})
-        set(ASCEND_ROOT "$ENV{ASCEND_HOME_PATH}" CACHE PATH "Path to Ascend root directory")
-    elseif(DEFINED ENV{ASCEND_TOOLKIT_HOME})
-        set(ASCEND_ROOT "$ENV{ASCEND_TOOLKIT_HOME}" CACHE PATH "Path to Ascend root directory")
-    else()
-        set(ASCEND_ROOT "/usr/local/Ascend/ascend-toolkit/latest" CACHE PATH "Path to Ascend root directory")
-    endif()
-endif()
-
-find_path(ASU_ASCEND_INCLUDE_DIR
-    NAMES acl/acl.h
-    HINTS
-        ${ASCEND_ROOT}/include
-        ${ASCEND_ROOT}/aarch64-linux/include
-        ${ASCEND_ROOT}/arm64-linux/include
-    NO_DEFAULT_PATH
-)
-if(NOT ASU_ASCEND_INCLUDE_DIR)
-    message(FATAL_ERROR "Cannot find acl/acl.h under ASCEND_ROOT=${ASCEND_ROOT}")
-endif()
-
-find_library(ASU_ASCENDCL_LIB
-    NAMES ascendcl
-    HINTS
-        ${ASCEND_ROOT}/lib64
-        ${ASCEND_ROOT}/aarch64-linux/lib64
-        ${ASCEND_ROOT}/aarch64-linux/devlib
-        ${ASCEND_ROOT}/arm64-linux/lib64
-        ${ASCEND_ROOT}/arm64-linux/devlib
-    NO_DEFAULT_PATH
-)
-if(NOT ASU_ASCENDCL_LIB)
-    message(FATAL_ERROR "Cannot find libascendcl under ASCEND_ROOT=${ASCEND_ROOT}")
-endif()
-
-message(STATUS "ASU ASCEND include: ${ASU_ASCEND_INCLUDE_DIR}")
-message(STATUS "ASU ASCENDCL lib: ${ASU_ASCENDCL_LIB}")
-
 option(BUILD_UCM_ASU_PROVIDER_AICPU "Build ASU AICPU trans provider" OFF)
 option(BUILD_UCM_ASU_PROVIDER_FAKE "Build ASU fake trans provider" ON)
 option(BUILD_UCM_ASU_PROVIDER_AIV "Build ASU AIV trans provider (requires external libumc.a)" OFF)
 
-add_library(asu_ascend_deps INTERFACE)
-target_include_directories(asu_ascend_deps INTERFACE ${ASU_ASCEND_INCLUDE_DIR})
-target_link_libraries(asu_ascend_deps INTERFACE ${ASU_ASCENDCL_LIB})
-
 file(GLOB ASU_COMMON_SOURCES CONFIGURE_DEPENDS common/*.cpp)
 file(GLOB ASU_TRANSPORT_SOURCES CONFIGURE_DEPENDS trans/src/*.cpp)
 file(GLOB LOGGER_SOURCES CONFIGURE_DEPENDS 
@@ -57,7 +15,6 @@ file(GLOB LOGGER_SOURCES CONFIGURE_DEPENDS
 )
 list(APPEND ASU_TRANSPORT_SOURCES ${ASU_COMMON_SOURCES})
 list(APPEND ASU_TRANSPORT_SOURCES ${LOGGER_SOURCES})
-list(APPEND ASU_TRANSPORT_SOURCES ${UCM_ROOT_DIR}/ucm/shared/trans/ascend/ascend_buffer.cc)
 add_library(asu_transport SHARED ${ASU_TRANSPORT_SOURCES})
 target_include_directories(asu_transport
     PUBLIC
@@ -74,9 +31,7 @@ target_include_directories(asu_transport
 target_compile_definitions(asu_transport PRIVATE SPDLOG_FMT_EXTERNAL)
 target_link_libraries(asu_transport
     PUBLIC
-        pthread fmt spdlog zlibstatic
-    PRIVATE
-        asu_ascend_deps
+        pthread fmt spdlog zlibstatic trans
 )
 
 if(BUILD_UCM_ASU_PROVIDER_AICPU)
@@ -150,7 +105,6 @@ if(BUILD_UNIT_TESTS)
     )
     target_link_libraries(asu.test PRIVATE
         asu_client
-        asu_ascend_deps
         gtest_main gtest
     )
     gtest_discover_tests(asu.test)

ucm/transport/kv/asu/client/src/asu_client_impl.cpp

@@ -250,7 +250,7 @@ Status AsuClientImpl::RegisterRegionsOnce(const std::vector<MemoryRegion>& regio
     std::vector<TransProvider::RegisterMemoryDesc> registerDescs;
     registerDescs.reserve(regions.size());
     for (const auto& region : regions) {
-        const auto memType = region.memoryType == MemoryType::ASCEND_DEVICE
+        const auto memType = region.memoryType == MemoryType::DEVICE
                                  ? TransProvider::MemType::MEM_DEVICE
                                  : TransProvider::MemType::MEM_HOST;
         registerDescs.push_back({memType, static_cast<std::uintptr_t>(region.addr),
@@ -624,7 +624,7 @@ Status AsuClientImpl::BindProviderRegions(const std::shared_ptr<TransProvider>&
     std::vector<TransProvider::BindMemoryDesc> bindDescs;
     bindDescs.reserve(registeredRegions.size());
     for (const auto& registeredRegion : registeredRegions) {
-        const auto memType = registeredRegion.region.memoryType == MemoryType::ASCEND_DEVICE
+        const auto memType = registeredRegion.region.memoryType == MemoryType::DEVICE
                                  ? TransProvider::MemType::MEM_DEVICE
                                  : TransProvider::MemType::MEM_HOST;
         bindDescs.push_back({memType, static_cast<std::uintptr_t>(registeredRegion.region.addr),

ucm/transport/kv/asu/test/transport/asu_submit_flow_test.cpp

@@ -21,7 +21,6 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  * */
-#include <acl/acl.h>
 #include <cstdint>
 #include <functional>
 #include <gtest/gtest.h>
@@ -33,6 +32,7 @@
 #include "asu_transport/trans_provider.h"
 #include "buffer_manager.h"
 #include "connection_internal.h"
+#include "trans/device.h"
 
 namespace UC::ASU {
 namespace {
@@ -134,14 +134,18 @@ class AsuSubmitFlowBufferTest : public ::testing::Test {
 protected:
     static void SetUpTestSuite()
     {
-        auto ret = aclInit(nullptr);
-        if (ret != ACL_SUCCESS && ret != ACL_ERROR_REPEAT_INITIALIZE) {
-            FAIL() << "aclInit failed: " << ret;
+        const auto initStatus = device_.Init();
+        if (initStatus.Failure() && initStatus != UC::Status::DuplicateKey()) {
+            FAIL() << "Device::Init failed: " << initStatus.ToString();
         }
-        ASSERT_EQ(aclrtSetDevice(0), ACL_SUCCESS);
+        ASSERT_TRUE(device_.Setup(0).Success());
     }
 
-    static void TearDownTestSuite() { aclrtResetDevice(0); }
+    static void TearDownTestSuite()
+    {
+        (void)device_.Reset(0);
+        (void)device_.Finalize();
+    }
 
     void SetUp() override
     {
@@ -151,6 +155,7 @@ class AsuSubmitFlowBufferTest : public ::testing::Test {
     }
 
     std::unique_ptr<AsuTransportImpl> transport_;
+    static inline Trans::Device device_;
 };
 
 }  // namespace
@@ -161,18 +166,20 @@ class AsuTransportBufferRegistrationTest : public ::testing::Test {
 protected:
     static void SetUpTestSuite()
     {
-        const auto ret = aclInit(nullptr);
-        if (ret != ACL_SUCCESS && ret != ACL_ERROR_REPEAT_INITIALIZE) {
-            FAIL() << "aclInit failed: " << ret;
+        const auto initStatus = device_.Init();
+        if (initStatus.Failure() && initStatus != UC::Status::DuplicateKey()) {
+            FAIL() << "Device::Init failed: " << initStatus.ToString();
         }
-        ASSERT_EQ(aclrtSetDevice(0), ACL_SUCCESS);
+        ASSERT_TRUE(device_.Setup(0).Success());
     }
 
     static void TearDownTestSuite()
     {
-        EXPECT_EQ(aclrtResetDevice(0), ACL_SUCCESS);
-        EXPECT_EQ(aclFinalize(), ACL_SUCCESS);
+        (void)device_.Reset(0);
+        (void)device_.Finalize();
     }
+
+    static inline Trans::Device device_;
 };
 
 TEST_F(AsuTransportBufferRegistrationTest, InitRegistersAndShutdownUnregistersBothBuffers)