add cmo code

integrations/vllm_ascend/tilexr_collectives/__init__.py

@@ -7,11 +7,23 @@
     TILEXR_DATA_TYPE_INT64,
     TILEXR_DATA_TYPE_UINT8,
     TILEXR_REDUCE_SUM,
+    TILEXR_CMO_FLUSH,
+    TILEXR_CMO_INVALIDATE,
+    TILEXR_CMO_PREFETCH,
     TILEXR_SUCCESS,
     TileXRCollectivesError,
     TileXRCollectivesRuntime,
 )
-from .torch_collectives import all_gather, all_reduce, all_to_all, broadcast, reduce_scatter
+from .torch_collectives import (
+    TileXRCmoMeta,
+    all_gather,
+    all_reduce,
+    all_to_all,
+    broadcast,
+    clear_cmo_meta,
+    mark_cmo,
+    reduce_scatter,
+)
 from .vllm_adapter import TileXRVllmCollectivesAdapter
 from .vllm_adapter import enabled as tilexr_vllm_collectives_enabled
 from .vllm_patch import patch_npu_communicator
@@ -25,14 +37,20 @@
     "TILEXR_DATA_TYPE_INT64",
     "TILEXR_DATA_TYPE_UINT8",
     "TILEXR_REDUCE_SUM",
+    "TILEXR_CMO_FLUSH",
+    "TILEXR_CMO_INVALIDATE",
+    "TILEXR_CMO_PREFETCH",
     "TILEXR_SUCCESS",
     "TileXRCollectivesError",
     "TileXRCollectivesRuntime",
+    "TileXRCmoMeta",
     "TileXRVllmCollectivesAdapter",
     "all_gather",
     "all_reduce",
     "all_to_all",
     "broadcast",
+    "clear_cmo_meta",
+    "mark_cmo",
     "patch_npu_communicator",
     "reduce_scatter",
     "tilexr_vllm_collectives_enabled",

integrations/vllm_ascend/tilexr_collectives/runtime.py

@@ -14,6 +14,9 @@
 TILEXR_DATA_TYPE_UINT8 = 7
 TILEXR_DATA_TYPE_BFP16 = 11
 TILEXR_REDUCE_SUM = 0
+TILEXR_CMO_PREFETCH = 0
+TILEXR_CMO_FLUSH = 1
+TILEXR_CMO_INVALIDATE = 2
 
 
 class TileXRCollectivesError(RuntimeError):
@@ -103,6 +106,18 @@ def _configure_symbols(self) -> None:
         self._comm_lib.TileXRCommInitRankLocal.restype = ctypes.c_int
         self._comm_lib.TileXRCommDestroy.argtypes = [ctypes.c_void_p]
         self._comm_lib.TileXRCommDestroy.restype = ctypes.c_int
