Add Cache IO aggregation

Author: naganomei

CMakeLists.txt

@@ -13,6 +13,7 @@ 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)
 option(UCM_ENABLE_GDR_STREAM "build GPUDirect RDMA stream support for CUDA runtime." OFF)
+option(UCM_ENABLE_ASCEND_IO_AGGREGATION "build Ascend Cache IO aggregation support." OFF)
 option(ASCEND_SUPPORTS_REGISTER_PIN "enable Ascend register pin optimization (requires CANN >= 8.5)" OFF)
 option(DOWNLOAD_DEPENDENCE "download dependence by cmake." ON)
 set(RUNTIME_ENVIRONMENT "simu" CACHE STRING "runtime: simu, ascend, musa or cuda.")

examples/ucm_config_example.yaml

@@ -24,6 +24,11 @@ ucm_connectors:
       # and building UCM with ENABLE_GDR=1.
       use_gdr: false
 
+      # Enable Cache IO aggregation for Ascend small-IO CacheStore H2D and D2H copies.
+      # This aggregates one shard into one IO object to reduce many small transfers.
+      # This requires an Ascend runtime and building UCM with UCM_ENABLE_ASCEND_IO_AGGREGATION=1.
+      cache_io_aggregation: false
+
 # When you use UcmNfsStore, you should set enable_event_sync to false.
 enable_event_sync: true
 # Enable UCM metrics so they can be monitored online via Grafana and Prometheus.

setup.py

@@ -35,7 +35,14 @@
 ENABLE_SPARSE = os.getenv("ENABLE_SPARSE")
 ENABLE_MINDIE = os.getenv("UCM_ENABLE_MINDIE", "0") not in ("", "0", "false", "False")
 ENABLE_GDR = os.getenv("ENABLE_GDR", "0") not in ("", "0", "false", "False")
-
+ENABLE_ASCEND_IO_AGGREGATION = os.getenv(
+    "UCM_ENABLE_ASCEND_IO_AGGREGATION", "0"
+) not in (
+    "",
+    "0",
+    "false",
+    "False",
+)
 
 def get_abi_flag_from_env() -> str:
     v = os.environ.get("UCM_CXX11_ABI")
@@ -171,6 +178,10 @@ def build_cmake(self, ext: CMakeExtension):
         if ENABLE_GDR:
             cmake_args += ["-DUCM_ENABLE_GDR_STREAM=ON"]
 
+        cmake_args += [
+            f"-DUCM_ENABLE_ASCEND_IO_AGGREGATION={'ON' if ENABLE_ASCEND_IO_AGGREGATION else 'OFF'}"
+        ]
+
         match PLATFORM:
             case "cuda":
                 cmake_args += ["-DRUNTIME_ENVIRONMENT=cuda"]

ucm/shared/trans/ascend/CMakeLists.txt

@@ -12,6 +12,54 @@ add_library(trans STATIC
 if(ASCEND_SUPPORTS_REGISTER_PIN)
     target_compile_definitions(trans PRIVATE ASCEND_SUPPORTS_REGISTER_PIN=1)
 endif()
+if(UCM_ENABLE_ASCEND_IO_AGGREGATION)
+    find_library(ASCEND_RUNTIME_LIBRARY runtime
+        PATHS
+            "${ASCEND_ROOT}/lib64"
+            "${ASCEND_ROOT}/lib"
+            "${ASCEND_ROOT}/runtime/lib64"
+            "${ASCEND_ROOT}/runtime/lib"
+            "${ASCEND_ROOT}/aarch64-linux/lib64"
+            "${ASCEND_ROOT}/aarch64-linux/lib"
+        NO_DEFAULT_PATH
+    )
+    find_path(ASCEND_FFTS_INCLUDE_DIR
+        NAMES runtime/rt_ffts_plus.h rt_external_ffts.h
+        PATHS
+            "${ASCEND_ROOT}/include"
+            "${ASCEND_ROOT}/pkg_inc"
+            "${ASCEND_ROOT}/pkg_inc/runtime"
+            "${ASCEND_ROOT}/aarch64-linux/pkg_inc"
+            "${ASCEND_ROOT}/aarch64-linux/pkg_inc/runtime"
+        NO_DEFAULT_PATH
+    )
+    if(NOT ASCEND_RUNTIME_LIBRARY OR NOT ASCEND_FFTS_INCLUDE_DIR)
+        message(FATAL_ERROR "UCM_ENABLE_ASCEND_IO_AGGREGATION requires FFTS headers and libruntime.")
+    endif()
+    get_filename_component(ASCEND_FFTS_INCLUDE_PARENT "${ASCEND_FFTS_INCLUDE_DIR}" DIRECTORY)
+    set(ASCEND_FFTS_INCLUDE_DIRS "${ASCEND_FFTS_INCLUDE_DIR}")
+    if(ASCEND_FFTS_INCLUDE_DIR MATCHES "/runtime$" AND EXISTS "${ASCEND_FFTS_INCLUDE_PARENT}")
+        list(APPEND ASCEND_FFTS_INCLUDE_DIRS "${ASCEND_FFTS_INCLUDE_PARENT}")
+    endif()
+    list(APPEND ASCEND_FFTS_INCLUDE_DIRS
+        "${ASCEND_ROOT}/pkg_inc"
+        "${ASCEND_ROOT}/pkg_inc/toolchain"
+        "${ASCEND_ROOT}/pkg_inc/profiling"
+        "${ASCEND_ROOT}/aarch64-linux/pkg_inc"
+        "${ASCEND_ROOT}/aarch64-linux/pkg_inc/toolchain"
+        "${ASCEND_ROOT}/aarch64-linux/pkg_inc/profiling"
+    )
+    list(REMOVE_DUPLICATES ASCEND_FFTS_INCLUDE_DIRS)
+    target_sources(trans PRIVATE
+        ascend_shard_io_aggregator.cc
+        ffts_d2d_dispatcher.cc
+    )
+    target_include_directories(trans PUBLIC ${ASCEND_FFTS_INCLUDE_DIRS})
+    target_compile_definitions(trans PUBLIC UCM_ENABLE_ASCEND_IO_AGGREGATION=1)
+    target_link_libraries(trans PUBLIC ${ASCEND_RUNTIME_LIBRARY})
+else()
+    target_compile_definitions(trans PUBLIC UCM_ENABLE_ASCEND_IO_AGGREGATION=0)
+endif()
 target_link_libraries(trans PUBLIC
     fmt
     Ascend::ascendcl