fix USE_ASCEND_DIRECT compile

Author: youxiao

mooncake-common/common.cmake

@@ -160,14 +160,15 @@ if (USE_ASCEND OR USE_ASCEND_DIRECT)
     file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
   endif()
   set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
-  set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib")
   set(ASCEND_INCLUDE_DIR "${ASCEND_TOOLKIT_ROOT}/include")
   add_compile_options(-Wno-ignored-qualifiers)
-  include_directories(/usr/local/include /usr/include ${ASCEND_INCLUDE_DIR})
-  link_directories(${ASCEND_LIB_DIR} ${ASCEND_DEVLIB_DIR})
+  include_directories(/usr/local/include /usr/include ${ASCEND_INCLUDE_DIR} ${ASCEND_INCLUDE_DIR}/experiment/runtime)
+  link_directories(${ASCEND_LIB_DIR})
 endif()
 
 if (USE_ASCEND)
+  set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib")
+  link_directories(${ASCEND_DEVLIB_DIR})
   add_compile_definitions(USE_ASCEND)
 endif()
 

mooncake-store/include/utils.h

@@ -102,7 +102,7 @@ void* allocate_buffer_allocator_memory(
     size_t total_size, const std::string& protocol = "",
     size_t alignment = facebook::cachelib::Slab::kSize);
 
-void free_memory(const std::string& protocol, void* ptr);
+void free_memory(const std::string& protocol, void* ptr, size_t size);
 
 [[nodiscard]] inline std::string byte_size_to_string(uint64_t bytes) {
     const double KB = 1024.0;

mooncake-store/src/client_buffer.cpp

@@ -24,7 +24,7 @@ ClientBufferAllocator::ClientBufferAllocator(size_t size,
 ClientBufferAllocator::~ClientBufferAllocator() {
     // Free the aligned allocated memory
     if (buffer_) {
-        free_memory(protocol, buffer_);
+        free_memory(protocol, buffer_, buffer_size_);
     }
 }
 

mooncake-store/src/utils.cpp

@@ -92,9 +92,9 @@ void *allocate_buffer_allocator_memory(size_t total_size,
     return aligned_alloc(alignment, total_size);
 }
 
-void free_memory(const std::string &protocol, void *ptr) {
+void free_memory(const std::string &protocol, void *ptr, size_t size) {
 #ifdef USE_ASCEND_DIRECT
-    if (protocol == "ascend") {
+    if (protocol == "ascend" && size > 0) {
         aclrtFreeHost(ptr);
         return;
     }

mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h

@@ -28,7 +28,8 @@
 #include <condition_variable>
 #include <set>
 
-#include <acl/acl.h>
+#include "runtime/rt.h"
+//#include <acl/acl.h>
 #include "transfer_metadata.h"
 #include "transport/transport.h"
 #include "adxl/adxl_engine.h"
@@ -81,21 +82,26 @@ class AscendDirectTransport : public Transport {
 
     void processSliceList(const std::vector<Slice *> &slice_list);
 
+    void connectAndTransfer(const std::string &target_adxl_engine_name,
+                            adxl::TransferOp operation,
+                            const std::vector<Slice *> &slice_list,
+                            int32_t times = 0);
+
     void localCopy(TransferRequest::OpCode opcode,
                    const std::vector<Slice *> &slice_list);
 
-    aclError copyWithBatch(TransferRequest::OpCode opcode,
-                           const std::vector<Slice *> &slice_list,
-                           aclrtMemcpyKind kind, size_t batch_num,
-                           size_t slice_index) const;
-
-    void copyWithSync(TransferRequest::OpCode opcode,
-                      const std::vector<Slice *> &slice_list,
-                      aclrtMemcpyKind kind);
-
-    void copyWithAsync(TransferRequest::OpCode opcode,
-                       const std::vector<Slice *> &slice_list,
-                       aclrtMemcpyKind kind);
+//    aclError copyWithBatch(TransferRequest::OpCode opcode,
+//                           const std::vector<Slice *> &slice_list,
+//                           aclrtMemcpyKind kind, size_t batch_num,
+//                           size_t slice_index) const;
+//
+//    void copyWithSync(TransferRequest::OpCode opcode,
+//                      const std::vector<Slice *> &slice_list,
+//                      aclrtMemcpyKind kind);
+//
+//    void copyWithAsync(TransferRequest::OpCode opcode,
+//                       const std::vector<Slice *> &slice_list,
+//                       aclrtMemcpyKind kind);
 
     uint16_t findAdxlListenPort() const;
 
@@ -105,7 +111,7 @@ class AscendDirectTransport : public Transport {
     int checkAndConnect(const std::string &target_adxl_engine_name);
 
     int disconnect(const std::string &target_adxl_engine_name,
-                   int32_t timeout_in_millis);
+                   int32_t timeout_in_millis, bool force = false);
 
     std::atomic_bool running_;
     std::unique_ptr<adxl::AdxlEngine> adxl_;
@@ -123,11 +129,12 @@ class AscendDirectTransport : public Transport {
     std::condition_variable queue_cv_;
 
     int32_t device_logic_id_{};
-    aclrtContext rt_context_{nullptr};
+    rtContext_t rt_context_{};
+//    aclrtContext rt_context_{nullptr};
     int32_t connect_timeout_ = 10000;
     int32_t transfer_timeout_ = 10000;
     std::string local_adxl_engine_name_{};
-    aclrtStream stream_{};
+//    aclrtStream stream_{};
     bool use_buffer_pool_{false};
 
     int32_t base_port_ = 20000;