+        self._comm_lib.TileXRSubmitCmoTask.argtypes = [
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_uint64,
+            ctypes.c_uint32,
+            ctypes.c_uint32,
+            ctypes.c_uint32,
+            ctypes.c_uint64,
+        ]
+        self._comm_lib.TileXRSubmitCmoTask.restype = ctypes.c_int
+        self._comm_lib.TileXRClearCmoTask.argtypes = [ctypes.c_void_p]
+        self._comm_lib.TileXRClearCmoTask.restype = ctypes.c_int
 
         self._collectives_lib.TileXRAllGather.argtypes = [
             ctypes.c_void_p,
@@ -258,6 +273,30 @@ def broadcast(
         )
         self._check("TileXRBroadcast", ret, f"rank={self.rank} count={count} dtype={tilexr_dtype} root={root}")
 
+    def submit_cmo_task(
+        self,
+        target_ptr: int,
+        total_bytes: int,
+        op_type: int = TILEXR_CMO_PREFETCH,
+        priority: int = 0,
+        chunk_bytes: int = 0,
+        expire_seq: int = 0,
+    ) -> None:
+        ret = self._comm_lib.TileXRSubmitCmoTask(
+            self._comm,
+            _void_p(target_ptr),
+            ctypes.c_uint64(int(total_bytes)),
+            ctypes.c_uint32(int(op_type)),
+            ctypes.c_uint32(int(priority)),
+            ctypes.c_uint32(int(chunk_bytes)),
+            ctypes.c_uint64(int(expire_seq)),
+        )
+        self._check("TileXRSubmitCmoTask", ret, f"rank={self.rank} bytes={total_bytes} op={op_type}")
+
+    def clear_cmo_task(self) -> None:
+        ret = self._comm_lib.TileXRClearCmoTask(self._comm)
+        self._check("TileXRClearCmoTask", ret, f"rank={self.rank}")
+
     def close(self) -> None:
         if self._closed:
             return

integrations/vllm_ascend/tilexr_collectives/torch_collectives.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from functools import lru_cache
 
 from .runtime import (
@@ -9,10 +10,88 @@
     TILEXR_DATA_TYPE_INT8,
     TILEXR_DATA_TYPE_INT32,
     TILEXR_DATA_TYPE_INT64,
+    TILEXR_CMO_PREFETCH,
     TileXRCollectivesRuntime,
 )
 
 
+_CMO_META_ATTR = "_tilexr_cmo_meta"
+_CMO_META_BY_ID: dict[int, "TileXRCmoMeta"] = {}
+
+
+@dataclass(frozen=True)
+class TileXRCmoMeta:
+    target_ptr: int
+    total_bytes: int
+    op_type: int = TILEXR_CMO_PREFETCH
+    priority: int = 0
+    chunk_bytes: int = 0
+    expire_seq: int = 0
+
+
+def mark_cmo(
+    tensor,
+    *,
+    target_tensor=None,
+    target_ptr: int | None = None,
+    total_bytes: int | None = None,
+    op_type: int = TILEXR_CMO_PREFETCH,
+    priority: int = 0,
+    chunk_bytes: int = 0,
+    expire_seq: int = 0,
+):
+    if target_tensor is not None and target_ptr is not None:
+        raise ValueError("target_tensor and target_ptr cannot both be set")
+    if target_tensor is None and target_ptr is None:
+        target_tensor = tensor
+    if target_tensor is not None:
+        _validate_npu_contiguous(target_tensor, "target_tensor")
+        target_ptr = int(target_tensor.data_ptr())
+        if total_bytes is None:
+            total_bytes = int(target_tensor.numel() * target_tensor.element_size())
+    if target_ptr is None or total_bytes is None or int(total_bytes) <= 0:
+        raise ValueError("CMO target_ptr and positive total_bytes are required")
+    meta = TileXRCmoMeta(
+        target_ptr=int(target_ptr),
+        total_bytes=int(total_bytes),
+        op_type=int(op_type),
+        priority=int(priority),
+        chunk_bytes=int(chunk_bytes),
+        expire_seq=int(expire_seq),
+    )
+    try:
+        setattr(tensor, _CMO_META_ATTR, meta)
+    except (AttributeError, TypeError):
+        _CMO_META_BY_ID[id(tensor)] = meta
+    return tensor
+
+
+def clear_cmo_meta(tensor):
+    if hasattr(tensor, _CMO_META_ATTR):
+        delattr(tensor, _CMO_META_ATTR)
+    _CMO_META_BY_ID.pop(id(tensor), None)
+    return tensor
+
+
+def _get_cmo_meta(tensor) -> TileXRCmoMeta | None:
+    return getattr(tensor, _CMO_META_ATTR, None) or _CMO_META_BY_ID.get(id(tensor))
+
+
+def _sync_cmo_task(runtime: TileXRCollectivesRuntime, tensor) -> None:
+    meta = _get_cmo_meta(tensor)
+    if meta is None:
+        runtime.clear_cmo_task()
+        return
+    runtime.submit_cmo_task(
+        target_ptr=meta.target_ptr,
+        total_bytes=meta.total_bytes,
+        op_type=meta.op_type,
+        priority=meta.priority,
+        chunk_bytes=meta.chunk_bytes,
+        expire_seq=meta.expire_seq,
+    )
+
+
 def _torch():
     import torch
 
@@ -139,6 +218,7 @@ def all_reduce(tensor, rank: int, world_size: int, install_prefix: str, runtime=
     device_index = _npu_device_index(tensor)
     _bind_npu_device(device_index)
     rt = _select_runtime(runtime, rank, world_size, install_prefix, device_index)
+    _sync_cmo_task(rt, tensor)
     output = torch.empty_like(tensor)
     rt.all_reduce(
         send_ptr=tensor.data_ptr(),

src/collectives/host/perf_trace_report.cpp

@@ -158,12 +158,23 @@ std::string EscapeHtml(const std::string &value)
     return out.str();
 }
 
-std::vector<TileXR::TileXRPerfCoreStageStats> NonEmptyStats(
+bool IsValidPerfStat(const TileXR::TileXRPerfTraceHeader &header,
+                     const TileXR::TileXRPerfCoreStageStats &stat)
+{
+    return stat.count != 0 &&
+        stat.rank < header.rankSize &&
+        stat.core < header.maxCoreCount &&
+        stat.stageId < header.stageCount &&
+        stat.stageId < TileXR::TILEXR_PERF_STAGE_COUNT;
+}
+
+std::vector<TileXR::TileXRPerfCoreStageStats> ValidStats(
+    const TileXR::TileXRPerfTraceHeader &header,
     const std::vector<TileXR::TileXRPerfCoreStageStats> &stats)
 {
     std::vector<TileXR::TileXRPerfCoreStageStats> result;
     for (const auto &stat : stats) {
-        if (stat.count != 0) {
+        if (IsValidPerfStat(header, stat)) {
             result.push_back(stat);
         }
     }
@@ -188,7 +199,7 @@ std::string BuildTraceJson(const TileXR::TileXRPerfTraceHeader &header,
                            const std::vector<TileXR::TileXRPerfCoreStageStats> &stats,
                            const PerfReportOptions &options)
 {
-    const auto nonEmptyStats = NonEmptyStats(stats);
+    const auto nonEmptyStats = ValidStats(header, stats);
     std::ostringstream out;
     out << "{\n";
     out << "  \"schema\": \"tilexr_perf_trace_report.v1\",\n";
@@ -291,7 +302,7 @@ std::string BuildHtmlReport(const TileXR::TileXRPerfTraceHeader &header,
                             const std::vector<TileXR::TileXRPerfCoreStageStats> &stats,
                             const PerfReportOptions &options)
 {
-    const auto nonEmptyStats = NonEmptyStats(stats);
+    const auto nonEmptyStats = ValidStats(header, stats);
     std::ostringstream out;
     out << "<!doctype html>\n<html><head><meta charset=\"utf-8\">";
     out << "<title>TileXR Collective Perf Report</title>";
@@ -414,7 +425,7 @@ std::vector<PerfStageSummary> SummarizePerfTrace(
     }
 
     for (const auto &stat : stats) {
-        if (stat.count == 0 || stat.stageId >= stageCount) {
+        if (!IsValidPerfStat(header, stat) || stat.stageId >= stageCount) {
             continue;
         }
 

src/collectives/kernels/CMakeLists.txt

@@ -53,6 +53,8 @@ include_directories(
         ${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/runtime/
         ${ASCEND_HOME_PATH}/${ARCH}-linux/include/
         ${ASCEND_HOME_PATH}/${ARCH}-linux/include/ascendc
+        ${ASCEND_HOME_PATH}/${ARCH}-linux/asc/include
+        ${ASCEND_HOME_PATH}/${ARCH}-linux/asc/impl
         ${ASCEND_HOME_PATH}/${ARCH}-linux/ascendc/include
         ${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw/
         ${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw/interface/
@@ -82,6 +84,7 @@ add_custom_command(
                 --static -o "${TILEXR_COLLECTIVES_OP}" --allow-multiple-definition
         COMMAND truncate -c -s ${TILEXR_COLLECTIVES_1OP_BIN_SIZE} "${TILEXR_COLLECTIVES_OP}"
         DEPENDS tilexr_collectives_op_tmp
+                "${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/tilexr_collectives_op_tmp.dir/tilexr_lccl_op.cpp.o"
         VERBATIM
 )
 

src/collectives/kernels/kernels/collectives.cce

@@ -20,6 +20,7 @@
 #include <type_traits>
 #include "kernel_operator.h"
 #include "comm_args.h"
+#include "tilexr_cmo.h"
 #include "perf_trace_kernel.h"
 #include "tilexr_types.h"
 #include "../datacopy_gm2gm.h"
@@ -72,6 +73,12 @@ magic, op, root, localRankSize, 0, nullptr, dumpAddr, shareAddrs[0], shareAddrs[
 shareAddrs[3], shareAddrs[4], shareAddrs[5], shareAddrs[6], shareAddrs[7], shareAddrs[8], shareAddrs[9], \
 shareAddrs[10], shareAddrs[11], shareAddrs[12], shareAddrs[13], shareAddrs[14], shareAddrs[15], perfTrace
 
+#define ALLREDUCE_ARGS_FUN_16P_CMO(T) \
+ALLREDUCE_ARGS_FUN_16P(T), GM_ADDR cmoTaskPtr
+
+#define ALLREDUCE_ARGS_CALL_16P_CMO(type) \
+ALLREDUCE_ARGS_CALL_16P(type), reinterpret_cast<__gm__ CommArgs *>(commArgs)->cmoTaskPtr
+
 #define ALLREDUCE_ARGS_FUN_16P_Origin(T) \
 __gm__ T *input, __gm__ T *output, int rank, int rankSize, int64_t len, int64_t magic, int op, int root, \
 int localRankSize, __gm__ int64_t *sendCountMatrix, GM_ADDR dumpAddr, __gm__ T* buff[MAX_RANK_NUM_OF_ONE_910B2C], \
@@ -648,7 +655,6 @@ __attribute__((always_inline)) inline __aicore__ void ProcessDataNew(int64_t dat
     return;
 }
 
-
 template <typename T>
 __attribute__((always_inline)) inline __aicore__ void ProcessDataNewNonBarrier(int64_t dataSizeRemain, __ubuf__ T *inputUB[2],
     __gm__ T *buff, int64_t dataOffsetNum, int64_t buffOffsetNum, __gm__ T *output, int64_t outputOffsetNum, int op)

src/collectives/kernels/kernels/lcal_allreduce_2npu_write.cce

@@ -58,4 +58,4 @@ __attribute__((always_inline)) inline __aicore__ void TileXRAllReduce2npuWrite(A
     ProcessData<T>(dataSizeRemain, inputUB[0], buff[rank], dataOffsetNum, buffOffsetNum, output, buffOffsetNum, op);
     DumpLcclLogInfo(dumpAddr, LogId::PROCESS, static_cast<Op>(op));
     DumpLcclLogInfo(dumpAddr, LogId::OVERALL, static_cast<Op>(op));
-}
+}