ascend enable fabric mem

Author: youxiao

mooncake-store/src/client_buffer.cpp

@@ -27,6 +27,9 @@ std::shared_ptr<ClientBufferAllocator> ClientBufferAllocator::create(
 ClientBufferAllocator::ClientBufferAllocator(size_t size,
                                              const std::string& protocol)
     : protocol(protocol), buffer_size_(size) {
+    if (size == 0) {
+        return;
+    }
     // Align to 64 bytes(cache line size) for better cache performance
     constexpr size_t alignment = 64;
     buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment);
@@ -55,6 +58,9 @@ ClientBufferAllocator::~ClientBufferAllocator() {
 }
 
 std::optional<BufferHandle> ClientBufferAllocator::allocate(size_t size) {
+    if (allocator_ == nullptr) {
+        return std::nullopt;
+    }
     auto handle = allocator_->allocate(size);
     if (!handle) {
         return std::nullopt;

mooncake-store/src/client_service.cpp

@@ -261,6 +261,13 @@ ErrorCode Client::InitTransferEngine(
         }
     }
 
+    if (protocol == "ascend") {
+        const char* ascend_use_fabric_mem =
+            std::getenv("ASCEND_ENABLE_USE_FABRIC_MEM");
+        if (ascend_use_fabric_mem) {
+            globalConfig().ascend_use_fabric_mem = true;
+        }
+    }
     auto [hostname, port] = parseHostNameWithPort(local_hostname);
     int rc = transfer_engine_->init(metadata_connstring, local_hostname,
                                     hostname, port);

mooncake-store/src/utils.cpp

@@ -10,6 +10,7 @@
 #include <random>
 #ifdef USE_ASCEND_DIRECT
 #include "acl/acl.h"
+#include "config.h"
 #endif
 
 #include <ylt/coro_http/coro_http_client.hpp>
@@ -79,13 +80,41 @@ void *allocate_buffer_allocator_memory(size_t total_size,
     }
 #ifdef USE_ASCEND_DIRECT
     if (protocol == "ascend" && total_size > 0) {
-        void *buffer = nullptr;
-        auto ret = aclrtMallocHost(&buffer, total_size);
-        if (ret != ACL_ERROR_NONE) {
-            LOG(ERROR) << "Failed to allocate memory: " << ret;
-            return nullptr;
+        if (globalConfig().ascend_use_fabric_mem) {
+            aclrtDrvMemHandle handle = nullptr;
+            aclrtPhysicalMemProp prop = {};
+            prop.handleType = ACL_MEM_HANDLE_TYPE_NONE;
+            prop.allocationType = ACL_MEM_ALLOCATION_TYPE_PINNED;
+            prop.memAttr = ACL_HBM_MEM_HUGE;
+            prop.location.type = ACL_MEM_LOCATION_TYPE_HOST;
+            prop.location.id = 0;
+            prop.reserve = 0;
+            auto ret = aclrtMallocPhysical(&handle, total_size, &prop, 0);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to allocate memory: " << ret;
+                return nullptr;
+            }
+            void *va;
+            ret = aclrtReserveMemAddress(&va, total_size, 0, nullptr, 1);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to reserve memory: " << ret;
+                return nullptr;
+            }
+            ret = aclrtMapMem(va, total_size, 0, handle, 0);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to map memory: " << ret;
+                return nullptr;
+            }
+            return va;
+        } else {
+            void *buffer = nullptr;
+            auto ret = aclrtMallocHost(&buffer, total_size);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to allocate memory: " << ret;
+                return nullptr;
+            }
+            return buffer;
         }
-        return buffer;
     }
 #endif
     // Allocate aligned memory
@@ -95,7 +124,29 @@ void *allocate_buffer_allocator_memory(size_t total_size,
 void free_memory(const std::string &protocol, void *ptr) {
 #ifdef USE_ASCEND_DIRECT
     if (protocol == "ascend") {
-        aclrtFreeHost(ptr);
+        if (globalConfig().ascend_use_fabric_mem) {
+            auto ret = aclrtUnmapMem(ptr);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to unmap memory: " << ptr;
+            }
+            aclrtReleaseMemAddress(ptr);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to release mem address: " << ptr;
+            }
+            aclrtDrvMemHandle handle;
+            ret = aclrtMemRetainAllocationHandle(ptr, &handle);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to retain allocation handle: " << ptr;
+                return;
+            }
+            ret = aclrtFreePhysical(handle);
+            if (ret != ACL_ERROR_NONE) {
+                LOG(ERROR) << "Failed to free physical handle: " << handle;
+                return;
+            }
+        } else {
+            aclrtFreeHost(ptr);
+        }
         return;
     }
 #endif

mooncake-transfer-engine/include/config.h

@@ -58,6 +58,7 @@ struct GlobalConfig {
     bool enable_dest_device_affinity = false;
     size_t eic_max_block_size = 64UL * 1024 * 1024;
     EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE;
+    bool ascend_use_fabric_mem = false;
 };
 
 void loadGlobalConfig(GlobalConfig &config);

mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp

@@ -32,6 +32,7 @@
 #include "transfer_metadata.h"
 #include "transfer_metadata_plugin.h"
 #include "transport/transport.h"
+#include "config.h"
 
 namespace mooncake {
 namespace {
@@ -67,18 +68,6 @@ AscendDirectTransport::~AscendDirectTransport() {
         }
         connected_segments_.clear();
     }
-
-    // Deregister all memory
-    std::lock_guard<std::mutex> mem_handle_lock(mem_handle_mutex_);
-    for (const auto &[addr, mem_handle] : addr_to_mem_handle_) {
-        auto status = adxl_->DeregisterMem(mem_handle);
-        if (status != adxl::SUCCESS) {
-            LOG(ERROR) << "Failed to deregister memory at address " << addr;
-        } else {
-            LOG(INFO) << "Deregistered memory at address " << addr;
-        }
-    }
-    addr_to_mem_handle_.clear();
     adxl_->Finalize();
 }
 
@@ -169,6 +158,12 @@ int AscendDirectTransport::InitAdxlEngine() {
             use_buffer_pool_ = false;
         }
     }
+    if (globalConfig().ascend_use_fabric_mem) {
+        options["EnableUseFabricMem"] = "1";
+        use_buffer_pool_ = false;
+        LOG(INFO) << "Use fabric mem is enabled.";
+        globalConfig().ascend_use_fabric_mem = false;
+    }
     auto adxl_engine_name =
         adxl::AscendString((host_ip + ":" + std::to_string(host_port)).c_str());
     auto status = adxl_->Initialize(adxl_engine_name, options);
@@ -178,7 +173,7 @@ int AscendDirectTransport::InitAdxlEngine() {
     }
     LOG(INFO) << "Success to initialize adxl engine:"
               << adxl_engine_name.GetString()
-              << " with device_id:" << device_logic_id_;
+              << " with device_id:" << device_logic_id_ << ", pid:" << getpid();
     char *connect_timeout_str = std::getenv("ASCEND_CONNECT_TIMEOUT");
     if (connect_timeout_str) {
         std::optional<int32_t> connect_timeout =