diff --git a/integrations/vllm_ascend/tilexr_collectives/__init__.py b/integrations/vllm_ascend/tilexr_collectives/__init__.py index 84b6283c..e45739df 100644 --- a/integrations/vllm_ascend/tilexr_collectives/__init__.py +++ b/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", diff --git a/integrations/vllm_ascend/tilexr_collectives/runtime.py b/integrations/vllm_ascend/tilexr_collectives/runtime.py index 1b80d6e1..d68958a5 100644 --- a/integrations/vllm_ascend/tilexr_collectives/runtime.py +++ b/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 diff --git a/integrations/vllm_ascend/tilexr_collectives/torch_collectives.py b/integrations/vllm_ascend/tilexr_collectives/torch_collectives.py index 51b4a3ad..84a50e6c 100644 --- a/integrations/vllm_ascend/tilexr_collectives/torch_collectives.py +++ b/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(), diff --git a/src/collectives/host/perf_trace_report.cpp b/src/collectives/host/perf_trace_report.cpp index 3c5da8fb..24f907fa 100644 --- a/src/collectives/host/perf_trace_report.cpp +++ b/src/collectives/host/perf_trace_report.cpp @@ -158,12 +158,23 @@ std::string EscapeHtml(const std::string &value) return out.str(); } -std::vector 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 ValidStats( + const TileXR::TileXRPerfTraceHeader &header, const std::vector &stats) { std::vector 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 &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 &stats, const PerfReportOptions &options) { - const auto nonEmptyStats = NonEmptyStats(stats); + const auto nonEmptyStats = ValidStats(header, stats); std::ostringstream out; out << "\n"; out << "TileXR Collective Perf Report"; @@ -414,7 +425,7 @@ std::vector SummarizePerfTrace( } for (const auto &stat : stats) { - if (stat.count == 0 || stat.stageId >= stageCount) { + if (!IsValidPerfStat(header, stat) || stat.stageId >= stageCount) { continue; } diff --git a/src/collectives/kernels/CMakeLists.txt b/src/collectives/kernels/CMakeLists.txt index ee113b84..ccf68244 100644 --- a/src/collectives/kernels/CMakeLists.txt +++ b/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 ) diff --git a/src/collectives/kernels/kernels/collectives.cce b/src/collectives/kernels/kernels/collectives.cce index 1111a77a..0dd16f2d 100644 --- a/src/collectives/kernels/kernels/collectives.cce +++ b/src/collectives/kernels/kernels/collectives.cce @@ -20,6 +20,7 @@ #include #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 __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) diff --git a/src/collectives/kernels/kernels/lcal_allreduce_2npu_write.cce b/src/collectives/kernels/kernels/lcal_allreduce_2npu_write.cce index 92053891..09370792 100644 --- a/src/collectives/kernels/kernels/lcal_allreduce_2npu_write.cce +++ b/src/collectives/kernels/kernels/lcal_allreduce_2npu_write.cce @@ -58,4 +58,4 @@ __attribute__((always_inline)) inline __aicore__ void TileXRAllReduce2npuWrite(A ProcessData(dataSizeRemain, inputUB[0], buff[rank], dataOffsetNum, buffOffsetNum, output, buffOffsetNum, op); DumpLcclLogInfo(dumpAddr, LogId::PROCESS, static_cast(op)); DumpLcclLogInfo(dumpAddr, LogId::OVERALL, static_cast(op)); -} \ No newline at end of file +} diff --git a/src/collectives/kernels/kernels/lcal_allreduce_big_data.cce b/src/collectives/kernels/kernels/lcal_allreduce_big_data.cce index 00818340..796516be 100644 --- a/src/collectives/kernels/kernels/lcal_allreduce_big_data.cce +++ b/src/collectives/kernels/kernels/lcal_allreduce_big_data.cce @@ -9,12 +9,17 @@ */ #include "collectives.cce" + +constexpr uint32_t TILEXR_ALLREDUCE_CMO_WAIT_MISS_THRESHOLD = 8; +constexpr uint32_t TILEXR_ALLREDUCE_CMO_CONSUMER_WORKER = 0; + template __attribute__((always_inline)) inline __aicore__ void TileXRAllReduceBigDataOrigin (__gm__ T* buff[16], __gm__ T *input, __gm__ T *output, int64_t blockNumPerGroup, uint32_t rank, uint32_t rankSize, uint64_t len, int64_t magic, __ubuf__ int64_t* ctrlFlagsUB, __ubuf__ int64_t* ctrlFlagsUB1, __ubuf__ int64_t* ctrlFlagsUB2, __ubuf__ T* inputUB[2], int64_t dataOffsetNum, int64_t flagOffset1st, - int64_t flagOffset2nd, int64_t x, int64_t corePerRank, int64_t coreSegmentedIdx, int op, GM_ADDR perfTrace) + int64_t flagOffset2nd, int64_t x, int64_t corePerRank, int64_t coreSegmentedIdx, int op, GM_ADDR perfTrace, + GM_ADDR cmoTaskPtr) { const uint32_t perfCore = static_cast(GetBlockIdx()); const uint32_t perfRank = static_cast(rank); @@ -67,6 +72,7 @@ __attribute__((always_inline)) inline __aicore__ void TileXRAllReduceBigDataOrig allDataSizeNeed2Add = (thisNPUProcessDataNum - coreSegmentedIdx * thisNPUCoreGroupAvgDMADataNum) * sizeof(T); } AscendC::PipeBarrier(); + uint32_t cmoWaitMissCount = 0; while (true) { if (*ctrlFlagsUB >= CeilDiv(allDataSizeNeed2Add, DMA_SIZE_PER_FLAG)) { break; @@ -88,15 +94,44 @@ __attribute__((always_inline)) inline __aicore__ void TileXRAllReduceBigDataOrig if (*ctrlFlagsUB1 == 0 || *ctrlFlagsUB2 == 0 || ((*ctrlFlagsUB1 >> 10) != (magic >> 10)) || ((*ctrlFlagsUB2 >> 10) != (magic >> 10))) { + cmoWaitMissCount += 1; + if (cmoWaitMissCount >= TILEXR_ALLREDUCE_CMO_WAIT_MISS_THRESHOLD) { + cmoWaitMissCount = 0; + const int64_t cmoPrefetchOffsetNum = rank * singleNPUProcessDataNum + + coreSegmentedIdx * thisNPUCoreGroupAvgDMADataNum + + (*ctrlFlagsUB) * DMA_SIZE_PER_FLAG / sizeof(T); + const int64_t cmoRemainBytes = allDataSizeNeed2Add - (*ctrlFlagsUB) * DMA_SIZE_PER_FLAG; + if (cmoRemainBytes > 0) { + TileXR::TileXRCmoRunWindow(cmoTaskPtr, TILEXR_ALLREDUCE_CMO_CONSUMER_WORKER, + reinterpret_cast((__gm__ T*)((__gm__ int64_t*)buff[x] + dataOffsetNum) + + cmoPrefetchOffsetNum), + static_cast(cmoRemainBytes)); + } + } continue; } int64_t preparedDataGroupCount = ((*ctrlFlagsUB1 & 0x3FF) <= (*ctrlFlagsUB2 & 0x3FF)) ? (*ctrlFlagsUB1 & 0x3FF) : (*ctrlFlagsUB2 & 0x3FF); if (*ctrlFlagsUB >= preparedDataGroupCount) { + cmoWaitMissCount += 1; + if (cmoWaitMissCount >= TILEXR_ALLREDUCE_CMO_WAIT_MISS_THRESHOLD) { + cmoWaitMissCount = 0; + const int64_t cmoPrefetchOffsetNum = rank * singleNPUProcessDataNum + + coreSegmentedIdx * thisNPUCoreGroupAvgDMADataNum + + (*ctrlFlagsUB) * DMA_SIZE_PER_FLAG / sizeof(T); + const int64_t cmoRemainBytes = allDataSizeNeed2Add - (*ctrlFlagsUB) * DMA_SIZE_PER_FLAG; + if (cmoRemainBytes > 0) { + TileXR::TileXRCmoRunWindow(cmoTaskPtr, TILEXR_ALLREDUCE_CMO_CONSUMER_WORKER, + reinterpret_cast((__gm__ T*)((__gm__ int64_t*)buff[x] + dataOffsetNum) + + cmoPrefetchOffsetNum), + static_cast(cmoRemainBytes)); + } + } continue; } + cmoWaitMissCount = 0; buffOffsetNum = rank * singleNPUProcessDataNum + coreSegmentedIdx * thisNPUCoreGroupAvgDMADataNum; dataSizeRemain = (preparedDataGroupCount - *ctrlFlagsUB) * DMA_SIZE_PER_FLAG; if (preparedDataGroupCount * DMA_SIZE_PER_FLAG > allDataSizeNeed2Add) { @@ -104,7 +139,8 @@ __attribute__((always_inline)) inline __aicore__ void TileXRAllReduceBigDataOrig } auto peerToOutput = TileXR::TileXRPerfStageBegin( perfTrace, TileXR::PerfStageId::PEER_IPC_TO_OUTPUT, TileXR::PerfBarrierPolicy::BARRIERED); - ProcessDataNew(dataSizeRemain, inputUB, buff[x], dataOffsetNum, buffOffsetNum + (*ctrlFlagsUB) * DMA_SIZE_PER_FLAG / sizeof(T), + ProcessDataNew(dataSizeRemain, inputUB, buff[x], dataOffsetNum, + buffOffsetNum + (*ctrlFlagsUB) * DMA_SIZE_PER_FLAG / sizeof(T), processOutput, buffOffsetNum + (*ctrlFlagsUB) * DMA_SIZE_PER_FLAG / sizeof(T), op); TileXR::TileXRPerfStageEnd( perfTrace, perfRank, perfCore, TileXR::PerfStageId::PEER_IPC_TO_OUTPUT, @@ -121,7 +157,6 @@ label0: return; } AscendC::PipeBarrier(); - for (int64_t i = 0; i < blockNumPerGroup; i++) { if (i / corePerRank == x) { continue; @@ -146,7 +181,7 @@ label0: } template -__attribute__((always_inline)) inline __aicore__ void TileXRAllReduceBigData(ALLREDUCE_ARGS_FUN_16P(T)) +__attribute__((always_inline)) inline __aicore__ void TileXRAllReduceBigData(ALLREDUCE_ARGS_FUN_16P_CMO(T)) { DumpLcclLogInfo(dumpAddr, LogId::OVERALL, static_cast(op)); DumpLcclLogInfo(dumpAddr, LogId::INIT, static_cast(op)); @@ -196,7 +231,8 @@ __attribute__((always_inline)) inline __aicore__ void TileXRAllReduceBigData(ALL TileXR::PerfStageId::POST_SYNC, postSync, TileXR::PerfBarrierPolicy::END_BARRIER_ONLY); TileXRAllReduceBigDataOrigin( buff, input + processedNum, output + processedNum, blockNumPerGroup, rank, rankSize, remainNum, (magic + i) * 1024, ctrlFlagsUB, ctrlFlagsUB1, - ctrlFlagsUB2, inputUB, dataOffsetNum, flagOffset1st, flagOffset2nd, x, corePerRank, coreSegmentedIdx, op, perfTrace); + ctrlFlagsUB2, inputUB, dataOffsetNum, flagOffset1st, flagOffset2nd, x, corePerRank, coreSegmentedIdx, op, perfTrace, + cmoTaskPtr); auto chunkBarrier = TileXR::TileXRPerfStageBegin( perfTrace, TileXR::PerfStageId::CHUNK_BARRIER, TileXR::PerfBarrierPolicy::NO_BARRIER); AscendC::PipeBarrier(); diff --git a/src/collectives/kernels/lccl_op.h b/src/collectives/kernels/lccl_op.h index 051c4884..e197f8cf 100644 --- a/src/collectives/kernels/lccl_op.h +++ b/src/collectives/kernels/lccl_op.h @@ -70,7 +70,6 @@ struct TileXRCoarsePerfToken { TileXR::TileXRPerfStageEnd( \ perfTrace, tokenName.perfRank, tokenName.perfCore, TileXR::PerfStageId::KERNEL_TOTAL, \ tokenName.kernelTotal, TileXR::PerfBarrierPolicy::NO_BARRIER) - #define CLASS_OP_QUANT_LAUNCH(name, outputType, inputType) \ do { \ name opKernel(localRank, localRankSize, extraFlag); \ @@ -160,10 +159,16 @@ extern "C" __global__ __aicore__ void TileXRAllReduce_##type##suffix(KERNELS_ARG __gm__ type * shareAddrs[TILEXR_MAX_RANK_SIZE]; \ GET_IPC_MEM_ARGS(type); \ if ((extraFlag & ExtraFlag::TOPO_PCIE) != 0) { \ - if (len * sizeof(type) < SIZE_OF_8M) { \ - TileXRAllReduce2npuWrite(ALLREDUCE_ARGS_CALL_16P(type)); \ + if (rankSize == quickOneshotRankSize) { \ + if (len * sizeof(type) < SIZE_OF_8M) { \ + TileXRAllReduce2npuWrite(ALLREDUCE_ARGS_CALL_16P(type)); \ + } else { \ + TileXRAllReduce2npuBigDataWrite(ALLREDUCE_ARGS_CALL_16P(type)); \ + } \ + } else if (len * sizeof(type) < cceSmallDataSize || lcalBlockNum == rankSize) { \ + TileXRAllReduceTwoShot(ALLREDUCE_ARGS_CALL_16P(type)); \ } else { \ - TileXRAllReduce2npuBigDataWrite(ALLREDUCE_ARGS_CALL_16P(type)); \ + TileXRAllReduceBigData(ALLREDUCE_ARGS_CALL_16P_CMO(type)); \ } \ } else if ((extraFlag & ExtraFlag::QUANT_FP16) != 0 && std::is_same_v) { \ if (len * sizeof(type) <= oneshotDataSize) { \ @@ -218,7 +223,7 @@ extern "C" __global__ __aicore__ void TileXRAllReduce_##type##suffix(KERNELS_ARG TileXRAllReduceTwoShot(ALLREDUCE_ARGS_CALL_16P(type)); \ } \ } else { \ - TileXRAllReduceBigData(ALLREDUCE_ARGS_CALL_16P(type)); \ + TileXRAllReduceBigData(ALLREDUCE_ARGS_CALL_16P_CMO(type)); \ } \ } \ TILEXR_COARSE_PERF_END(coarsePerf); \ diff --git a/src/comm/CMakeLists.txt b/src/comm/CMakeLists.txt index a7c9ceda..15ea88f9 100644 --- a/src/comm/CMakeLists.txt +++ b/src/comm/CMakeLists.txt @@ -183,6 +183,7 @@ install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma_types.h ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma.h ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma_compat.h + ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_cmo.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) install(CODE "string(CONCAT _tilexr_collective_header \"tilexr_\" \"collectives.h\") string(CONCAT _tilexr_collective_perf_header \"tilexr_\" \"collectives\" \"_perf.h\") diff --git a/src/comm/comm_wrap.cpp b/src/comm/comm_wrap.cpp index 9d594323..6c9e4861 100644 --- a/src/comm/comm_wrap.cpp +++ b/src/comm/comm_wrap.cpp @@ -225,6 +225,29 @@ int TileXRGetSDMAWorkspaceDev(TileXRCommPtr comm, GM_ADDR *workspace) return TILEXR_SUCCESS; } + +int TileXRSubmitCmoTask(TileXRCommPtr comm, GM_ADDR targetAddr, uint64_t totalBytes, + uint32_t opType, uint32_t priority, uint32_t chunkBytes, + uint64_t expireSeq) +{ + if (comm == nullptr) { + TILEXR_LOG(ERROR) << "TileXRSubmitCmoTask invalid comm"; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + auto* c = static_cast(comm); + return c->SubmitCmoTask(targetAddr, totalBytes, opType, priority, chunkBytes, expireSeq); +} + +int TileXRClearCmoTask(TileXRCommPtr comm) +{ + if (comm == nullptr) { + TILEXR_LOG(ERROR) << "TileXRClearCmoTask invalid comm"; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + auto* c = static_cast(comm); + return c->ClearCmoTask(); +} + void TileXRPrintDFX2Log(TileXRCommPtr comm) { if (comm == nullptr) { diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 9b1d3d4e..47c8e58f 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -54,6 +54,7 @@ static std::mutex g_mtx; static std::mutex g_sdmaMtx; static bool g_sdmaUnavailable = false; +<<<<<<< HEAD namespace { bool IsEnvEnabled(const char* name, bool defaultValue) @@ -75,6 +76,18 @@ bool IsEnvEnabled(const char* name, bool defaultValue) } // namespace +======= +static bool EnvBoolTrue(const char *name) +{ + const char *value = std::getenv(name); + if (value == nullptr) { + return false; + } + const string text(value); + return text == "1" || text == "true" || text == "TRUE" || text == "yes" || text == "YES"; +} + +>>>>>>> c13c44a (add cmo code) // 如果是互联的链路，返回false； 对910B2C那些不互联的链路，返回true bool SkipUnusedChannel910B2C(int curRank, int peerRank, ChipName chipName) @@ -351,6 +364,101 @@ int TileXRComm::UpdateCommArgsDev() return TILEXR_SUCCESS; } +int TileXRComm::SubmitCmoTask(GM_ADDR targetAddr, uint64_t totalBytes, uint32_t opType, uint32_t priority, + uint32_t chunkBytes, uint64_t expireSeq) +{ + if (!inited_) { + TILEXR_LOG(ERROR) << "TileXRSubmitCmoTask requires initialized communicator"; + return TILEXR_ERROR_NOT_INITIALIZED; + } + if (targetAddr == nullptr || totalBytes == 0 || !TileXRCmoOpTypeValid(opType)) { + TILEXR_LOG(ERROR) << "TileXRSubmitCmoTask invalid args, target " << static_cast(targetAddr) + << ", bytes " << totalBytes << ", op " << opType; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + if ((commArgs_.extraFlag & ExtraFlag::TOPO_PCIE) != 0 && + !EnvBoolTrue("TILEXR_CMO_ALLOW_PCIE_FALLBACK")) { + TILEXR_LOG(WARN) << "TileXRSubmitCmoTask skipped on PCIe fallback without TILEXR_CMO_ALLOW_PCIE_FALLBACK=1"; + return TILEXR_ERROR_NOT_SUPPORT; + } + + bool allocated = false; + if (cmoTaskDev_ == nullptr) { + int ret = aclrtMalloc(reinterpret_cast(&cmoTaskDev_), sizeof(TileXRCmoTaskDesc), + ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + TILEXR_LOG(ERROR) << "aclrtMalloc cmo task err " << ret; + cmoTaskDev_ = nullptr; + return TILEXR_ERROR_INTERNAL; + } + allocated = true; + } + + TileXRCmoTaskDesc task {}; + task.targetAddr = targetAddr; + task.totalBytes = totalBytes; + task.offset = 0; + task.opType = static_cast(opType); + task.priority = priority; + task.chunkBytes = chunkBytes; + task.state = CmoTaskState::PENDING; + task.expireSeq = expireSeq; + task.executedBytes = 0; + + int ret = aclrtMemcpy(cmoTaskDev_, sizeof(task), &task, sizeof(task), ACL_MEMCPY_HOST_TO_DEVICE); + if (ret != ACL_SUCCESS) { + TILEXR_LOG(ERROR) << "aclrtMemcpy cmo task err " << ret; + if (allocated) { + FreePeerMem(cmoTaskDev_); + } + return TILEXR_ERROR_INTERNAL; + } + + const uint32_t oldExtraFlag = commArgs_.extraFlag; + const GM_ADDR oldCmoTaskPtr = commArgs_.cmoTaskPtr; + const TileXRCmoTaskDesc oldTask = cmoTaskHost_; + + cmoTaskHost_ = task; + commArgs_.cmoTaskPtr = cmoTaskDev_; + commArgs_.extraFlag |= ExtraFlag::CMO; + + ret = UpdateCommArgsDev(); + if (ret != TILEXR_SUCCESS) { + commArgs_.extraFlag = oldExtraFlag; + commArgs_.cmoTaskPtr = oldCmoTaskPtr; + cmoTaskHost_ = oldTask; + if (allocated) { + FreePeerMem(cmoTaskDev_); + } + return ret; + } + return TILEXR_SUCCESS; +} + +int TileXRComm::ClearCmoTask() +{ + if (!inited_) { + TILEXR_LOG(ERROR) << "TileXRClearCmoTask requires initialized communicator"; + return TILEXR_ERROR_NOT_INITIALIZED; + } + + const uint32_t oldExtraFlag = commArgs_.extraFlag; + const GM_ADDR oldCmoTaskPtr = commArgs_.cmoTaskPtr; + + commArgs_.extraFlag &= ~ExtraFlag::CMO; + commArgs_.cmoTaskPtr = nullptr; + int ret = UpdateCommArgsDev(); + if (ret != TILEXR_SUCCESS) { + commArgs_.extraFlag = oldExtraFlag; + commArgs_.cmoTaskPtr = oldCmoTaskPtr; + return ret; + } + + FreePeerMem(cmoTaskDev_); + cmoTaskHost_ = {}; + return TILEXR_SUCCESS; +} + int TileXRComm::RegisterUDMAMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMemHandle *handle) { if (!inited_) { @@ -1097,7 +1205,11 @@ TileXRComm::~TileXRComm() } FreePeerMem(commArgs_.dumpAddr); FreePeerMem(peerMem_[rank_]); +<<<<<<< HEAD FreePeerMem(creditIpcMem_[rank_]); +======= + FreePeerMem(cmoTaskDev_); +>>>>>>> c13c44a (add cmo code) FreePeerMem(commArgsPtr_); ResetSDMAState(); } diff --git a/src/comm/tilexr_comm.h b/src/comm/tilexr_comm.h index 6d9b152a..b3b9a8f5 100644 --- a/src/comm/tilexr_comm.h +++ b/src/comm/tilexr_comm.h @@ -19,6 +19,7 @@ #include "../include/tilexr_types.h" #include "../include/tilexr_api.h" #include "../include/comm_args.h" +#include "../include/tilexr_cmo.h" namespace TileXR { constexpr int IPC_NAME_SIZE = 65; @@ -51,6 +52,9 @@ class TileXRComm { const TileXRUDMARegistry* GetUDMARegistryHost() const; bool IsSDMAAvailable() const; GM_ADDR GetSDMAWorkspacePtr() const; + int SubmitCmoTask(GM_ADDR targetAddr, uint64_t totalBytes, uint32_t opType, uint32_t priority, + uint32_t chunkBytes, uint64_t expireSeq); + int ClearCmoTask(); SDMAInitStatus GetSDMAInitStatus() const; std::string PrintDFX(); friend class Lccl; @@ -112,6 +116,8 @@ class TileXRComm { bool isEnableMsprofOp_ = false; std::unique_ptr udmaContext_; GM_ADDR sdmaWorkspaceDev_ = nullptr; + GM_ADDR cmoTaskDev_ = nullptr; + TileXRCmoTaskDesc cmoTaskHost_ = {}; SDMAInitStatus sdmaInitStatus_ = SDMAInitStatus::DISABLED_BY_ENV; std::unique_ptr sdmaTransport_; }; diff --git a/src/include/comm_args.h b/src/include/comm_args.h index 8d71accd..3f5bcf8f 100644 --- a/src/include/comm_args.h +++ b/src/include/comm_args.h @@ -100,6 +100,7 @@ struct ExtraFlag { static constexpr uint32_t ATOMIC_ENABLE = 1 << 15; // 表示在910A5算子中启用atomic实现 static constexpr uint32_t IS_GREATER_THAN_40_AIV = 1 << 16; static constexpr uint32_t PERF_CYCLE_A5 = 1 << 17; + static constexpr uint32_t CMO = 1 << 18; }; struct CommArgs { @@ -122,6 +123,7 @@ struct CommArgs { GM_ADDR udmaInfoPtr = nullptr; // device-side TileXR::UDMAInfo*; nullptr 表示 UDMA 不可用 GM_ADDR udmaRegistryPtr = nullptr; // device-side TileXRUDMARegistry* for user-registered UDMA memory GM_ADDR sdmaWorkspacePtr = nullptr; // device-side SDMA workspace; nullptr 表示 SDMA 不可用 + GM_ADDR cmoTaskPtr = nullptr; // device-side TileXRCmoTaskDesc*; nullptr means no pending CMO task }; struct LcclDumpBlockInfo { diff --git a/src/include/tilexr_api.h b/src/include/tilexr_api.h index a7b4c98b..254cec22 100644 --- a/src/include/tilexr_api.h +++ b/src/include/tilexr_api.h @@ -20,6 +20,13 @@ extern "C" { typedef void *TileXRCommPtr; typedef uint32_t TileXRUDMAMemHandle; + +typedef enum { + TILEXR_CMO_PREFETCH = 0, + TILEXR_CMO_FLUSH = 1, + TILEXR_CMO_INVALIDATE = 2, +} TileXRCmoOpType; + #define TILEXRUNIQUE_ID_BYTES 128 typedef struct { char internal[TILEXRUNIQUE_ID_BYTES]; } TileXRUniqueId; @@ -53,6 +60,12 @@ int TileXRSDMAAvailable(TileXRCommPtr comm, bool *available); int TileXRGetSDMAWorkspaceDev(TileXRCommPtr comm, GM_ADDR *workspace); +int TileXRSubmitCmoTask(TileXRCommPtr comm, GM_ADDR targetAddr, uint64_t totalBytes, + uint32_t opType, uint32_t priority, uint32_t chunkBytes, + uint64_t expireSeq); + +int TileXRClearCmoTask(TileXRCommPtr comm); + void TileXRPrintDFX2Log(TileXRCommPtr comm); int TileXRCommInit(int rank, int rankSize, TileXRCommPtr *comms); diff --git a/tests/collectives/tilexr-tests/tilexr_collective_perf.cpp b/tests/collectives/tilexr-tests/tilexr_collective_perf.cpp index 6ec5d51d..23906881 100644 --- a/tests/collectives/tilexr-tests/tilexr_collective_perf.cpp +++ b/tests/collectives/tilexr-tests/tilexr_collective_perf.cpp @@ -50,6 +50,11 @@ enum class CommMode { SOCKET, }; +enum class CmoTargetMode { + SEND, + SCRATCH, +}; + struct DataTypeInfo { TileXR::TileXRDataType type = TileXR::TILEXR_DATA_TYPE_INT32; std::string name = "int32"; @@ -77,6 +82,12 @@ struct Options { std::string profileDir; bool profileAiPrompt = false; int profileSampleEvery = 1; + bool cmo = false; + uint32_t cmoOpType = TILEXR_CMO_PREFETCH; + int64_t cmoBytes = 0; + int64_t cmoMinBytes = 8 * 1024 * 1024; + uint32_t cmoChunkBytes = 0; + CmoTargetMode cmoTargetMode = CmoTargetMode::SEND; }; struct Measurement { @@ -202,7 +213,10 @@ void PrintUsage(const char *program) << " --check 0|1 [--csv path]\n" << " [--min-algbw GB/s] [--max-latency-us us]\n" << " [--profile 0|1] [--profile-dir path]\n" - << " [--profile-ai-prompt 0|1] [--profile-sample-every N]\n"; + << " [--profile-ai-prompt 0|1] [--profile-sample-every N]\n" + << " [--cmo 0|1] [--cmo-op prefetch|flush|invalidate|N]\n" + << " [--cmo-bytes N] [--cmo-min-bytes N] [--cmo-chunk-bytes N]\n" + << " [--cmo-target send|scratch]\n"; } bool ParseBool(const std::string &value, bool &out) @@ -276,6 +290,42 @@ bool ParseInt(const std::string &text, int &out) return true; } +bool ParseCmoOpType(const std::string &value, uint32_t &opType) +{ + if (value == "prefetch") { + opType = TILEXR_CMO_PREFETCH; + return true; + } + if (value == "flush") { + opType = TILEXR_CMO_FLUSH; + return true; + } + if (value == "invalidate") { + opType = TILEXR_CMO_INVALIDATE; + return true; + } + int64_t numeric = 0; + if (!ParseInt64(value, numeric) || numeric < TILEXR_CMO_PREFETCH || numeric > TILEXR_CMO_INVALIDATE) { + return false; + } + opType = static_cast(numeric); + return true; +} + + +bool ParseCmoTargetMode(const std::string &value, CmoTargetMode &mode) +{ + if (value == "send") { + mode = CmoTargetMode::SEND; + return true; + } + if (value == "scratch") { + mode = CmoTargetMode::SCRATCH; + return true; + } + return false; +} + bool ParseDouble(const std::string &text, double &out) { char *end = nullptr; @@ -605,6 +655,44 @@ bool ParseOptions(int argc, char **argv, Options &options) std::cerr << "ERROR: invalid --profile-sample-every" << std::endl; return false; } + } else if (arg == "--cmo") { + const char *value = requireValue(arg); + if (value == nullptr || !ParseBool(value, options.cmo)) { + std::cerr << "ERROR: --cmo must be 0 or 1" << std::endl; + return false; + } + } else if (arg == "--cmo-op") { + const char *value = requireValue(arg); + if (value == nullptr || !ParseCmoOpType(value, options.cmoOpType)) { + std::cerr << "ERROR: --cmo-op must be prefetch, flush, invalidate, or 0..2" << std::endl; + return false; + } + } else if (arg == "--cmo-bytes") { + const char *value = requireValue(arg); + if (value == nullptr || !ParseInt64(value, options.cmoBytes)) { + std::cerr << "ERROR: invalid --cmo-bytes" << std::endl; + return false; + } + } else if (arg == "--cmo-min-bytes") { + const char *value = requireValue(arg); + if (value == nullptr || !ParseInt64(value, options.cmoMinBytes)) { + std::cerr << "ERROR: invalid --cmo-min-bytes" << std::endl; + return false; + } + } else if (arg == "--cmo-chunk-bytes") { + const char *value = requireValue(arg); + int parsed = 0; + if (value == nullptr || !ParseInt(value, parsed) || parsed < 0) { + std::cerr << "ERROR: invalid --cmo-chunk-bytes" << std::endl; + return false; + } + options.cmoChunkBytes = static_cast(parsed); + } else if (arg == "--cmo-target") { + const char *value = requireValue(arg); + if (value == nullptr || !ParseCmoTargetMode(value, options.cmoTargetMode)) { + std::cerr << "ERROR: --cmo-target must be send or scratch" << std::endl; + return false; + } } else if (arg == "--help" || arg == "-h") { PrintUsage(argv[0]); std::exit(0); @@ -625,6 +713,14 @@ bool ParseOptions(int argc, char **argv, Options &options) std::cerr << "ERROR: --profile-sample-every must be positive" << std::endl; return false; } + if (options.cmoBytes < 0) { + std::cerr << "ERROR: --cmo-bytes must be non-negative" << std::endl; + return false; + } + if (options.cmoMinBytes < 0) { + std::cerr << "ERROR: --cmo-min-bytes must be non-negative" << std::endl; + return false; + } if (options.check && options.dtype.name != "int32" && IsReductionOp(options.op)) { std::cerr << "ERROR: --check=1 for allreduce/reducescatter requires --datatype int32" << std::endl; return false; @@ -653,9 +749,69 @@ bool CheckTileXR(int rank, const std::string &step, int ret) return false; } +bool EnvBoolTrue(const char *name) +{ + const char *value = std::getenv(name); + if (value == nullptr) { + return false; + } + const std::string text(value); + return text == "1" || text == "true" || text == "yes"; +} + +bool AllowCmoOnCurrentEnv() +{ + if (EnvBoolTrue("TILEXR_CMO_ALLOW_PCIE_FALLBACK")) { + return true; + } + const char *soc = std::getenv("TILEXR_SOC_NAME"); + if (soc == nullptr) { + return true; + } + const std::string socName(soc); + return socName.find("910b") == std::string::npos && socName.find("910B") == std::string::npos; +} + +bool ShouldSubmitCmoForMessage(const Options &options, int64_t messageBytes) +{ + if (!options.cmo || options.op != CollectiveOp::ALLREDUCE || !AllowCmoOnCurrentEnv()) { + return false; + } + return messageBytes >= options.cmoMinBytes; +} + +uint64_t ResolveCmoBytesForMessage(const Options &options, int64_t messageBytes) +{ + if (!ShouldSubmitCmoForMessage(options, messageBytes)) { + return 0; + } + uint64_t bytes = static_cast(options.cmoBytes > 0 ? options.cmoBytes : messageBytes); + if (options.cmoTargetMode == CmoTargetMode::SEND && bytes > static_cast(messageBytes)) { + bytes = static_cast(messageBytes); + } + return bytes; +} + +bool SubmitCmoTaskIfEnabled(const Options &options, TileXRCommPtr comm, void *cmoTarget, uint64_t cmoBytes) +{ + if (!options.cmo || cmoBytes == 0) { + return true; + } + if (cmoTarget == nullptr) { + std::cerr << "[rank " << options.rank << "] ERROR: CMO is enabled but target is empty" << std::endl; + return false; + } + return CheckTileXR(options.rank, "TileXRSubmitCmoTask", + TileXRSubmitCmoTask(comm, static_cast(cmoTarget), cmoBytes, options.cmoOpType, 0, + options.cmoChunkBytes, 0)); +} + bool CallCollective(const Options &options, void *sendBuf, void *recvBuf, int64_t count, TileXRCommPtr comm, - aclrtStream stream) + aclrtStream stream, void *cmoTarget = nullptr, uint64_t cmoBytes = 0) { + if (!SubmitCmoTaskIfEnabled(options, comm, cmoTarget, cmoBytes)) { + return false; + } switch (options.op) { case CollectiveOp::ALLGATHER: return CheckTileXR(options.rank, "TileXRAllGather", @@ -962,7 +1118,8 @@ double ComputeBusBandwidthGbps(CollectiveOp op, int rankSize, double algBwGbps) } bool MeasureOnce(const Options &options, void *devSend, void *devRecv, int64_t count, TileXRCommPtr comm, - aclrtStream stream, uint64_t profileLaunchIndex, int &totalErrors, double &us) + aclrtStream stream, uint64_t profileLaunchIndex, int &totalErrors, double &us, + void *cmoTarget, uint64_t cmoBytes) { aclrtEvent start = nullptr; aclrtEvent stop = nullptr; @@ -985,7 +1142,7 @@ bool MeasureOnce(const Options &options, void *devSend, void *devRecv, int64_t c return false; } std::string incompleteReason; - if (!CallCollective(options, devSend, devRecv, count, comm, stream)) { + if (!CallCollective(options, devSend, devRecv, count, comm, stream, cmoTarget, cmoBytes)) { incompleteReason = "CallCollective failed before measured launch completed"; } else if (!CheckAcl(options.rank, "aclrtRecordEvent stop", aclrtRecordEvent(stop, stream))) { incompleteReason = "aclrtRecordEvent stop failed"; @@ -1013,10 +1170,11 @@ bool MeasureOnce(const Options &options, void *devSend, void *devRecv, int64_t c } bool Measure(const Options &options, void *devSend, void *devRecv, int64_t count, TileXRCommPtr comm, - aclrtStream stream, uint64_t &profileLaunchIndex, int &totalErrors, Measurement &measurement) + aclrtStream stream, uint64_t &profileLaunchIndex, int &totalErrors, Measurement &measurement, + void *cmoTarget, uint64_t cmoBytes) { for (int i = 0; i < options.warmupIters; ++i) { - if (!CallCollective(options, devSend, devRecv, count, comm, stream)) { + if (!CallCollective(options, devSend, devRecv, count, comm, stream, cmoTarget, cmoBytes)) { return false; } } @@ -1030,7 +1188,7 @@ bool Measure(const Options &options, void *devSend, void *devRecv, int64_t count double us = 0.0; const uint64_t currentProfileLaunchIndex = profileLaunchIndex++; if (!MeasureOnce(options, devSend, devRecv, count, comm, stream, - currentProfileLaunchIndex, totalErrors, us)) { + currentProfileLaunchIndex, totalErrors, us, cmoTarget, cmoBytes)) { return false; } samples.push_back(us); @@ -1401,6 +1559,9 @@ int main(int argc, char **argv) void *devSend = nullptr; void *devRecv = nullptr; + void *devCmo = nullptr; + void *activeCmoTarget = nullptr; + uint64_t activeCmoBytes = 0; bool ok = CheckAcl(options.rank, "aclrtMalloc send", aclrtMalloc(&devSend, static_cast(sendBytes), ACL_MEM_MALLOC_HUGE_FIRST)) && CheckAcl(options.rank, "aclrtMemcpy H2D send", @@ -1418,10 +1579,27 @@ int main(int argc, char **argv) } } + if (ok && ShouldSubmitCmoForMessage(options, actualSendBytesPerRank)) { + activeCmoBytes = ResolveCmoBytesForMessage(options, actualSendBytesPerRank); + if (activeCmoBytes == 0 || activeCmoBytes > static_cast(kMaxHostBufferBytes)) { + std::cerr << "ERROR: invalid active CMO bytes " << activeCmoBytes << std::endl; + ok = false; + } else if (options.cmoTargetMode == CmoTargetMode::SCRATCH) { + ok = CheckAcl(options.rank, "aclrtMalloc cmo", + aclrtMalloc(&devCmo, static_cast(activeCmoBytes), ACL_MEM_MALLOC_HUGE_FIRST)) && + CheckAcl(options.rank, "aclrtMemset cmo", + aclrtMemset(devCmo, static_cast(activeCmoBytes), 0, + static_cast(activeCmoBytes))); + activeCmoTarget = devCmo; + } else { + activeCmoTarget = devSend; + } + } + Measurement measurement; if (ok) { ok = Measure(options, devSend, devRecv, count, comm, stream, - profileLaunchIndex, totalErrors, measurement); + profileLaunchIndex, totalErrors, measurement, activeCmoTarget, activeCmoBytes); } int errors = 0; @@ -1437,7 +1615,7 @@ int main(int argc, char **argv) static_cast(sendBytes), ACL_MEMCPY_HOST_TO_DEVICE)); } ok = ok && - CallCollective(options, devSend, devRecv, count, comm, stream) && + CallCollective(options, devSend, devRecv, count, comm, stream, activeCmoTarget, activeCmoBytes) && CheckAcl(options.rank, "aclrtSynchronizeStream check", aclrtSynchronizeStream(stream)) && CheckAcl(options.rank, "aclrtMemcpy D2H recv", aclrtMemcpy(hostRecv.data(), static_cast(recvBytes), devRecv, @@ -1450,6 +1628,10 @@ int main(int argc, char **argv) } } + + if (devCmo != nullptr) { + aclrtFree(devCmo); + } if (devSend != nullptr) { aclrtFree(devSend); } diff --git a/tests/collectives/unit/test_collective_profile_report.py b/tests/collectives/unit/test_collective_profile_report.py index 3f656d91..9289264c 100644 --- a/tests/collectives/unit/test_collective_profile_report.py +++ b/tests/collectives/unit/test_collective_profile_report.py @@ -191,9 +191,11 @@ def test_writes_perfetto_trace_for_ui_perfetto_dev(self): trace_json = json.loads((root / "trace.json").read_text(encoding="utf-8")) self.assertIn("traceEvents", trace_json) + self.assertNotIn("displayTimeUnit", trace_json) perfetto = json.loads((root / "perfetto_trace.json").read_text(encoding="utf-8")) self.assertIn("traceEvents", perfetto) + self.assertNotIn("displayTimeUnit", perfetto) self.assertIn({"name": "process_name", "ph": "M", "pid": 0, "args": {"name": "rank0"}}, perfetto["traceEvents"]) self.assertIn({"name": "thread_name", "ph": "M", "pid": 0, "tid": 0, "args": {"name": "rank0/core0"}}, perfetto["traceEvents"]) diff --git a/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp b/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp index 55fd34c9..fe5045a8 100644 --- a/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp +++ b/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp @@ -152,6 +152,7 @@ void TestCollectivesOwnsCceBuild() CheckContains(kernelsCmakePath, kernelsCmake, "CONFIGURE_DEPENDS"); CheckContains(kernelsCmakePath, kernelsCmake, "OBJECT_DEPENDS"); CheckContains(kernelsCmakePath, kernelsCmake, "tilexr_collectives_op.o"); + CheckContains(kernelsCmakePath, kernelsCmake, "tilexr_lccl_op.cpp.o"); CheckContains(kernelsCmakePath, kernelsCmake, "tilexr_collectives_op"); CheckContains(kernelsCmakePath, kernelsCmake, "TILEXR_COLLECTIVES_ENABLE_PROFILING"); CheckContains(kernelsCmakePath, kernelsCmake, "TILEXR_COLLECTIVES_1OP_BIN_SIZE 10485760"); @@ -182,6 +183,8 @@ void TestCollectivesKernelSourcesAreScoped() CheckContains(perfTraceKernelPath, perfTraceKernel, "TILEXR_PERF_TRACE_STATS_OFFSET"); CheckContains(perfTraceKernelPath, perfTraceKernel, "CpGM2UB"); CheckContains(perfTraceKernelPath, perfTraceKernel, "CpUB2GM"); + CheckContains(perfTraceKernelPath, perfTraceKernel, "TILEXR_PERF_TRACE_STATS_OFFSET"); + CheckContains(perfTraceKernelPath, perfTraceKernel, "TILEXR_PERF_STAGE_COUNT"); CheckContains(perfTraceKernelPath, perfTraceKernel, "GetBlockNum()"); CheckDoesNotContain(perfTraceKernelPath, perfTraceKernel, "header->statsOffset"); CheckDoesNotContain(perfTraceKernelPath, perfTraceKernel, "header->maxCoreCount"); diff --git a/tools/collectives/tilexr_collective_profile_report.py b/tools/collectives/tilexr_collective_profile_report.py index bfbb5fd2..da28d5d7 100755 --- a/tools/collectives/tilexr_collective_profile_report.py +++ b/tools/collectives/tilexr_collective_profile_report.py @@ -804,7 +804,6 @@ def render_perfetto_trace(index): }) return { - "displayTimeUnit": "us", "traceEvents": events, }