diff --git a/CMakeLists.txt b/CMakeLists.txt index 6f80a96c..3914b0f8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -14,8 +14,27 @@ project(tilexr LANGUAGES CXX) set(CMAKE_CXX_STANDARD 14) option(TILEXR_BUILD_COLLECTIVES "Build optional TileXR collectives library" OFF) option(TILEXR_BUILD_EP "Build TileXR EP communication library" OFF) +option(TILEXR_EP_BUILD_URMA_COMBINE "Build the Ascend950 URMA EP combine operator" OFF) option(TILEXR_BUILD_TESTS "Build TileXR tests" OFF) option(TILEXR_BUILD_CHECKER "Build TileXR no-NPU checker" OFF) +set(TILEXR_UDMA_QP_COUNT "1" CACHE STRING "Base UDMA queue-pair count per EID") +set(TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT "22" CACHE STRING + "Send AIV count for the Ascend950 URMA EP combine operator") + +if(NOT TILEXR_UDMA_QP_COUNT MATCHES "^[1-9][0-9]*$") + message(FATAL_ERROR "TILEXR_UDMA_QP_COUNT must be a positive integer") +endif() +if(NOT TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT MATCHES "^[1-9][0-9]*$" OR + TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT GREATER_EQUAL 64) + message(FATAL_ERROR + "TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT must be an integer in [1, 63]") +endif() + +set(TILEXR_UDMA_EFFECTIVE_QP_COUNT "${TILEXR_UDMA_QP_COUNT}") +if(TILEXR_BUILD_EP AND TILEXR_EP_BUILD_URMA_COMBINE AND + TILEXR_UDMA_EFFECTIVE_QP_COUNT LESS TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT) + set(TILEXR_UDMA_EFFECTIVE_QP_COUNT "${TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT}") +endif() include(CTest) if(BUILD_TESTING OR TILEXR_BUILD_TESTS) enable_testing() diff --git a/docs/EP_URMA_COMBINE_S22.md b/docs/EP_URMA_COMBINE_S22.md new file mode 100644 index 00000000..74db1a21 --- /dev/null +++ b/docs/EP_URMA_COMBINE_S22.md @@ -0,0 +1,79 @@ +# EP URMA Combine for Ascend 950 + +This change adds a TileXR-native Mixture-of-Experts combine operator backed by +URMA/UDMA. The production implementation is fixed to the validated S22 layout: + +- 42 Pack/Receive AIVs and 22 Send AIVs; +- 22 independent UDMA queue pairs; +- one WQE per doorbell; +- QDC-v3 quant/dequant; +- round-robin receive scheduling with sticky ready state; +- parallel release publication and deferred same-parity credit checks; +- first-launch start gate after a stream synchronization; +- BiSheng `-O2`, with profiling disabled by default. + +The fixed configuration avoids carrying experiment-only variant switches into +the production build. The kernel contains no experiment-selection macros or +alternate protocol branches; only the Send AIV count and the explicit +profiling build remain configurable. + +## Build + +Configure and build through CMake. No variant or wrapper script is required. + +```bash +cmake -S . -B build-s22 \ + -DTILEXR_BUILD_EP=ON \ + -DTILEXR_EP_BUILD_URMA_COMBINE=ON \ + -DTILEXR_EP_SOC_TYPE=ascend950 \ + -DBUILD_TESTING=OFF \ + -DCMAKE_INSTALL_PREFIX="$PWD/install-s22" +cmake --build build-s22 -j --target tilexr-ep +cmake --install build-s22 +``` + +`ASCEND_HOME_PATH` and `ASCEND_DRIVER_PATH` must point to the active CANN +toolkit and driver before configuration. A profiling build is available only +when explicitly requested with `-DTILEXR_EP_ENABLE_PROFILING=ON`; it is not the +production default. + +The URMA Combine target is opt-in. A Dispatch-only EP build leaves +`TILEXR_EP_BUILD_URMA_COMBINE=OFF` and uses the generic UDMA default of one QP +per EID. When URMA Combine is enabled, CMake raises the effective QP count to +the configured Send AIV count; `comm/udma` does not contain an Ascend950- or +operator-specific QP constant. The validated default is 22 Send AIVs and can +be changed with `-DTILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT=`. + +## API + +The public header `src/include/tilexr_ep.h` provides: + +- `TileXRMoeEpCombineUrmaGetWorkspaceSize` for the registered workspace size; +- `TileXRMoeEpCombineUrma` for the production launch; +- `TileXRMoeEpCombineUrmaProfile` for an explicitly profiled launch. + +The workspace must be aligned and registered with TileXR UDMA before launch. +The operator supports an eight-rank Ascend 950 deployment and uses the existing +TileXR communicator for rank information and registered-memory exchange. + +## Performance Evidence + +The production comparison metric is profiling-free `strictKernelCycles`, +aggregated as max core per rank, max rank per launch, then median across 100 +launches. For BS128, H=5120, top-k=6, rank-size=8 and enqueue-window=1, S22 +measured 88,343.5 cycles (88.3435 us using 1000 cycles/us). The paired baseline +was 88,602.5 cycles, a 0.29% reduction. BS32 improved from 39,753 to 38,541 +cycles (3.05%). + +The single retained detailed report is +[BS128 S22 profile](performance/tilexr_ep_urma_combine_s22_bs128.html). Its +charts, stage maxima, heatmaps and kernel KPIs exclude the first-launch start +gate and rebase the axis to steady-state device work. The raw embedded capture +retains the gate samples only for provenance. + +This boundary explains the large host/device timing difference. The host round +starts before launch and ends after stream synchronization, so it includes the +start-gate wait, API launch/synchronization overhead and cross-rank scheduling +skew. Device `kernel_total` begins after the gate releases. The values therefore +describe different intervals and must not be subtracted to infer kernel work; +production comparisons use steady-state `strictKernelCycles`. diff --git a/docs/performance/tilexr_ep_urma_combine_s22_bs128.html b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html new file mode 100644 index 00000000..ad99b7ba --- /dev/null +++ b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html @@ -0,0 +1,71 @@ + + + + + +TileXR EP URMA Combine - S22 P42/S22/QP22 O2 - BS128 steady-state profile + + + +
steady-state capture

TileXR EP URMA Combine · S22 P42/S22/QP22 O2 · BS128

+
8 ranks · 64 AIV · URMA · 10 paired profile samples · representative launch0 · median cluster kernel 83.87 us (coarse attribution profile) · coarse/fine attribution medians 83.87/92.32 us
+
42 Pack/Receive + 22 Send · 22 QPs · 1-WQE doorbell · Q/D v3 · bs=128, topK=6, h=5120, selfSendCnt=768, routeStride=5632B, enqueueWindow=1
+
+
+
+

Host and device timing boundary

The synchronization preamble is intentionally omitted from every chart, stage maximum, heatmap, and kernel KPI in this report. The device kernel_total interval begins after that preamble has released, while the host round timer starts before launch and ends after stream synchronization. The host value therefore includes the first-launch synchronization wait, API launch and synchronization overhead, and cross-rank scheduling skew; the device value does not.

The omitted wait has a 227.67 us executed-only median in this capture and is the dominant reason that host latency is much larger than the 83.87 us coarse device critical path. The two values have different boundaries and must not be subtracted to infer unaccounted kernel work. Production comparisons use profiling-free strictKernelCycles; the final BS128 median is 88,343.5 cycles.

+

Coarse Stage Maxima Across AIVs

Each value is the maximum coarse sum_us for that stage across the selected rank's AIVs. Maxima can come from different cores, overlap in time, and must not be added as a critical path.

+

Fine Explicit DCCI Breakdown

Fine diagnostic capture; not the coarse timing launch.

+

64-Core Coarse Timeline

Outlined bars are the exact per-stage maxima listed above. Receive bar width and position are coarse; its internal colors are fine-derived composition, not chronology. The axis begins at the first steady-state device work; the synchronization preamble is not drawn.

+

Core x Fine Exposed-Time Heatmap

Fine diagnostic capture; absolute durations include instrumentation overhead.

+

Fine Route Transport AIV Diagnostics

Fine diagnostic capture; route counters and exposed stage time are not from the coarse launch.

+

Rank Critical Paths

+
The total-only capture provides the timing headline when supplied; it disables stage attribution inside the measured kernel. Coarse bars are per-core activity envelopes for attribution. The synchronization preamble remains in the embedded raw capture for provenance but is excluded from all rendered steady-state views. Each receive_total envelope is divided into RX flag wait, unpack exposed wait, dequant+clear, and output using the same core's fine exposed wall-time. Fine profiling preserves the quant/dequant double-buffer pipeline: hidden asynchronous transfer time is not charged to an exposed-wait stage. Dequant+clear includes next-route unpack submission bookkeeping, but not that transfer's asynchronous in-flight time. Explicit DCCI spans cover combine-kernel UDMACleanCacheLines calls and exclude their surrounding barriers; SQ/CQ cache maintenance inside the UDMA helper remains part of UDMA post/quiet. DCCI category values are accumulated core-time, not additional coarse wall-time. The colored parts show composition, not chronological ordering, and must not be added to infer hardware-engine occupancy. Legacy fine captures without rx_unpack_wait predate this contract. Repeated fine stages use accumulated exposed wall-time in the heatmap, and their durations include material instrumentation overhead. The detailed timeline uses launch0, selected as the coarse sample nearest the median after maximum core per rank and maximum rank per launch. The sample CSV contains all 10 paired launches. Each rank is independently normalized; the report does not assume synchronized device clocks across NPUs. Open perfetto_trace.json for the standard TileXR drill-down.
+
+ + diff --git a/src/comm/CMakeLists.txt b/src/comm/CMakeLists.txt index e0fd5253..222d19f3 100644 --- a/src/comm/CMakeLists.txt +++ b/src/comm/CMakeLists.txt @@ -122,6 +122,9 @@ set(TILEXR_SOURCE_FILE tilexr_comm.cpp add_library(tile-comm SHARED ${TILEXR_SOURCE_FILE}) +target_compile_definitions(tile-comm PRIVATE + TILEXR_UDMA_QP_COUNT_VALUE=${TILEXR_UDMA_EFFECTIVE_QP_COUNT}) + if(TILEXR_HAVE_PTO_SDMA) target_compile_definitions(tile-comm PRIVATE TILEXR_HAVE_PTO_SDMA=1) else() diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index 401b5f6e..2c4c8cc1 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -10,8 +10,6 @@ namespace TileXR { namespace { -constexpr uint32_t TILEXR_UDMA_QP_NUM = 1; - template void CopyVector(std::vector& dst, size_t offset, const std::vector& src) { @@ -32,9 +30,24 @@ int BuildUDMAInfoImage( UDMAInfo& info, std::vector& bytes) { - const size_t rankCount = sq.size(); - if (rankCount == 0 || rq.size() != rankCount || scq.size() != rankCount || - rcq.size() != rankCount || mem.size() != rankCount) { + return BuildUDMAInfoImage(deviceBase, 1, sq, rq, scq, rcq, mem, info, bytes); +} + +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpCount, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + UDMAInfo& info, + std::vector& bytes) +{ + const size_t queueEntryCount = sq.size(); + if (qpCount == 0 || queueEntryCount == 0 || queueEntryCount % qpCount != 0 || + rq.size() != queueEntryCount || scq.size() != queueEntryCount || + rcq.size() != queueEntryCount || mem.size() != queueEntryCount) { return TILEXR_UDMA_LAYOUT_INVALID; } @@ -46,7 +59,7 @@ int BuildUDMAInfoImage( const size_t totalBytes = memOffset + mem.size() * sizeof(UDMAMemInfo); info = {}; - info.qpNum = TILEXR_UDMA_QP_NUM; + info.qpNum = qpCount; info.sqPtr = deviceBase + sqOffset; info.rqPtr = deviceBase + rqOffset; info.scqPtr = deviceBase + scqOffset; diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index fefd49e4..39c94e43 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -26,6 +26,17 @@ int BuildUDMAInfoImage( UDMAInfo& info, std::vector& bytes); +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpCount, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + UDMAInfo& info, + std::vector& bytes); + } // namespace TileXR #endif // TILEXR_UDMA_LAYOUT_H diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 7bafa5f1..5f8dcfbb 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -23,6 +23,11 @@ namespace TileXR { namespace { +#ifndef TILEXR_UDMA_QP_COUNT_VALUE +#define TILEXR_UDMA_QP_COUNT_VALUE 1 +#endif +constexpr uint32_t kTileXRUdmaQpCount = TILEXR_UDMA_QP_COUNT_VALUE; + uint32_t Log2Uint64(uint64_t value) { uint32_t result = 0; @@ -288,11 +293,7 @@ bool ResolveLocalEidRoute( } // namespace -struct TileXRUDMATransport::PerEidState { - uint32_t eidIndex = 0; - void* ctxHandle = nullptr; - void* tokenHandle = nullptr; - void* chanHandle = nullptr; +struct PerQpState { void* cqHandle = nullptr; void* qpHandle = nullptr; CqInfoT cqInfo {}; @@ -309,6 +310,14 @@ struct TileXRUDMATransport::PerEidState { UDMACQCtx localCq {}; }; +struct TileXRUDMATransport::PerEidState { + uint32_t eidIndex = 0; + void* ctxHandle = nullptr; + void* tokenHandle = nullptr; + void* chanHandle = nullptr; + std::vector qps; +}; + TileXRUDMATransport::TileXRUDMATransport() = default; TileXRUDMATransport::~TileXRUDMATransport() @@ -582,12 +591,11 @@ void TileXRUDMATransport::FreeDeviceScalar(void*& ptr) const int TileXRUDMATransport::CreateQueues() { for (const auto& ctxEntry : ctxHandleByEid_) { - PerEidState state {}; + auto& state = states_[ctxEntry.first]; state.eidIndex = ctxEntry.first; state.ctxHandle = ctxEntry.second; state.tokenHandle = tokenHandleByEid_[ctxEntry.first]; - state.remoteQpHandles.assign(options_.rankSize, nullptr); - state.tpnList.assign(options_.rankSize, 0); + state.qps.resize(kTileXRUdmaQpCount); ChanInfoT chanInfo {}; chanInfo.in.dataPlaneFlag.bs.poolCqCstm = 1; @@ -596,91 +604,103 @@ int TileXRUDMATransport::CreateQueues() return TILEXR_ERROR_INTERNAL; } - state.cqInfo.in.chanHandle = state.chanHandle; - state.cqInfo.in.depth = TILEXR_UDMA_CQ_DEPTH; - state.cqInfo.in.ub.mode = JFC_MODE_USER_CTL_NORMAL; - ret = loader_.RaCtxCqCreate(state.ctxHandle, &state.cqInfo, &state.cqHandle); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; - } - state.localCq.cqn = 0; - state.localCq.bufAddr = state.cqInfo.out.bufAddr; - state.localCq.baseBkShift = Log2Uint64(state.cqInfo.out.cqeSize); - state.localCq.depth = state.cqInfo.in.depth; - if (AllocDeviceScalar(&state.cqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.cqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS) { - return TILEXR_ERROR_INTERNAL; - } - state.localCq.headAddr = reinterpret_cast(state.cqPiAddr); - state.localCq.tailAddr = reinterpret_cast(state.cqCiAddr); - state.localCq.dbMode = UDMADBMode::SW_DB; - state.localCq.dbAddr = state.cqInfo.out.swdbAddr; - - QpCreateAttr qpAttr {}; - qpAttr.scqHandle = state.cqHandle; - qpAttr.rcqHandle = state.cqHandle; - qpAttr.srqHandle = state.cqHandle; - qpAttr.sqDepth = TILEXR_UDMA_SQ_DEPTH; - qpAttr.rqDepth = TILEXR_UDMA_RQ_DEPTH_DEFAULT; - qpAttr.transportMode = CONN_RM; - qpAttr.ub.mode = JETTY_MODE_USER_CTL_NORMAL; - qpAttr.ub.flag.value = 1; - qpAttr.ub.jfsFlag.value = 2; - qpAttr.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - qpAttr.ub.rnrRetry = 7; - qpAttr.ub.extMode.piType = 0; - qpAttr.ub.extMode.cstmFlag.bs.sqCstm = 0; - qpAttr.ub.extMode.sqebbNum = TILEXR_UDMA_SQ_DEPTH; - qpAttr.ub.tokenIdHandle = state.tokenHandle; - ret = loader_.RaCtxQpCreate(state.ctxHandle, &qpAttr, &state.qpInfo, &state.qpHandle); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; - } - state.localWq.wqn = 0; - state.localWq.bufAddr = state.qpInfo.ub.sqBuffVa; - state.localWq.baseBkShift = Log2Uint64(state.qpInfo.ub.wqebbSize); - state.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; - if (AllocDeviceScalar(&state.sqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.sqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.wqeCntAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.amoAddr, sizeof(uint64_t)) != TILEXR_SUCCESS) { - return TILEXR_ERROR_INTERNAL; + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + auto& qp = state.qps[qpIdx]; + qp.remoteQpHandles.assign(options_.rankSize, nullptr); + qp.tpnList.assign(options_.rankSize, 0); + + qp.cqInfo.in.chanHandle = state.chanHandle; + qp.cqInfo.in.depth = TILEXR_UDMA_CQ_DEPTH; + qp.cqInfo.in.ub.mode = JFC_MODE_USER_CTL_NORMAL; + ret = loader_.RaCtxCqCreate(state.ctxHandle, &qp.cqInfo, &qp.cqHandle); + if (ret != 0) { + return TILEXR_ERROR_INTERNAL; + } + qp.localCq.cqn = 0; + qp.localCq.bufAddr = qp.cqInfo.out.bufAddr; + qp.localCq.baseBkShift = Log2Uint64(qp.cqInfo.out.cqeSize); + qp.localCq.depth = qp.cqInfo.in.depth; + if (AllocDeviceScalar(&qp.cqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.cqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + qp.localCq.headAddr = reinterpret_cast(qp.cqPiAddr); + qp.localCq.tailAddr = reinterpret_cast(qp.cqCiAddr); + qp.localCq.dbMode = UDMADBMode::SW_DB; + qp.localCq.dbAddr = qp.cqInfo.out.swdbAddr; + + QpCreateAttr qpAttr {}; + qpAttr.scqHandle = qp.cqHandle; + qpAttr.rcqHandle = qp.cqHandle; + qpAttr.srqHandle = qp.cqHandle; + qpAttr.sqDepth = TILEXR_UDMA_SQ_DEPTH; + qpAttr.rqDepth = TILEXR_UDMA_RQ_DEPTH_DEFAULT; + qpAttr.transportMode = CONN_RM; + qpAttr.ub.mode = JETTY_MODE_USER_CTL_NORMAL; + qpAttr.ub.flag.value = 1; + qpAttr.ub.jfsFlag.value = 2; + qpAttr.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + qpAttr.ub.rnrRetry = 7; + qpAttr.ub.extMode.piType = 0; + qpAttr.ub.extMode.cstmFlag.bs.sqCstm = 0; + qpAttr.ub.extMode.sqebbNum = TILEXR_UDMA_SQ_DEPTH; + qpAttr.ub.tokenIdHandle = state.tokenHandle; + ret = loader_.RaCtxQpCreate(state.ctxHandle, &qpAttr, &qp.qpInfo, &qp.qpHandle); + if (ret != 0) { + return TILEXR_ERROR_INTERNAL; + } + qp.localWq.wqn = 0; + qp.localWq.bufAddr = qp.qpInfo.ub.sqBuffVa; + qp.localWq.baseBkShift = Log2Uint64(qp.qpInfo.ub.wqebbSize); + qp.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; + if (AllocDeviceScalar(&qp.sqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.sqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.wqeCntAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.amoAddr, sizeof(uint64_t)) != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + qp.localWq.headAddr = reinterpret_cast(qp.sqPiAddr); + qp.localWq.tailAddr = reinterpret_cast(qp.sqCiAddr); + qp.localWq.dbMode = UDMADBMode::SW_DB; + qp.localWq.dbAddr = qp.qpInfo.ub.dbAddr; + qp.localWq.wqeCntAddr = reinterpret_cast(qp.wqeCntAddr); + qp.localWq.amoAddr = reinterpret_cast(qp.amoAddr); } - state.localWq.headAddr = reinterpret_cast(state.sqPiAddr); - state.localWq.tailAddr = reinterpret_cast(state.sqCiAddr); - state.localWq.dbMode = UDMADBMode::SW_DB; - state.localWq.dbAddr = state.qpInfo.ub.dbAddr; - state.localWq.wqeCntAddr = reinterpret_cast(state.wqeCntAddr); - state.localWq.amoAddr = reinterpret_cast(state.amoAddr); - states_[state.eidIndex] = state; } return states_.empty() ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; } int TileXRUDMATransport::ImportQueues() { - std::vector localImports(eidCount_); - std::vector localKeys(eidCount_); + const size_t entriesPerRank = static_cast(eidCount_) * kTileXRUdmaQpCount; + std::vector localImports(entriesPerRank); + std::vector localKeys(entriesPerRank); for (const auto& stateEntry : states_) { const auto& state = stateEntry.second; if (state.eidIndex >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - localImports[state.eidIndex].in.ub.mode = JETTY_IMPORT_MODE_NORMAL; - localImports[state.eidIndex].in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - localImports[state.eidIndex].in.ub.policy = JETTY_GRP_POLICY_RR; - localImports[state.eidIndex].in.ub.type = TARGET_TYPE_JETTY; - localImports[state.eidIndex].in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; - localImports[state.eidIndex].in.ub.tpType = 1; - localKeys[state.eidIndex] = state.qpInfo.key; + if (state.qps.size() != kTileXRUdmaQpCount) { + return TILEXR_ERROR_INTERNAL; + } + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + const size_t index = static_cast(state.eidIndex) * kTileXRUdmaQpCount + qpIdx; + localImports[index].in.ub.mode = JETTY_IMPORT_MODE_NORMAL; + localImports[index].in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + localImports[index].in.ub.policy = JETTY_GRP_POLICY_RR; + localImports[index].in.ub.type = TARGET_TYPE_JETTY; + localImports[index].in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; + localImports[index].in.ub.tpType = 1; + localKeys[index] = state.qps[qpIdx].qpInfo.key; + } } - std::vector allImports(options_.rankSize * eidCount_); + std::vector allImports(options_.rankSize * entriesPerRank); int ret = options_.exchange->AllGather(localImports.data(), localImports.size(), allImports.data()); if (ret != TILEXR_SUCCESS) { return ret; } - std::vector allKeys(options_.rankSize * eidCount_); + std::vector allKeys(options_.rankSize * entriesPerRank); ret = options_.exchange->AllGather(localKeys.data(), localKeys.size(), allKeys.data()); if (ret != TILEXR_SUCCESS) { return ret; @@ -700,13 +720,18 @@ int TileXRUDMATransport::ImportQueues() if (remoteEid >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - QpImportInfoT importInfo = allImports[peer * eidCount_ + remoteEid]; - importInfo.in.key = allKeys[peer * eidCount_ + remoteEid]; - ret = loader_.RaCtxQpImport(state.ctxHandle, &importInfo, &state.remoteQpHandles[peer]); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + const size_t index = static_cast(peer) * entriesPerRank + + static_cast(remoteEid) * kTileXRUdmaQpCount + qpIdx; + QpImportInfoT importInfo = allImports[index]; + importInfo.in.key = allKeys[index]; + auto& qp = state.qps[qpIdx]; + ret = loader_.RaCtxQpImport(state.ctxHandle, &importInfo, &qp.remoteQpHandles[peer]); + if (ret != 0 || qp.remoteQpHandles[peer] == nullptr) { + return TILEXR_ERROR_INTERNAL; + } + qp.tpnList[peer] = importInfo.out.ub.tpn; } - state.tpnList[peer] = importInfo.out.ub.tpn; } } return TILEXR_SUCCESS; @@ -768,11 +793,12 @@ int TileXRUDMATransport::RefreshUDMAInfo() return TILEXR_ERROR_INTERNAL; } - std::vector sq(options_.rankSize); - std::vector rq(options_.rankSize); - std::vector scq(options_.rankSize); - std::vector rcq(options_.rankSize); - std::vector mem(options_.rankSize); + const size_t queueEntryCount = static_cast(options_.rankSize) * kTileXRUdmaQpCount; + std::vector sq(queueEntryCount); + std::vector rq(queueEntryCount); + std::vector scq(queueEntryCount); + std::vector rcq(queueEntryCount); + std::vector mem(queueEntryCount); for (int rank = 0; rank < options_.rankSize; ++rank) { uint32_t localEid = fallbackEid; @@ -786,26 +812,37 @@ int TileXRUDMATransport::RefreshUDMAInfo() stateIt = fallbackIt; } const auto& state = stateIt->second; - sq[rank] = state.localWq; - rq[rank] = state.localWq; - scq[rank] = state.localCq; - rcq[rank] = state.localCq; + if (state.qps.size() != kTileXRUdmaQpCount) { + return TILEXR_ERROR_INTERNAL; + } + UDMAMemInfo rankMem {}; if (rank == options_.rank) { const auto localMemIt = localMemInfoByEid_.find(localEid); if (localMemIt != localMemInfoByEid_.end()) { - mem[rank] = localMemIt->second; + rankMem = localMemIt->second; } } else { - mem[rank] = allMem[rank * eidCount_ + remoteEid]; - mem[rank].tpn = state.tpnList[rank]; + rankMem = allMem[rank * eidCount_ + remoteEid]; } - mem[rank].eidAddr = reinterpret_cast( + rankMem.eidAddr = reinterpret_cast( eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + const size_t index = static_cast(rank) * kTileXRUdmaQpCount + qpIdx; + const auto& qp = state.qps[qpIdx]; + sq[index] = qp.localWq; + rq[index] = qp.localWq; + scq[index] = qp.localCq; + rcq[index] = qp.localCq; + mem[index] = rankMem; + if (rank != options_.rank) { + mem[index].tpn = qp.tpnList[rank]; + } + } } if (udmaInfoDev_ == nullptr) { - const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo); - udmaInfoSize_ = static_cast(sizeof(UDMAInfo) + oneRankSize * options_.rankSize); + const size_t oneQueueSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo); + udmaInfoSize_ = static_cast(sizeof(UDMAInfo) + oneQueueSize * queueEntryCount); ret = aclrtMalloc(reinterpret_cast(&udmaInfoDev_), udmaInfoSize_, ACL_MEM_MALLOC_HUGE_FIRST); if (ret != ACL_SUCCESS) { return TILEXR_ERROR_INTERNAL; @@ -814,7 +851,8 @@ int TileXRUDMATransport::RefreshUDMAInfo() UDMAInfo info {}; std::vector image; - ret = BuildUDMAInfoImage(reinterpret_cast(udmaInfoDev_), sq, rq, scq, rcq, mem, info, image); + ret = BuildUDMAInfoImage(reinterpret_cast(udmaInfoDev_), kTileXRUdmaQpCount, + sq, rq, scq, rcq, mem, info, image); if (ret != TILEXR_UDMA_LAYOUT_SUCCESS) { return TILEXR_ERROR_PARA_CHECK_FAIL; } @@ -1002,26 +1040,28 @@ void TileXRUDMATransport::CleanupQueues() { for (auto& stateEntry : states_) { auto& state = stateEntry.second; - for (void* remoteQp : state.remoteQpHandles) { - if (remoteQp != nullptr && state.ctxHandle != nullptr) { - loader_.RaCtxQpUnimport(state.ctxHandle, remoteQp); + for (auto& qp : state.qps) { + for (void* remoteQp : qp.remoteQpHandles) { + if (remoteQp != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxQpUnimport(state.ctxHandle, remoteQp); + } } - } - if (state.qpHandle != nullptr) { - loader_.RaCtxQpDestroy(state.qpHandle); - } - if (state.cqHandle != nullptr && state.ctxHandle != nullptr) { - loader_.RaCtxCqDestroy(state.ctxHandle, state.cqHandle); + if (qp.qpHandle != nullptr) { + loader_.RaCtxQpDestroy(qp.qpHandle); + } + if (qp.cqHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxCqDestroy(state.ctxHandle, qp.cqHandle); + } + FreeDeviceScalar(qp.cqPiAddr); + FreeDeviceScalar(qp.cqCiAddr); + FreeDeviceScalar(qp.sqPiAddr); + FreeDeviceScalar(qp.sqCiAddr); + FreeDeviceScalar(qp.wqeCntAddr); + FreeDeviceScalar(qp.amoAddr); } if (state.chanHandle != nullptr && state.ctxHandle != nullptr) { loader_.RaCtxChanDestroy(state.ctxHandle, state.chanHandle); } - FreeDeviceScalar(state.cqPiAddr); - FreeDeviceScalar(state.cqCiAddr); - FreeDeviceScalar(state.sqPiAddr); - FreeDeviceScalar(state.sqCiAddr); - FreeDeviceScalar(state.wqeCntAddr); - FreeDeviceScalar(state.amoAddr); } states_.clear(); } diff --git a/src/ep/CMakeLists.txt b/src/ep/CMakeLists.txt index c2fc93d5..8cfbc919 100644 --- a/src/ep/CMakeLists.txt +++ b/src/ep/CMakeLists.txt @@ -1,5 +1,14 @@ include(GNUInstallDirs) +option(TILEXR_EP_ENABLE_PROFILING "Enable EP URMA combine kernel profiling" OFF) + +# The protocol constants live in ep_urma_combine_config.h and are fixed +# to the validated Ascend 950 S22 production configuration. +set(TILEXR_EP_PROFILE_COMPILE_DEFINITION) +if(TILEXR_EP_ENABLE_PROFILING) + set(TILEXR_EP_PROFILE_COMPILE_DEFINITION -DTILEXR_EP_ENABLE_PROFILING=1) +endif() + find_program(BISHENG_EXECUTABLE bisheng) if(NOT BISHENG_EXECUTABLE) message(FATAL_ERROR "bisheng not found; source scripts/common_env.sh before configuring with -DTILEXR_BUILD_EP=ON") @@ -12,7 +21,9 @@ endif() set(TILEXR_EP_SOC_TYPE "${TILEXR_EP_DEFAULT_SOC_TYPE}" CACHE STRING "TileXR EP kernel SOC type") string(TOLOWER "${TILEXR_EP_SOC_TYPE}" TILEXR_EP_SOC_TYPE_LOWER) +set(TILEXR_EP_URMA_COMBINE_SUPPORTED OFF) if(TILEXR_EP_SOC_TYPE_LOWER MATCHES "ascend950|ascend910_9|ascend910-9|a5") + set(TILEXR_EP_URMA_COMBINE_SUPPORTED ON) set(TILEXR_EP_NPU_ARCH "dav-3510") set(TILEXR_EP_AICORE_ARCH "--cce-aicore-arch=dav-c310-vec") set(TILEXR_EP_CATLASS_ARCH "3510") @@ -25,6 +36,10 @@ else() set(TILEXR_EP_AICORE_ARCH "--cce-aicore-arch=dav-c220-vec") set(TILEXR_EP_CATLASS_ARCH "2201") endif() +if(TILEXR_EP_BUILD_URMA_COMBINE AND NOT TILEXR_EP_URMA_COMBINE_SUPPORTED) + message(FATAL_ERROR + "TILEXR_EP_BUILD_URMA_COMBINE requires an Ascend950/A5 EP target") +endif() set(TILEXR_EP_KERNEL_COMPILE_OPTIONS -xcce @@ -40,8 +55,10 @@ set(TILEXR_EP_KERNEL_LINK_OPTIONS ${TILEXR_EP_AICORE_ARCH}) set(TILEXR_EP_DISPATCH_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_dispatch_kernel.cpp") set(TILEXR_EP_COMBINE_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_combine_kernel.cpp") +set(TILEXR_EP_URMA_COMBINE_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_urma_combine_kernel.cpp") set(TILEXR_EP_DISPATCH_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_dispatch_kernel.so") set(TILEXR_EP_COMBINE_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_combine_kernel.so") +set(TILEXR_EP_URMA_COMBINE_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_urma_combine_kernel.so") set(TILEXR_EP_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/compiler/tikcpp -I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw @@ -55,6 +72,8 @@ set(TILEXR_EP_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/runtime -I${ASCEND_HOME_PATH}/${ARCH}-linux/runtime/include -I${ASCEND_HOME_PATH}/${ARCH}-linux/include + -I${ASCEND_HOME_PATH}/${ARCH}-linux/asc/include + -I${ASCEND_HOME_PATH}/${ARCH}-linux/asc/impl -I${ASCEND_DRIVER_PATH}/kernel/inc -I${CMAKE_SOURCE_DIR}/3rdparty -I${CMAKE_SOURCE_DIR}/src/include @@ -135,15 +154,82 @@ add_custom_command( ) add_custom_target(tilexr_ep_combine_kernel ALL DEPENDS "${TILEXR_EP_COMBINE_KERNEL_SO}") -add_library(tilexr-ep SHARED +if(TILEXR_EP_BUILD_URMA_COMBINE) + add_custom_command( + OUTPUT "${TILEXR_EP_URMA_COMBINE_KERNEL_SO}" + COMMAND ${BISHENG_EXECUTABLE} + ${TILEXR_EP_KERNEL_COMPILE_OPTIONS} + -O2 + -std=gnu++17 + -fPIC + -shared + ${TILEXR_EP_KERNEL_LINK_OPTIONS} + -DCATLASS_ARCH=${TILEXR_EP_CATLASS_ARCH} + -DTILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT=${TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT} + ${TILEXR_EP_PROFILE_COMPILE_DEFINITION} + ${TILEXR_EP_KERNEL_INCLUDES} + "${TILEXR_EP_URMA_COMBINE_KERNEL_SOURCE}" + -L${ASCEND_DRIVER_PATH}/lib64/driver + -L${ASCEND_HOME_PATH}/${ARCH}-linux/lib64 + -L${ASCEND_HOME_PATH}/${ARCH}-linux/devlib + -lruntime + -lascendcl + -lstdc++ + -lm + -ltiling_api + -lplatform + -lc_sec + -ldl + -lnnopbase + -lpthread + -o "${TILEXR_EP_URMA_COMBINE_KERNEL_SO}" + DEPENDS + "${TILEXR_EP_URMA_COMBINE_KERNEL_SOURCE}" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_combine_config.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_combine.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_perf_trace.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_udma.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_window.h" + "${CMAKE_SOURCE_DIR}/src/include/comm_args.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_data_as_flag.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_perf_trace.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_udma.h" + VERBATIM + COMMENT "Building TileXR EP URMA combine kernel (Ascend 950 S22, O2)" + ) + add_custom_target(tilexr_ep_urma_combine_kernel ALL + DEPENDS "${TILEXR_EP_URMA_COMBINE_KERNEL_SO}") +endif() + +set(TILEXR_EP_HOST_SOURCES host/ep_layout.cpp host/ep_dispatch_host.cpp host/ep_launch_context.cpp host/ep_kernel_launch.cpp host/tilexr_ep_dispatch.cpp ) +if(TILEXR_EP_BUILD_URMA_COMBINE) + list(APPEND TILEXR_EP_HOST_SOURCES + host/ep_urma_combine_layout.cpp + host/ep_urma_combine_host.cpp + host/tilexr_ep_urma_combine.cpp) +endif() +add_library(tilexr-ep SHARED ${TILEXR_EP_HOST_SOURCES}) -add_dependencies(tilexr-ep tilexr_ep_dispatch_kernel tilexr_ep_combine_kernel) +add_dependencies(tilexr-ep + tilexr_ep_dispatch_kernel + tilexr_ep_combine_kernel) +if(TILEXR_EP_BUILD_URMA_COMBINE) + add_dependencies(tilexr-ep tilexr_ep_urma_combine_kernel) +endif() + +if(TILEXR_EP_BUILD_URMA_COMBINE) + target_compile_definitions(tilexr-ep PRIVATE + TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT=${TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT}) +endif() +if(TILEXR_EP_BUILD_URMA_COMBINE AND TILEXR_EP_ENABLE_PROFILING) + target_compile_definitions(tilexr-ep PRIVATE TILEXR_EP_ENABLE_PROFILING=1) +endif() target_include_directories(tilexr-ep PUBLIC @@ -170,6 +256,9 @@ target_link_libraries(tilexr-ep runtime ascend_hal ) +if(TILEXR_EP_BUILD_URMA_COMBINE) + target_link_libraries(tilexr-ep PRIVATE -l:libtilexr_ep_urma_combine_kernel.so) +endif() set_target_properties(tilexr-ep PROPERTIES BUILD_RPATH "${CMAKE_CURRENT_BINARY_DIR}" @@ -182,4 +271,9 @@ install(FILES "${TILEXR_EP_COMBINE_KERNEL_SO}" DESTINATION ${CMAKE_INSTALL_LIBDIR} ) -install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_ep.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +if(TILEXR_EP_BUILD_URMA_COMBINE) + install(FILES "${TILEXR_EP_URMA_COMBINE_KERNEL_SO}" + DESTINATION ${CMAKE_INSTALL_LIBDIR}) +endif() +install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_ep.h + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) diff --git a/src/ep/common/ep_urma_combine.h b/src/ep/common/ep_urma_combine.h new file mode 100644 index 00000000..8b889b9c --- /dev/null +++ b/src/ep/common/ep_urma_combine.h @@ -0,0 +1,66 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_COMBINE_H +#define TILEXR_EP_COMMON_EP_URMA_COMBINE_H + +#include + +#include "ep_urma_combine_config.h" + +namespace TileXREp { + +constexpr int64_t kEpUrmaCombineWorkspaceAlignment = 2 * 1024 * 1024; +constexpr int64_t kEpUrmaCombineHeaderBytes = 512; +constexpr int64_t kEpUrmaCombineCacheLineBytes = 64; +constexpr int64_t kEpUrmaCombineDataBlockBytes = 512; +constexpr int64_t kEpUrmaCombinePayloadBytes = 480; +constexpr int64_t kEpUrmaCombineFlagBytes = 32; +constexpr int64_t kEpUrmaCombineQuantHeaderBytes = 32; +constexpr int64_t kEpUrmaCombineAivCount = kEpUrmaCombineProfileCoreCount; +constexpr int64_t kEpUrmaCombinePackLaneCount = kEpUrmaCombineProfilePackReceiveCoreCount; +constexpr int64_t kEpUrmaCombineSendLaneCount = kEpUrmaCombineProfileSendCoreCount; +constexpr int64_t kEpUrmaCombineRequiredQpCount = kEpUrmaCombineSendLaneCount; +constexpr int64_t kEpUrmaCombineMaxHidden = 8192; +constexpr int64_t kEpUrmaCombineMaxTopK = 16; +constexpr int64_t kEpUrmaCombineMaxBlocksPerRoute = 255; +constexpr int64_t kEpUrmaCombineQuantModeInt8PerRoute = 1; + +constexpr uint32_t kEpUrmaCombineTxRouteReady = 1; +constexpr uint32_t kEpUrmaCombineRxLaneDone = 2; +constexpr uint32_t kEpUrmaCombineSenderDone = 3; +constexpr uint32_t kEpUrmaCombineRxBufferReleased = 4; +constexpr uint32_t kEpUrmaCombinePublishDone = 5; +constexpr uint32_t kEpUrmaCombineStartLocalReady = 6; +constexpr uint32_t kEpUrmaCombineStartRankReady = 7; +constexpr uint32_t kEpUrmaCombineStartPublishDone = 8; +constexpr uint32_t kEpUrmaCombineStartRun = 9; +constexpr uint32_t kEpUrmaCombineCreditExpectedReady = 10; +constexpr uint32_t kEpUrmaCombineCreditShardDone = 11; +constexpr uint32_t kEpUrmaCombineCreditRun = 12; +constexpr uint32_t kEpUrmaCombineRxReleaseShardDone = 13; + +constexpr uint64_t kEpUrmaCombineStatusOk = 0; +constexpr uint64_t kEpUrmaCombineStatusInvalidRoute = 1; +constexpr uint64_t kEpUrmaCombineStatusInvalidQuantHeader = 2; + +struct alignas(kEpUrmaCombineQuantHeaderBytes) EpUrmaCombineQuantHeader { + float scale; + int32_t quantMode; + int64_t reserved0; + int64_t reserved1; + int64_t reserved2; +}; + +struct alignas(kEpUrmaCombineCacheLineBytes) EpUrmaCombineControlLine { + uint64_t value; + uint64_t reserved[7]; +}; + +static_assert(sizeof(EpUrmaCombineQuantHeader) == kEpUrmaCombineQuantHeaderBytes, + "URMA combine quant header must be 32 bytes"); +static_assert(sizeof(EpUrmaCombineControlLine) == kEpUrmaCombineCacheLineBytes, + "URMA combine control line must be one cache line"); +static_assert(kEpUrmaCombinePackLaneCount + kEpUrmaCombineSendLaneCount == kEpUrmaCombineAivCount, + "URMA combine must occupy all configured AIV cores"); + +} // namespace TileXREp + +#endif // TILEXR_EP_COMMON_EP_URMA_COMBINE_H diff --git a/src/ep/common/ep_urma_combine_config.h b/src/ep/common/ep_urma_combine_config.h new file mode 100644 index 00000000..9b7de5c1 --- /dev/null +++ b/src/ep/common/ep_urma_combine_config.h @@ -0,0 +1,91 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_COMBINE_CONFIG_H +#define TILEXR_EP_COMMON_EP_URMA_COMBINE_CONFIG_H + +#include + +namespace TileXREp { + +// Validated Ascend 950 production protocol. CMake may override only the Send AIV count. +#ifndef TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT +#define TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT 22 +#endif + +constexpr uint32_t kEpUrmaCombineProfileCoreCount = 64; +constexpr uint32_t kEpUrmaCombineProfileSendCoreCount = TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT; +constexpr uint32_t kEpUrmaCombineProfilePackReceiveCoreCount = + kEpUrmaCombineProfileCoreCount - kEpUrmaCombineProfileSendCoreCount; + +enum class EpUrmaCombinePerfStage : uint32_t { + KERNEL_TOTAL = 0, + PACK_TOTAL = 1, + PACK_INPUT_WAIT = 2, + PACK_QUANTIZE = 3, + PACK_TX_PUBLISH = 4, + RECEIVE_TOTAL = 5, + RX_FLAG_POLL_WAIT = 6, + RX_UNPACK_WAIT = 7, + RX_UNPACK_DEQUANT_CLEAR = 8, + RX_OUTPUT = 9, + SEND_TOTAL = 10, + TX_META_SCAN = 11, + TX_READY_POLL = 12, + SELF_COPY = 13, + UDMA_POST = 14, + UDMA_QUIET = 15, + LOCAL_SENDER_WAIT = 16, + LOCAL_RX_WAIT = 17, + ROUND_PUBLISH = 18, + GLOBAL_ROUND_WAIT = 19, + DCCI_TOTAL = 20, + START_GATE = 21, + PACK_TX_DATA_SUBMIT = 22, + PACK_FIRST_TX_READY = 23, + PACK_MTE3_EXPOSED_WAIT = 24, + RX_READY_MTE2_WAIT = 25, + RX_READY_VECTOR = 26, +}; + +constexpr uint32_t kEpUrmaCombinePerfStageCount = 27; + +#if !defined(__CCE__) || !defined(__CCE_IS_AICORE__) +constexpr const char *const kEpUrmaCombinePerfStageNames[kEpUrmaCombinePerfStageCount] = { + "kernel_total", + "pack_total", + "pack_input_wait", + "pack_quantize", + "pack_tx_publish", + "receive_total", + "rx_flag_poll_wait", + "rx_unpack_wait", + "rx_unpack_dequant_clear", + "rx_output", + "send_total", + "tx_meta_scan", + "tx_ready_poll", + "self_copy", + "udma_post", + "udma_quiet", + "local_sender_wait", + "local_rx_wait", + "round_publish", + "global_round_wait", + "dcci_total", + "start_gate", + "pack_tx_data_submit", + "pack_first_tx_ready", + "pack_mte3_exposed_wait", + "rx_ready_mte2_wait", + "rx_ready_vector", +}; +#endif + +static_assert(kEpUrmaCombineProfilePackReceiveCoreCount + kEpUrmaCombineProfileSendCoreCount == + kEpUrmaCombineProfileCoreCount, + "URMA combine profile core roles must cover all AIV cores"); +static_assert(kEpUrmaCombineProfileSendCoreCount > 0 && + kEpUrmaCombineProfileSendCoreCount < kEpUrmaCombineProfileCoreCount, + "URMA combine Send core count must leave at least one Pack/Receive core"); + +} // namespace TileXREp + +#endif // TILEXR_EP_COMMON_EP_URMA_COMBINE_CONFIG_H diff --git a/src/ep/common/ep_urma_perf_trace.h b/src/ep/common/ep_urma_perf_trace.h new file mode 100644 index 00000000..7433984c --- /dev/null +++ b/src/ep/common/ep_urma_perf_trace.h @@ -0,0 +1,176 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_PERF_TRACE_H +#define TILEXR_EP_COMMON_EP_URMA_PERF_TRACE_H + +#include "kernel_operator.h" +#include "tilexr_perf_trace.h" + +namespace TileXR { + +constexpr uint32_t TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT = 27; +constexpr uint32_t TILEXR_PERF_TRACE_LOCAL_STATS_UB_END = 195712; +constexpr uint32_t TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET = + TILEXR_PERF_TRACE_LOCAL_STATS_UB_END - + TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT * sizeof(TileXRPerfCoreStageStats); +constexpr uint32_t TILEXR_PERF_TRACE_MIN_UB_BYTES = 192 * 1024; + +static_assert(TILEXR_PERF_TRACE_LOCAL_STATS_UB_END <= TILEXR_PERF_TRACE_MIN_UB_BYTES, + "local perf stats must fit the minimum supported AIV UB"); +static_assert(TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET % 32 == 0, + "local perf stats must remain 32-byte aligned"); + +#if defined(TILEXR_EP_ENABLE_PROFILING) + +__attribute__((always_inline)) inline __aicore__ bool TileXRPerfTraceEnabled(GM_ADDR trace) +{ + return trace != nullptr; +} + +__attribute__((always_inline)) inline __aicore__ uint64_t TileXRPerfCycleNow(GM_ADDR trace) +{ + return trace == nullptr ? 0 : static_cast(AscendC::GetSystemCycle()); +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsInit( + GM_ADDR trace, __ubuf__ TileXRPerfCoreStageStats *statsUB, uint32_t rank, uint32_t core, + uint32_t stageCount) +{ + if (trace == nullptr || stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT) { + return; + } + for (uint32_t stage = 0; stage < stageCount; ++stage) { + statsUB[stage].rank = rank; + statsUB[stage].core = core; + statsUB[stage].stageId = stage; + statsUB[stage].reserved = 0; + statsUB[stage].count = 0; + statsUB[stage].sumCycles = 0; + statsUB[stage].minCycles = 0; + statsUB[stage].maxCycles = 0; + statsUB[stage].firstStartCycle = 0; + statsUB[stage].lastEndCycle = 0; + statsUB[stage].aux0 = 0; + statsUB[stage].aux1 = 0; + statsUB[stage].aux2 = 0; + statsUB[stage].aux3 = 0; + } +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalRecord( + GM_ADDR trace, __ubuf__ TileXRPerfCoreStageStats *statsUB, uint32_t stageCount, + uint32_t stage, uint64_t startCycle, uint64_t endCycle) +{ + if (trace == nullptr || stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT || + stage >= stageCount || endCycle < startCycle) { + return; + } + __ubuf__ TileXRPerfCoreStageStats *stat = statsUB + stage; + const uint64_t duration = endCycle - startCycle; + if (stat->count == 0) { + stat->minCycles = duration; + stat->maxCycles = duration; + stat->firstStartCycle = startCycle; + } else { + if (duration < stat->minCycles) { + stat->minCycles = duration; + } + if (duration > stat->maxCycles) { + stat->maxCycles = duration; + } + if (startCycle < stat->firstStartCycle) { + stat->firstStartCycle = startCycle; + } + } + stat->count += 1; + stat->sumCycles += duration; + if (endCycle > stat->lastEndCycle) { + stat->lastEndCycle = endCycle; + } +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalAddAux( + GM_ADDR trace, __ubuf__ TileXRPerfCoreStageStats *statsUB, uint32_t stageCount, + uint32_t stage, uint32_t auxIndex, uint64_t value) +{ + if (trace == nullptr || stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT || stage >= stageCount) { + return; + } + __ubuf__ TileXRPerfCoreStageStats *stat = statsUB + stage; + if (auxIndex == 0) { + stat->aux0 += value; + } else if (auxIndex == 1) { + stat->aux1 += value; + } else if (auxIndex == 2) { + stat->aux2 += value; + } else if (auxIndex == 3) { + stat->aux3 += value; + } +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsFlush( + GM_ADDR trace, uint32_t rank, uint32_t core, uint32_t maxCoreCount, uint32_t stageCount, + __ubuf__ TileXRPerfCoreStageStats *statsUB) +{ + if (trace == nullptr || stageCount == 0 || + stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT) { + return; + } + const size_t firstSlot = PerfTraceStatsOffset(rank, core, 0, maxCoreCount, stageCount); + __gm__ uint8_t *dstAddr = reinterpret_cast<__gm__ uint8_t *>( + reinterpret_cast<__gm__ TileXRPerfCoreStageStats *>( + trace + TILEXR_PERF_TRACE_STATS_OFFSET) + firstSlot); + AscendC::LocalTensor local; + AscendC::TBuffAddr localAddr; + localAddr.logicPos = static_cast(AscendC::TPosition::VECIN); + localAddr.bufferAddr = reinterpret_cast(statsUB); + local.SetAddr(localAddr); + AscendC::GlobalTensor dst; + dst.SetGlobalBuffer(dstAddr); + AscendC::DataCopyExtParams params { + 1, stageCount * static_cast(sizeof(TileXRPerfCoreStageStats)), 0, 0, 0}; + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad(dst, local, params); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + (void)statsUB; +} + +#else + +__attribute__((always_inline)) inline __aicore__ bool TileXRPerfTraceEnabled(GM_ADDR trace) +{ + (void)trace; + return false; +} + +__attribute__((always_inline)) inline __aicore__ uint64_t TileXRPerfCycleNow(GM_ADDR trace) +{ + (void)trace; + return 0; +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsInit( + GM_ADDR, __ubuf__ TileXRPerfCoreStageStats *, uint32_t, uint32_t, uint32_t) +{ +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalRecord( + GM_ADDR, __ubuf__ TileXRPerfCoreStageStats *, uint32_t, uint32_t, uint64_t, uint64_t) +{ +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalAddAux( + GM_ADDR, __ubuf__ TileXRPerfCoreStageStats *, uint32_t, uint32_t, uint32_t, uint64_t) +{ +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsFlush( + GM_ADDR, uint32_t, uint32_t, uint32_t, uint32_t, __ubuf__ TileXRPerfCoreStageStats *) +{ +} + +#endif + +} // namespace TileXR + +#endif // TILEXR_EP_COMMON_EP_URMA_PERF_TRACE_H diff --git a/src/ep/common/ep_urma_udma.h b/src/ep/common/ep_urma_udma.h new file mode 100644 index 00000000..5d575925 --- /dev/null +++ b/src/ep/common/ep_urma_udma.h @@ -0,0 +1,160 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_UDMA_H +#define TILEXR_EP_COMMON_EP_URMA_UDMA_H + +#include "tilexr_udma.h" + +namespace TileXREp { + +// S22 gives each Send AIV an independent QP. Keep the optimized WQE/CQE path +// private to the operator instead of changing TileXR's shared UDMA helpers. +__aicore__ inline __gm__ TileXR::UDMAMemInfo *EpUrmaGetRemoteMemInfo( + __gm__ TileXR::UDMAInfo *info, uint32_t rank, uint32_t qpIdx) +{ + return reinterpret_cast<__gm__ TileXR::UDMAMemInfo *>( + info->memPtr + sizeof(TileXR::UDMAMemInfo) * (rank * info->qpNum + qpIdx)); +} + +__aicore__ inline uint32_t EpUrmaPollCQ( + __gm__ TileXR::UDMAInfo *info, uint32_t rank, uint32_t qpIdx, uint32_t target) +{ + if (target == 0) { + return 0; + } + __gm__ TileXR::UDMACQCtx *cq = TileXR::UDMAGetSCQCtx(info, rank, qpIdx); + __gm__ TileXR::UDMAWQCtx *sq = TileXR::UDMAGetWQCtx(info, rank, qpIdx); + const uint64_t cqBase = cq->bufAddr; + const uint32_t cqeSize = 1U << cq->baseBkShift; + uint32_t tail = ld_dev(reinterpret_cast<__gm__ uint32_t *>(cq->tailAddr), 0); + while (tail != target) { + __gm__ uint32_t *cqeWords = reinterpret_cast<__gm__ uint32_t *>( + cqBase + cqeSize * (tail & (TileXR::TILEXR_UDMA_CQ_DEPTH - 1))); + const bool validOwner = ((tail / TileXR::TILEXR_UDMA_CQ_DEPTH) & 1U) != 0; + uint32_t word0 = 0; + uint32_t retries = 0; + while (retries < TileXR::TILEXR_UDMA_MAX_RETRY_TIMES) { + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(cqeWords), sizeof(TileXR::UDMACqeCtx)); + __asm__ __volatile__("" ::: "memory"); + word0 = cqeWords[0]; + __asm__ __volatile__("" ::: "memory"); + if ((validOwner ^ ((word0 & (1U << 2U)) != 0)) != 0) { + break; + } + ++retries; + } + if (retries >= TileXR::TILEXR_UDMA_MAX_RETRY_TIMES) { + return 0xFFU; + } + const uint8_t status = static_cast((word0 >> 24U) & 0xFFU); + const uint8_t subStatus = static_cast((word0 >> 16U) & 0xFFU); + if (status != 0 || subStatus != 0) { + return (static_cast(status) << 8U) | subStatus; + } + ++tail; + } + TileXR::UDMAPollCQUpdateInfo(tail, cq, sq); + return 0; +} + +__aicore__ inline void EpUrmaFillWriteWqe( + __gm__ uint8_t *wqe, __gm__ uint8_t *remoteAddr, + __gm__ TileXR::UDMAMemInfo *remote, uint32_t head, uint32_t depth, + uint64_t bytes, __gm__ uint8_t *localAddr) +{ + __gm__ uint32_t *sqe = reinterpret_cast<__gm__ uint32_t *>(wqe); + const uint32_t owner = (head & depth) == 0U ? 1U : 0U; + sqe[0] = (head % depth) | (0b00100010U << 16U) | + ((remote->tokenValueValid ? 1U : 0U) << 28U) | + ((remote->rmtJettyType & 0x3U) << 29U) | (owner << 31U); + sqe[1] = (static_cast(remote->targetHint) & 0xFFU) | + (static_cast(TileXR::UDMAOpcode::WRITE) << 8U); + sqe[2] = (remote->tpn & 0xFFFFFFU) | (1U << 24U); + sqe[3] = remote->tid & 0xFFFFFU; + __gm__ uint32_t *eid = reinterpret_cast<__gm__ uint32_t *>(remote->eidAddr); + sqe[4] = eid[0]; + sqe[5] = eid[1]; + sqe[6] = eid[2]; + sqe[7] = eid[3]; + sqe[8] = remote->rmtTokenValue; + sqe[9] = 0; + const uint64_t remoteValue = reinterpret_cast(remoteAddr); + sqe[10] = static_cast(remoteValue); + sqe[11] = static_cast(remoteValue >> 32U); + + __gm__ uint32_t *sge = sqe + sizeof(TileXR::UDMASqeCtx) / sizeof(uint32_t); + const uint64_t localValue = reinterpret_cast(localAddr); + sge[0] = static_cast(bytes); + sge[1] = 0; + sge[2] = static_cast(localValue); + sge[3] = static_cast(localValue >> 32U); +} + +__aicore__ inline void EpUrmaPostWrite( + __gm__ TileXR::UDMAInfo *info, __gm__ uint8_t *remoteAddr, + __gm__ uint8_t *localAddr, uint32_t rank, uint32_t qpIdx, uint64_t bytes) +{ + __gm__ TileXR::UDMAWQCtx *sq = TileXR::UDMAGetWQCtx(info, rank, qpIdx); + const uint32_t wqeSize = 1U << sq->baseBkShift; + uint32_t head = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->headAddr), 0); + uint32_t count = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->wqeCntAddr), 0); + const uint32_t tail = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->tailAddr), 0); + if ((count + 10U) % TileXR::TILEXR_UDMA_SQ_BB_COUNT == + tail % TileXR::TILEXR_UDMA_SQ_BB_COUNT) { + const uint32_t target = tail + TileXR::TILEXR_UDMA_NUM_CQE_PER_POLL > count ? + count : tail + TileXR::TILEXR_UDMA_NUM_CQE_PER_POLL; + (void)EpUrmaPollCQ(info, rank, qpIdx, target); + } + + __gm__ uint8_t *wqe = reinterpret_cast<__gm__ uint8_t *>( + sq->bufAddr + wqeSize * (head % TileXR::TILEXR_UDMA_SQ_BB_COUNT)); + EpUrmaFillWriteWqe(wqe, remoteAddr, EpUrmaGetRemoteMemInfo(info, rank, qpIdx), + head, sq->depth, bytes, localAddr); + AscendC::PipeBarrier(); + TileXR::UDMACleanCacheLines(wqe, wqeSize); + AscendC::PipeBarrier(); + ++head; + st_dev(head, reinterpret_cast<__gm__ uint32_t *>(sq->dbAddr), 0); + st_dev(head, reinterpret_cast<__gm__ uint32_t *>(sq->headAddr), 0); + st_dev(++count, reinterpret_cast<__gm__ uint32_t *>(sq->wqeCntAddr), 0); +} + +template +__aicore__ inline void EpUrmaUDMAPutNbi( + const __gm__ TileXR::CommArgs *args, int targetRank, const __gm__ T *localSrc, + uint64_t byteOffset, uint32_t byteCount, uint32_t qpIdx = 0) +{ + if (!TileXR::UDMARegistryEnabled(args)) { + return; + } + __gm__ TileXR::UDMAInfo *info = TileXR::GetUDMAInfo(args); + if (qpIdx >= info->qpNum) { + return; + } + __gm__ TileXR::TileXRUDMARegistry *registry = TileXR::GetUDMARegistry(args); + if (!TileXR::UDMARegisteredRangeValid(registry, targetRank, byteOffset, byteCount)) { + return; + } + EpUrmaPostWrite(info, TileXR::UDMARegisteredRemoteAddr(registry, targetRank, byteOffset), + reinterpret_cast<__gm__ uint8_t *>(const_cast<__gm__ T *>(localSrc)), + static_cast(targetRank), qpIdx, byteCount); +} + +__aicore__ inline uint32_t EpUrmaUDMAQuiet( + const __gm__ TileXR::CommArgs *args, int targetRank, uint32_t qpIdx = 0) +{ + if (!TileXR::UDMAEnabled(args)) { + return 0; + } + __gm__ TileXR::UDMAInfo *info = TileXR::GetUDMAInfo(args); + if (qpIdx >= info->qpNum) { + return 0; + } + __gm__ TileXR::UDMAWQCtx *sq = + TileXR::UDMAGetWQCtx(info, static_cast(targetRank), qpIdx); + const uint32_t count = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->wqeCntAddr), 0); + return EpUrmaPollCQ(info, static_cast(targetRank), qpIdx, count); +} + +} // namespace TileXREp + +#endif // TILEXR_EP_COMMON_EP_URMA_UDMA_H diff --git a/src/ep/host/ep_urma_combine_host.cpp b/src/ep/host/ep_urma_combine_host.cpp new file mode 100644 index 00000000..ec67f6bc --- /dev/null +++ b/src/ep/host/ep_urma_combine_host.cpp @@ -0,0 +1,175 @@ +#include "ep_urma_combine_host.h" + +#include +#include + + +#include "ep_urma_combine.h" +#include "tilexr_perf_trace.h" +#include "tilexr_udma_reg.h" + +extern void launch_tilexr_ep_urma_combine_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, + GM_ADDR expertOut, GM_ADDR assistInfoForCombine, GM_ADDR topKWeights, GM_ADDR yOut, GM_ADDR workspace, + int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, int64_t workspaceBytes, int64_t magic, + int64_t commBytes, int64_t blockCount, int64_t routeStride, int64_t rxWindowBytes, int64_t rxWindowOffset0, + int64_t rxWindowOffset1, int64_t roundDoneOffset0, int64_t roundDoneOffset1, int64_t rxLaneDoneOffset, + int64_t senderDoneOffset, int64_t roundPublishOffset, int64_t roundCreditOffset, + int64_t startGateOffset, int64_t runStartGate, + int64_t errorStatusOffset, int64_t txReadyOffset, int64_t txDataOffset, GM_ADDR perfTrace, int64_t perfTraceBytes, + GM_ADDR strictKernelCycles); + +namespace TileXREp { + + +int TileXREpGetUrmaCombineProfileSize(int64_t rankSize, int64_t *profileBytes) +{ + if (profileBytes == nullptr || rankSize <= 0 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + constexpr int64_t kStatsPerRank = kEpUrmaCombineAivCount * kEpUrmaCombinePerfStageCount; + constexpr int64_t kStatsBytesPerRank = + kStatsPerRank * static_cast(sizeof(TileXR::TileXRPerfCoreStageStats)); + if (rankSize > (std::numeric_limits::max() - + static_cast(TileXR::TILEXR_PERF_TRACE_STATS_OFFSET)) / kStatsBytesPerRank) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *profileBytes = static_cast(TileXR::TILEXR_PERF_TRACE_STATS_OFFSET) + + rankSize * kStatsBytesPerRank; + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpValidateBasicUrmaCombineParams(const EpUrmaCombineParams ¶ms) +{ + if ((params.selfSendCnt > 0 && (params.expertOut == nullptr || params.assistInfoForCombine == nullptr)) || + params.topKWeights == nullptr || params.comm == nullptr || params.yOut == nullptr || + params.workspace == nullptr || params.stream == nullptr || params.selfSendCnt < 0 || params.bs <= 0 || + params.h <= 0 || params.h > kEpUrmaCombineMaxHidden || params.topK <= 0 || + params.topK > kEpUrmaCombineMaxTopK || params.workspaceBytes <= 0 || params.perfTraceBytes < 0 || + ((params.perfTrace == nullptr) != (params.perfTraceBytes == 0)) || + params.strictKernelCyclesBytes < 0 || + ((params.strictKernelCycles == nullptr) != (params.strictKernelCyclesBytes == 0)) || + params.dtype != TileXR::TILEXR_DATA_TYPE_FP16) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if ((reinterpret_cast(params.workspace) % kEpUrmaCombineWorkspaceAlignment) != 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (params.perfTrace != nullptr && (reinterpret_cast(params.perfTrace) % 32) != 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + constexpr int64_t kStrictKernelCyclesBytes = + kEpUrmaCombineAivCount * static_cast(sizeof(uint64_t)); + if (params.strictKernelCycles != nullptr && (params.perfTrace != nullptr || + params.strictKernelCyclesBytes < kStrictKernelCyclesBytes || + (reinterpret_cast(params.strictKernelCycles) % 32) != 0)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpPrepareUrmaCombineLaunchContext( + const EpUrmaCombineParams ¶ms, EpUrmaCombineLaunchContext *context) +{ + if (context == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *context = EpUrmaCombineLaunchContext {}; + + int ret = TileXREpValidateBasicUrmaCombineParams(params); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + ret = TileXRGetCommArgsHost(params.comm, context->hostArgs); + if (ret != TileXR::TILEXR_SUCCESS || context->hostArgs == nullptr) { + *context = EpUrmaCombineLaunchContext {}; + return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + } + ret = TileXRGetCommArgsDev(params.comm, context->devArgs); + if (ret != TileXR::TILEXR_SUCCESS || context->devArgs == nullptr) { + *context = EpUrmaCombineLaunchContext {}; + return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + } + + const TileXR::CommArgs &args = *context->hostArgs; + if (args.rankSize <= 0 || args.rankSize > TileXR::TILEXR_MAX_RANK_SIZE || args.rank < 0 || + args.rank >= args.rankSize || (args.rankSize > 1 && + ((args.extraFlag & TileXR::ExtraFlag::UDMA) == 0 || args.udmaInfoPtr == nullptr || + args.udmaRegistryPtr == nullptr))) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + + if (params.perfTrace != nullptr) { + int64_t requiredProfileBytes = 0; + ret = TileXREpGetUrmaCombineProfileSize(args.rankSize, &requiredProfileBytes); + if (ret != TileXR::TILEXR_SUCCESS || params.perfTraceBytes < requiredProfileBytes) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + + ret = TileXREpBuildUrmaCombineWorkspaceConfig( + args.rankSize, params.bs, params.h, params.topK, params.selfSendCnt, &context->workspace); + if (ret != TileXR::TILEXR_SUCCESS || context->workspace.requiredBytes > params.workspaceBytes) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + if (args.rankSize > 1) { + const TileXR::TileXRUDMARegistry *registry = nullptr; + ret = TileXRGetUDMARegistryHost(params.comm, ®istry); + if (ret != TileXR::TILEXR_SUCCESS || !TileXR::UDMARegistryValid(registry, args.rankSize)) { + *context = EpUrmaCombineLaunchContext {}; + return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + } + if (registry->regions[args.rank].base != static_cast(params.workspace) || + !TileXR::UDMARegionContains(registry, args.rank, 0, + static_cast(context->workspace.requiredBytes))) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + for (int32_t peer = 0; peer < args.rankSize; ++peer) { + if (!TileXR::UDMARegionContains( + registry, peer, 0, static_cast(context->workspace.fixedBytes))) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + } + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpLaunchPreparedUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context, int64_t magic, + bool runStartGate) +{ + launch_tilexr_ep_urma_combine_kernel(static_cast(kEpUrmaCombineAivCount), + params.stream, context.devArgs, + static_cast(params.expertOut), reinterpret_cast(params.assistInfoForCombine), + reinterpret_cast(params.topKWeights), static_cast(params.yOut), + static_cast(params.workspace), params.selfSendCnt, params.bs, params.h, params.topK, + params.workspaceBytes, magic, context.workspace.commBytes, context.workspace.blockCount, + context.workspace.routeStride, context.workspace.rxWindowBytes, context.workspace.rxWindowOffsets[0], + context.workspace.rxWindowOffsets[1], context.workspace.roundDoneOffsets[0], + context.workspace.roundDoneOffsets[1], context.workspace.rxLaneDoneOffset, + context.workspace.senderDoneOffset, context.workspace.roundPublishOffset, + context.workspace.roundCreditOffset, + context.workspace.startGateOffset, runStartGate ? 1 : 0, context.workspace.errorStatusOffset, + context.workspace.txReadyOffset, context.workspace.txDataOffset, + static_cast(params.perfTrace), params.perfTraceBytes, + static_cast(params.strictKernelCycles)); + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpLaunchUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context) +{ + int64_t magic = 0; + const int ret = TileXRCommNextMagic(params.comm, &magic); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + return TileXREpLaunchPreparedUrmaCombineKernel(params, context, magic, true); +} + +} // namespace TileXREp diff --git a/src/ep/host/ep_urma_combine_host.h b/src/ep/host/ep_urma_combine_host.h new file mode 100644 index 00000000..6e7b531b --- /dev/null +++ b/src/ep/host/ep_urma_combine_host.h @@ -0,0 +1,51 @@ +#ifndef TILEXR_EP_HOST_EP_URMA_COMBINE_HOST_H +#define TILEXR_EP_HOST_EP_URMA_COMBINE_HOST_H + +#include + +#include "acl/acl_base.h" +#include "ep_urma_combine_layout.h" +#include "tilexr_api.h" +#include "tilexr_types.h" + +namespace TileXREp { + +struct EpUrmaCombineParams { + void *expertOut = nullptr; + int32_t *assistInfoForCombine = nullptr; + float *topKWeights = nullptr; + TileXRCommPtr comm = nullptr; + int64_t selfSendCnt = 0; + int64_t bs = 0; + int64_t h = 0; + int64_t topK = 0; + void *yOut = nullptr; + void *workspace = nullptr; + int64_t workspaceBytes = 0; + void *perfTrace = nullptr; + int64_t perfTraceBytes = 0; + void *strictKernelCycles = nullptr; + int64_t strictKernelCyclesBytes = 0; + TileXR::TileXRDataType dtype = TileXR::TILEXR_DATA_TYPE_RESERVED; + aclrtStream stream = nullptr; +}; + +struct EpUrmaCombineLaunchContext { + TileXR::CommArgs *hostArgs = nullptr; + GM_ADDR devArgs = nullptr; + EpUrmaCombineWorkspaceConfig workspace {}; +}; + +int TileXREpGetUrmaCombineProfileSize(int64_t rankSize, int64_t *profileBytes); +int TileXREpValidateBasicUrmaCombineParams(const EpUrmaCombineParams ¶ms); +int TileXREpPrepareUrmaCombineLaunchContext( + const EpUrmaCombineParams ¶ms, EpUrmaCombineLaunchContext *context); +int TileXREpLaunchPreparedUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context, int64_t magic, + bool runStartGate = true); +int TileXREpLaunchUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context); + +} // namespace TileXREp + +#endif // TILEXR_EP_HOST_EP_URMA_COMBINE_HOST_H diff --git a/src/ep/host/ep_urma_combine_layout.cpp b/src/ep/host/ep_urma_combine_layout.cpp new file mode 100644 index 00000000..f3a2fe14 --- /dev/null +++ b/src/ep/host/ep_urma_combine_layout.cpp @@ -0,0 +1,117 @@ +#include "ep_urma_combine_layout.h" + +#include + +#include "comm_args.h" +#include "ep_urma_combine.h" +#include "tilexr_types.h" + +namespace TileXREp { +namespace { + +bool AddInt64(int64_t lhs, int64_t rhs, int64_t *out) +{ + if (out == nullptr || lhs < 0 || rhs < 0 || rhs > std::numeric_limits::max() - lhs) { + return false; + } + *out = lhs + rhs; + return true; +} + +bool MulInt64(int64_t lhs, int64_t rhs, int64_t *out) +{ + if (out == nullptr || lhs < 0 || rhs < 0 || + (lhs != 0 && rhs > std::numeric_limits::max() / lhs)) { + return false; + } + *out = lhs * rhs; + return true; +} + +bool AlignUpInt64(int64_t value, int64_t alignment, int64_t *out) +{ + if (out == nullptr || value < 0 || alignment <= 0) { + return false; + } + const int64_t remainder = value % alignment; + return remainder == 0 ? (*out = value, true) : AddInt64(value, alignment - remainder, out); +} + +bool AddAlignedRegion(int64_t *offset, int64_t bytes, int64_t alignment, int64_t *regionOffset) +{ + int64_t aligned = 0; + if (offset == nullptr || regionOffset == nullptr || bytes < 0 || + !AlignUpInt64(*offset, alignment, &aligned) || !AddInt64(aligned, bytes, offset)) { + return false; + } + *regionOffset = aligned; + return true; +} + +} // namespace + +int TileXREpBuildUrmaCombineWorkspaceConfig(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCnt, EpUrmaCombineWorkspaceConfig *out) +{ + if (out == nullptr || rankSize <= 0 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || bs <= 0 || h <= 0 || + h > kEpUrmaCombineMaxHidden || topK <= 0 || topK > kEpUrmaCombineMaxTopK || selfSendCnt < 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + EpUrmaCombineWorkspaceConfig next {}; + next.rankSize = rankSize; + next.bs = bs; + next.h = h; + next.topK = topK; + next.selfSendCnt = selfSendCnt; + + if (!AlignUpInt64(h, kEpUrmaCombineQuantHeaderBytes, &next.quantDataBytes) || + !AddInt64(kEpUrmaCombineQuantHeaderBytes, next.quantDataBytes, &next.commBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.blockCount = (next.commBytes + kEpUrmaCombinePayloadBytes - 1) / kEpUrmaCombinePayloadBytes; + if (next.blockCount <= 0 || next.blockCount > kEpUrmaCombineMaxBlocksPerRoute || + !MulInt64(next.blockCount, kEpUrmaCombineDataBlockBytes, &next.routeStride) || + !MulInt64(bs, topK, &next.routeCount) || + !MulInt64(next.routeCount, next.routeStride, &next.rxWindowBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t offset = kEpUrmaCombineHeaderBytes; + int64_t roundDoneBytes = 0; + int64_t rxLaneDoneBytes = 0; + int64_t senderDoneBytes = 0; + int64_t txReadyBytes = 0; + int64_t txDataBytes = 0; + if (!MulInt64(rankSize, kEpUrmaCombineCacheLineBytes, &roundDoneBytes) || + !MulInt64(kEpUrmaCombinePackLaneCount, kEpUrmaCombineCacheLineBytes, &rxLaneDoneBytes) || + !MulInt64(kEpUrmaCombineSendLaneCount, kEpUrmaCombineCacheLineBytes, &senderDoneBytes) || + !MulInt64(selfSendCnt, kEpUrmaCombineCacheLineBytes, &txReadyBytes) || + !MulInt64(selfSendCnt, next.routeStride, &txDataBytes) || + !AddAlignedRegion(&offset, next.rxWindowBytes, kEpUrmaCombineDataBlockBytes, &next.rxWindowOffsets[0]) || + !AddAlignedRegion(&offset, next.rxWindowBytes, kEpUrmaCombineDataBlockBytes, &next.rxWindowOffsets[1]) || + !AddAlignedRegion(&offset, roundDoneBytes, kEpUrmaCombineCacheLineBytes, &next.roundDoneOffsets[0]) || + !AddAlignedRegion(&offset, roundDoneBytes, kEpUrmaCombineCacheLineBytes, &next.roundDoneOffsets[1]) || + !AddAlignedRegion(&offset, rxLaneDoneBytes, kEpUrmaCombineCacheLineBytes, &next.rxLaneDoneOffset) || + !AddAlignedRegion(&offset, senderDoneBytes, kEpUrmaCombineCacheLineBytes, &next.senderDoneOffset) || + !AddAlignedRegion(&offset, kEpUrmaCombineCacheLineBytes, kEpUrmaCombineCacheLineBytes, + &next.roundPublishOffset) || + !AddAlignedRegion(&offset, kEpUrmaCombineCacheLineBytes, kEpUrmaCombineCacheLineBytes, + &next.roundCreditOffset) || + !AddAlignedRegion(&offset, roundDoneBytes, kEpUrmaCombineCacheLineBytes, &next.startGateOffset) || + !AddAlignedRegion(&offset, kEpUrmaCombineCacheLineBytes, kEpUrmaCombineCacheLineBytes, + &next.errorStatusOffset)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + next.fixedBytes = offset; + if (!AddAlignedRegion(&offset, txReadyBytes, kEpUrmaCombineCacheLineBytes, &next.txReadyOffset) || + !AddAlignedRegion(&offset, txDataBytes, kEpUrmaCombineDataBlockBytes, &next.txDataOffset)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.requiredBytes = offset; + *out = next; + return TileXR::TILEXR_SUCCESS; +} + +} // namespace TileXREp diff --git a/src/ep/host/ep_urma_combine_layout.h b/src/ep/host/ep_urma_combine_layout.h new file mode 100644 index 00000000..fc71bd85 --- /dev/null +++ b/src/ep/host/ep_urma_combine_layout.h @@ -0,0 +1,39 @@ +#ifndef TILEXR_EP_HOST_EP_URMA_COMBINE_LAYOUT_H +#define TILEXR_EP_HOST_EP_URMA_COMBINE_LAYOUT_H + +#include + +namespace TileXREp { + +struct EpUrmaCombineWorkspaceConfig { + int64_t rankSize = 0; + int64_t bs = 0; + int64_t h = 0; + int64_t topK = 0; + int64_t selfSendCnt = 0; + int64_t quantDataBytes = 0; + int64_t commBytes = 0; + int64_t blockCount = 0; + int64_t routeStride = 0; + int64_t routeCount = 0; + int64_t rxWindowBytes = 0; + int64_t rxWindowOffsets[2] = {0, 0}; + int64_t roundDoneOffsets[2] = {0, 0}; + int64_t rxLaneDoneOffset = 0; + int64_t senderDoneOffset = 0; + int64_t roundPublishOffset = 0; + int64_t roundCreditOffset = 0; + int64_t startGateOffset = 0; + int64_t errorStatusOffset = 0; + int64_t fixedBytes = 0; + int64_t txReadyOffset = 0; + int64_t txDataOffset = 0; + int64_t requiredBytes = 0; +}; + +int TileXREpBuildUrmaCombineWorkspaceConfig(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCnt, EpUrmaCombineWorkspaceConfig *out); + +} // namespace TileXREp + +#endif // TILEXR_EP_HOST_EP_URMA_COMBINE_LAYOUT_H diff --git a/src/ep/host/ep_urma_combine_start_gate_window.h b/src/ep/host/ep_urma_combine_start_gate_window.h new file mode 100644 index 00000000..78c9d907 --- /dev/null +++ b/src/ep/host/ep_urma_combine_start_gate_window.h @@ -0,0 +1,26 @@ +#ifndef TILEXR_EP_HOST_EP_URMA_COMBINE_START_GATE_WINDOW_H +#define TILEXR_EP_HOST_EP_URMA_COMBINE_START_GATE_WINDOW_H + +namespace TileXREp { + +class EpUrmaCombineStartGateWindow { +public: + bool BeginLaunch() + { + const bool runStartGate = firstLaunchAfterSynchronization_; + firstLaunchAfterSynchronization_ = false; + return runStartGate; + } + + void StreamSynchronized() + { + firstLaunchAfterSynchronization_ = true; + } + +private: + bool firstLaunchAfterSynchronization_ = true; +}; + +} // namespace TileXREp + +#endif // TILEXR_EP_HOST_EP_URMA_COMBINE_START_GATE_WINDOW_H diff --git a/src/ep/host/tilexr_ep_urma_combine.cpp b/src/ep/host/tilexr_ep_urma_combine.cpp new file mode 100644 index 00000000..f7df518f --- /dev/null +++ b/src/ep/host/tilexr_ep_urma_combine.cpp @@ -0,0 +1,110 @@ +#include "tilexr_ep.h" + +#include + +#include "ep_urma_combine.h" +#include "ep_urma_combine_host.h" + +int TileXRMoeEpCombineUrmaGetWorkspaceSize(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCapacity, int64_t *workspaceBytes) +{ + if (workspaceBytes == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + TileXREp::EpUrmaCombineWorkspaceConfig layout {}; + const int ret = TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig( + rankSize, bs, h, topK, selfSendCapacity, &layout); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + const int64_t alignment = TileXREp::kEpUrmaCombineWorkspaceAlignment; + const int64_t remainder = layout.requiredBytes % alignment; + if (remainder != 0 && layout.requiredBytes > + std::numeric_limits::max() - (alignment - remainder)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *workspaceBytes = remainder == 0 ? layout.requiredBytes : layout.requiredBytes + alignment - remainder; + return TileXR::TILEXR_SUCCESS; +} + +int TileXRMoeEpCombineUrmaGetProfileSize(int64_t rankSize, int64_t *profileBytes) +{ + return TileXREp::TileXREpGetUrmaCombineProfileSize(rankSize, profileBytes); +} + +namespace { + +int CombineUrmaImpl(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, void *perfTrace, int64_t perfTraceBytes, + TileXR::TileXRDataType dtype, aclrtStream stream) +{ + TileXREp::EpUrmaCombineParams params {}; + params.expertOut = expertOut; + params.assistInfoForCombine = assistInfoForCombine; + params.topKWeights = topKWeights; + params.comm = comm; + params.selfSendCnt = selfSendCnt; + params.bs = bs; + params.h = h; + params.topK = topK; + params.yOut = yOut; + params.workspace = workspace; + params.workspaceBytes = workspaceBytes; + params.perfTrace = perfTrace; + params.perfTraceBytes = perfTraceBytes; + params.dtype = dtype; + params.stream = stream; + + int ret = TileXREp::TileXREpValidateBasicUrmaCombineParams(params); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + TileXREp::EpUrmaCombineLaunchContext context {}; + ret = TileXREp::TileXREpPrepareUrmaCombineLaunchContext(params, &context); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + return TileXREp::TileXREpLaunchUrmaCombineKernel(params, context); +} + +} // namespace + +int TileXRMoeEpCombineUrma(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, TileXR::TileXRDataType dtype, aclrtStream stream) +{ + return CombineUrmaImpl(expertOut, assistInfoForCombine, topKWeights, comm, selfSendCnt, bs, h, topK, + yOut, workspace, workspaceBytes, nullptr, 0, dtype, stream); +} + +int TileXRMoeEpCombineUrmaProfile(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, void *perfTrace, int64_t perfTraceBytes, + TileXR::TileXRDataType dtype, aclrtStream stream) +{ +#if !defined(TILEXR_EP_ENABLE_PROFILING) + (void)expertOut; + (void)assistInfoForCombine; + (void)topKWeights; + (void)comm; + (void)selfSendCnt; + (void)bs; + (void)h; + (void)topK; + (void)yOut; + (void)workspace; + (void)workspaceBytes; + (void)perfTrace; + (void)perfTraceBytes; + (void)dtype; + (void)stream; + return TileXR::TILEXR_ERROR_NOT_SUPPORT; +#else + if (perfTrace == nullptr || perfTraceBytes <= 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + return CombineUrmaImpl(expertOut, assistInfoForCombine, topKWeights, comm, selfSendCnt, bs, h, topK, + yOut, workspace, workspaceBytes, perfTrace, perfTraceBytes, dtype, stream); +#endif +} diff --git a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp new file mode 100644 index 00000000..2f15a6c9 --- /dev/null +++ b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp @@ -0,0 +1,1486 @@ +#include "comm_args.h" +#include "ep_urma_combine.h" +#include "ep_urma_udma.h" +#include "ep_window.h" +#include "kernel_operator.h" +#include "tilexr_data_as_flag.h" +#include "ep_urma_perf_trace.h" +#include "tilexr_udma.h" + +namespace { + +constexpr uint32_t kCursorAlignment = 32; +constexpr int64_t kPipelineBufferCount = 2; +using PackInputQueue = AscendC::TQue; +constexpr int64_t kRxTokenScheduleWindow = 3; + +using PerfStats = TileXR::TileXRPerfCoreStageStats; +using PerfStage = TileXREp::EpUrmaCombinePerfStage; + +static_assert(TileXREp::kEpUrmaCombinePerfStageCount <= + TileXR::TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT, + "URMA combine profile stages must fit the reserved UB trace region"); +static_assert(TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET + + TileXREp::kEpUrmaCombinePerfStageCount * sizeof(PerfStats) <= + TileXR::TILEXR_PERF_TRACE_MIN_UB_BYTES, + "URMA combine profile stats exceed the minimum supported AIV UB"); + +__aicore__ inline uint32_t PerfStageId(PerfStage stage) +{ + return static_cast(stage); +} + +__aicore__ inline uint64_t ProfileBegin(GM_ADDR perfTrace) +{ + return TileXR::TileXRPerfCycleNow(perfTrace); +} + +__aicore__ inline void ProfileEnd( + GM_ADDR perfTrace, __ubuf__ PerfStats *stats, PerfStage stage, uint64_t startCycle) +{ + TileXR::TileXRPerfLocalRecord(perfTrace, stats, TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(stage), startCycle, TileXR::TileXRPerfCycleNow(perfTrace)); +} + +__aicore__ inline void ProfileAux( + GM_ADDR perfTrace, __ubuf__ PerfStats *stats, PerfStage stage, uint32_t auxIndex, uint64_t value) +{ + TileXR::TileXRPerfLocalAddAux(perfTrace, stats, TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(stage), auxIndex, value); +} + +__aicore__ inline uint64_t ProfileKernelTimingBegin(GM_ADDR perfTrace) +{ + if (TileXR::TileXRPerfTraceEnabled(perfTrace)) { + AscendC::PipeBarrier(); + return static_cast(AscendC::GetSystemCycle()); + } + return 0; +} + +__aicore__ inline void ProfileFinish( + GM_ADDR perfTrace, uint32_t rank, uint32_t core, + __ubuf__ PerfStats *stats, uint64_t kernelStart) +{ + if (TileXR::TileXRPerfTraceEnabled(perfTrace)) { + AscendC::PipeBarrier(); + const uint64_t kernelEnd = static_cast(AscendC::GetSystemCycle()); + TileXR::TileXRPerfLocalRecord(perfTrace, stats, + TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(PerfStage::KERNEL_TOTAL), kernelStart, kernelEnd); + } + TileXR::TileXRPerfLocalStatsFlush(perfTrace, rank, core, + static_cast(TileXREp::kEpUrmaCombineAivCount), + TileXREp::kEpUrmaCombinePerfStageCount, stats); +} + +__aicore__ inline uint64_t StrictKernelTimingBegin( + GM_ADDR strictKernelCycles) +{ + if (strictKernelCycles == nullptr) { + return 0; + } + AscendC::PipeBarrier(); + return static_cast(AscendC::GetSystemCycle()); +} + +__aicore__ inline void StrictKernelTimingFinish( + GM_ADDR strictKernelCycles, uint32_t core, uint64_t startCycle) +{ + if (strictKernelCycles == nullptr) { + return; + } + AscendC::PipeBarrier(); + const uint64_t endCycle = static_cast(AscendC::GetSystemCycle()); + const uint64_t duration = endCycle >= startCycle ? endCycle - startCycle : 0; + AscendC::WriteGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(strictKernelCycles) + core, duration); + AscendC::DataSyncBarrier(); +} + +__aicore__ inline bool ProfileBufferValid( + GM_ADDR perfTrace, int64_t perfTraceBytes, int32_t rank, int32_t rankSize) +{ + if (perfTrace == nullptr || perfTraceBytes < + static_cast(sizeof(TileXR::TileXRPerfTraceHeader)) || rank < 0 || rankSize <= 0) { + return false; + } + const uint64_t requiredStatsBytes = static_cast(rankSize) * + TileXREp::kEpUrmaCombineAivCount * TileXREp::kEpUrmaCombinePerfStageCount * sizeof(PerfStats); + const uint64_t requiredBytes = TileXR::TILEXR_PERF_TRACE_STATS_OFFSET + requiredStatsBytes; + if (requiredBytes > static_cast(perfTraceBytes)) { + return false; + } + const auto header = reinterpret_cast<__gm__ TileXR::TileXRPerfTraceHeader *>(perfTrace); + return header->magic == TileXR::TILEXR_PERF_TRACE_MAGIC && + header->version == TileXR::TILEXR_PERF_TRACE_VERSION && + header->headerSize == sizeof(TileXR::TileXRPerfTraceHeader) && + header->coreStageStatsSize == sizeof(PerfStats) && + header->rank == static_cast(rank) && + header->rankSize == static_cast(rankSize) && + header->blockDim == TileXREp::kEpUrmaCombineAivCount && + header->maxCoreCount == TileXREp::kEpUrmaCombineAivCount && + header->stageCount == TileXREp::kEpUrmaCombinePerfStageCount && + header->flags <= 2 && header->cycleToUsDivisor != 0 && + header->statsOffset == TileXR::TILEXR_PERF_TRACE_STATS_OFFSET && + header->statsBytes == requiredStatsBytes; +} + +static_assert(TileXREp::kEpUrmaCombineDataBlockBytes == TileXR::DATA_AS_FLAG_BLOCK_BYTES, + "URMA combine and DataAsFlag block sizes must match"); +static_assert(TileXREp::kEpUrmaCombinePayloadBytes == TileXR::DATA_AS_FLAG_PAYLOAD_BYTES, + "URMA combine and DataAsFlag payload sizes must match"); +static_assert(TileXREp::kEpUrmaCombineFlagBytes == TileXR::DATA_AS_FLAG_FLAG_BYTES, + "URMA combine and DataAsFlag flag sizes must match"); +constexpr int64_t kMaxInputBytes = TileXREp::kEpUrmaCombineMaxHidden * sizeof(half); +constexpr int64_t kMaxFloatBytes = TileXREp::kEpUrmaCombineMaxHidden * sizeof(float); +constexpr float kMaxFiniteHalf = 65504.0f; +constexpr int64_t kMaxLogicalBytes = + ((TileXREp::kEpUrmaCombineQuantHeaderBytes + TileXREp::kEpUrmaCombineMaxHidden + + TileXREp::kEpUrmaCombinePayloadBytes - 1) / + TileXREp::kEpUrmaCombinePayloadBytes) * TileXREp::kEpUrmaCombinePayloadBytes; +constexpr int64_t kMaxBlockCount = kMaxLogicalBytes / TileXREp::kEpUrmaCombinePayloadBytes; +constexpr int64_t kMaxRouteStride = kMaxBlockCount * TileXREp::kEpUrmaCombineDataBlockBytes; +constexpr int64_t kMaxPackUbBytes = kPipelineBufferCount * kMaxInputBytes + + kPipelineBufferCount * kMaxRouteStride + + kMaxFloatBytes + 2 * kMaxInputBytes + kCursorAlignment + kMaxLogicalBytes; +constexpr int64_t kMaxRouteFlagBytes = + kMaxBlockCount * TileXREp::kEpUrmaCombineFlagBytes; +constexpr int64_t kMaxReadyFlagBufferBytes = kMaxRouteFlagBytes; +constexpr int64_t kMaxReceiveUbBytes = kPipelineBufferCount * kMaxLogicalBytes + + kMaxReadyFlagBufferBytes + kMaxRouteFlagBytes + + 2 * kMaxFloatBytes + kMaxInputBytes + + TileXREp::kEpUrmaCombineMaxTopK * static_cast(sizeof(float)); +static_assert(kMaxPackUbBytes <= TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET, + "URMA combine Pack ping-pong buffers overlap the profiling UB region"); +static_assert(kMaxReceiveUbBytes <= TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET, + "URMA combine Receive ping-pong buffers overlap the profiling UB region"); + +__aicore__ inline int64_t AlignUpInt64(int64_t value, int64_t alignment) +{ + const int64_t remainder = value % alignment; + return remainder == 0 ? value : value + alignment - remainder; +} + +template +__aicore__ inline void SetNoCacheRead(AscendC::GlobalTensor &tensor) +{ + tensor.template SetL2CacheHint( + AscendC::CacheMode::CACHE_MODE_DISABLE); +} + +template +__aicore__ inline void SetNoCacheWrite(AscendC::GlobalTensor &tensor) +{ + tensor.template SetL2CacheHint( + AscendC::CacheMode::CACHE_MODE_DISABLE); +} + +__aicore__ inline void CachelessAcquireBarrier() +{ + AscendC::DataSyncBarrier(); + AscendC::PipeBarrier(); +} + +__aicore__ inline uint64_t EncodeControlValue(int64_t magic, uint32_t step) +{ + return (static_cast(static_cast(magic)) << 32) | static_cast(step); +} + +__aicore__ inline void StoreControlValue(GM_ADDR lineAddr, uint64_t value) +{ + AscendC::PipeBarrier(); + AscendC::WriteGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(lineAddr), value); + AscendC::DataSyncBarrier(); +} + + +__aicore__ inline uint64_t LoadControlValue(GM_ADDR lineAddr) +{ + const uint64_t value = AscendC::ReadGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(lineAddr)); + return value; +} + +__aicore__ inline void StoreError( + GM_ADDR workspaceGM, int64_t errorStatusOffset, int64_t magic, uint64_t status) +{ + if (status != TileXREp::kEpUrmaCombineStatusOk) { + StoreControlValue(workspaceGM + errorStatusOffset, + EncodeControlValue(magic, static_cast(status))); + } +} + +__aicore__ inline TileXREp::EpAssistTuple LoadRouteMetaBypass( + GM_ADDR assistInfoGM, int64_t index) +{ + __gm__ int32_t *src = reinterpret_cast<__gm__ int32_t *>(assistInfoGM) + + index * TileXREp::kEpAssistTupleInts; + TileXREp::EpAssistTuple tuple; + tuple.srcRank = AscendC::ReadGmByPassDCache(src); + tuple.tokenId = AscendC::ReadGmByPassDCache(src + 1); + tuple.topKId = AscendC::ReadGmByPassDCache(src + 2); + tuple.expertId = 0; + return tuple; +} + +__aicore__ inline TileXREp::EpAssistTuple LoadRouteMeta( + GM_ADDR assistInfoGM, int64_t index) +{ + return LoadRouteMetaBypass(assistInfoGM, index); +} + + +__aicore__ inline bool RouteMetaValid( + const TileXREp::EpAssistTuple &tuple, int32_t rankSize, int64_t bs, int64_t topK) +{ + return tuple.srcRank >= 0 && tuple.srcRank < rankSize && tuple.tokenId >= 0 && tuple.tokenId < bs && + tuple.topKId >= 0 && tuple.topKId < topK; +} + +__aicore__ inline void LoadTopKWeights( + GM_ADDR weightsGM, int64_t begin, int64_t topK, AscendC::LocalTensor &local) +{ + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(weightsGM) + begin, topK); + const uint32_t alignedElements = static_cast( + AlignUpInt64(topK * static_cast(sizeof(float)), kCursorAlignment) / sizeof(float)); + AscendC::DataCopyExtParams copyParams { + 1, static_cast(topK * static_cast(sizeof(float))), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams { + true, 0, static_cast(alignedElements - static_cast(topK)), 0.0f}; + AscendC::DataCopyPad(local, src, copyParams, padParams); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + +__aicore__ inline event_t PipelineEvent(int64_t slot) +{ + return slot == 0 ? EVENT_ID0 : EVENT_ID1; +} + +__aicore__ inline void EnqueuePackInput( + const AscendC::GlobalTensor &expertOut, int64_t route, + int64_t h, uint8_t inputPadding, PackInputQueue &inputQueue) +{ + AscendC::LocalTensor input = inputQueue.AllocTensor(); + AscendC::DataCopyExtParams inputParams { + 1, static_cast(h * static_cast(sizeof(half))), 0, 0, 0}; + AscendC::DataCopyPadExtParams inputPad {true, 0, inputPadding, static_cast(0.0f)}; + AscendC::DataCopyPad(input, expertOut[route * h], inputParams, inputPad); + inputQueue.EnQue(input); +} + +__aicore__ inline void StartPackPublish( + GM_ADDR txRouteAddr, int64_t routeStride, + AscendC::LocalTensor &packed, event_t eventId) +{ + AscendC::GlobalTensor txRoute; + txRoute.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(txRouteAddr), routeStride / sizeof(float)); + SetNoCacheWrite(txRoute); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(txRoute, packed, static_cast(routeStride / sizeof(float))); + AscendC::SetFlag(eventId); +} + +__aicore__ inline void FinishPackCopy( + GM_ADDR txRouteAddr, int64_t routeStride, event_t eventId, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t waitStart = ProfileBegin(perfTrace); + AscendC::WaitFlag(eventId); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_MTE3_EXPOSED_WAIT, waitStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_MTE3_EXPOSED_WAIT, 0, + static_cast(routeStride)); +} + +__aicore__ inline void PublishTxReadyBatch( + GM_ADDR workspaceGM, int64_t txReadyOffset, + int64_t firstRoute, int64_t routeCount, uint64_t readyValue) +{ + if (routeCount <= 0) { + return; + } + GM_ADDR firstLineAddr = workspaceGM + txReadyOffset + + firstRoute * TileXREp::kEpUrmaCombineCacheLineBytes; + for (int64_t index = 0; index < routeCount; ++index) { + AscendC::WriteGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(firstLineAddr + + index * TileXREp::kEpUrmaCombineCacheLineBytes), + readyValue); + } + AscendC::DataSyncBarrier(); +} + +__aicore__ __attribute__((always_inline)) inline void StartPackQuantization( + AscendC::LocalTensor &routeInput, + AscendC::LocalTensor &routePacked, + AscendC::LocalTensor &quantFloat, + AscendC::LocalTensor &absHalf, + AscendC::LocalTensor &reduceOut, + AscendC::LocalTensor &reduceTmp, + AscendC::LocalTensor &logical, + int64_t h, int64_t logicalBytes, int64_t routeStride, AscendC::TEventID quantizeEvent) +{ + (void)routePacked; + (void)logical; + (void)logicalBytes; + (void)routeStride; + AscendC::Abs(absHalf, routeInput, static_cast(h)); + AscendC::PipeBarrier(); + AscendC::ReduceMax(reduceOut, absHalf, reduceTmp, static_cast(h), false); + (void)quantFloat; + AscendC::SetFlag(quantizeEvent); +} + +__aicore__ __attribute__((always_inline)) inline void FinishPackQuantization( + AscendC::LocalTensor &routeInput, + AscendC::LocalTensor &routePacked, + AscendC::LocalTensor &quantFloat, + AscendC::LocalTensor &reduceOut, + AscendC::LocalTensor &logical, + AscendC::LocalTensor &logicalFloat, + AscendC::LocalTensor &logicalInt, + int64_t h, int64_t blockCount, uint8_t repeats, + AscendC::TEventID quantizeEvent, AscendC::TEventID scalarToVectorEvent) +{ + AscendC::WaitFlag(quantizeEvent); + const float maxAbs = static_cast(reduceOut(0)); + const float scale = maxAbs > 0.0f ? maxAbs / 127.0f : 1.0f; + const float inverseScale = maxAbs > 0.0f ? 127.0f / maxAbs : 1.0f; + logicalFloat.SetValue(0, scale); + logicalInt.SetValue(1, static_cast(TileXREp::kEpUrmaCombineQuantModeInt8PerRoute)); + const bool halfScaleSafe = maxAbs >= 0.0f && maxAbs <= kMaxFiniteHalf && + inverseScale > 0.0f && inverseScale <= kMaxFiniteHalf; + if (halfScaleSafe) { + AscendC::Muls(routeInput, routeInput, static_cast(inverseScale), + static_cast(h)); + AscendC::PipeBarrier(); + } else { + AscendC::Cast(quantFloat, routeInput, AscendC::RoundMode::CAST_NONE, + static_cast(h)); + AscendC::PipeBarrier(); + AscendC::Muls(quantFloat, quantFloat, inverseScale, static_cast(h)); + AscendC::PipeBarrier(); + AscendC::Cast(routeInput, quantFloat, AscendC::RoundMode::CAST_RINT, + static_cast(h)); + AscendC::PipeBarrier(); + } + AscendC::Cast(logical[TileXREp::kEpUrmaCombineQuantHeaderBytes], routeInput, + AscendC::RoundMode::CAST_RINT, static_cast(h)); + AscendC::PipeBarrier(); + + AscendC::SetFlag(scalarToVectorEvent); + AscendC::WaitFlag(scalarToVectorEvent); + AscendC::Copy(routePacked, logicalFloat, static_cast(64), repeats, {1, 1, 16, 15}); + AscendC::Copy(routePacked[64], logicalFloat[64], static_cast(56), repeats, {1, 1, 16, 15}); + AscendC::PipeBarrier(); +} + +__aicore__ inline void PackRoutes( + GM_ADDR expertOutGM, GM_ADDR workspaceGM, int64_t selfSendCnt, int64_t h, + int64_t magic, int64_t blockCount, int64_t routeStride, int64_t txReadyOffset, int64_t txDataOffset, + int64_t laneId, AscendC::TPipe &pipe, GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t firstReadyStart = ProfileBegin(perfTrace); + const int64_t inputBytes = AlignUpInt64(h * static_cast(sizeof(half)), kCursorAlignment); + const int64_t floatBytes = AlignUpInt64(h * static_cast(sizeof(float)), kCursorAlignment); + const int64_t logicalBytes = blockCount * TileXREp::kEpUrmaCombinePayloadBytes; + + PackInputQueue inputQueue; + AscendC::TBuf quantFloatBuf; + AscendC::TBuf absHalfBuf; + AscendC::TBuf reduceOutBuf; + AscendC::TBuf reduceTmpBuf; + AscendC::TBuf logicalBuf; + AscendC::TBuf packedBuf; + pipe.InitBuffer(inputQueue, static_cast(kPipelineBufferCount), static_cast(inputBytes)); + pipe.InitBuffer(quantFloatBuf, static_cast(floatBytes)); + pipe.InitBuffer(absHalfBuf, static_cast(inputBytes)); + pipe.InitBuffer(reduceOutBuf, kCursorAlignment); + pipe.InitBuffer(reduceTmpBuf, static_cast(inputBytes)); + pipe.InitBuffer(logicalBuf, static_cast(logicalBytes)); + pipe.InitBuffer(packedBuf, static_cast(kPipelineBufferCount * routeStride)); + AscendC::LocalTensor packed0 = packedBuf.GetWithOffset( + static_cast(routeStride / sizeof(float)), 0); + AscendC::LocalTensor packed1 = packedBuf.GetWithOffset( + static_cast(routeStride / sizeof(float)), static_cast(routeStride)); + AscendC::LocalTensor quantFloat = quantFloatBuf.Get(); + AscendC::LocalTensor absHalf = absHalfBuf.Get(); + AscendC::LocalTensor reduceOutHalf = reduceOutBuf.Get(); + AscendC::LocalTensor reduceTmpHalf = reduceTmpBuf.Get(); + AscendC::LocalTensor logical = logicalBuf.Get(); + AscendC::LocalTensor logicalFloat = logical.template ReinterpretCast(); + AscendC::LocalTensor logicalInt = logical.template ReinterpretCast(); + AscendC::GlobalTensor expertOut; + expertOut.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(expertOutGM), selfSendCnt * h); + + const int64_t begin = selfSendCnt * laneId / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t end = selfSendCnt * (laneId + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const uint64_t readyValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineTxRouteReady); + const uint32_t inputAlignedElements = static_cast(inputBytes / sizeof(half)); + const uint8_t inputPadding = static_cast(inputAlignedElements - static_cast(h)); + const uint8_t repeats = static_cast(blockCount); + const AscendC::TEventID quantizeEvent = pipe.FetchEventID(); + const AscendC::TEventID scalarToVectorEvent = pipe.FetchEventID(); + bool copyPending = false; + int64_t pendingRoute = 0; + int64_t pendingSlot = 0; + int64_t readyBatchFirstRoute = 0; + int64_t readyBatchCount = 0; + bool firstReadyPublished = false; + + AscendC::Duplicate(logical, static_cast(0), static_cast(logicalBytes)); + AscendC::Duplicate(packed0, TileXR::DATA_AS_FLAG_READY_VALUE, + static_cast(routeStride / sizeof(float))); + AscendC::Duplicate(packed1, TileXR::DATA_AS_FLAG_READY_VALUE, + static_cast(routeStride / sizeof(float))); + + for (int64_t route = begin; route < end; ++route) { + const int64_t slot = (route - begin) % kPipelineBufferCount; + const event_t eventId = PipelineEvent(slot); + EnqueuePackInput(expertOut, route, h, inputPadding, inputQueue); + const uint64_t inputWaitStart = ProfileBegin(perfTrace); + AscendC::LocalTensor input = inputQueue.DeQue(); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_INPUT_WAIT, inputWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_INPUT_WAIT, 0, + static_cast(h * static_cast(sizeof(half)))); + + const uint64_t quantizeHeadStart = ProfileBegin(perfTrace); + if (slot == 0) { + StartPackQuantization(input, packed0, quantFloat, absHalf, reduceOutHalf, reduceTmpHalf, + logical, h, logicalBytes, routeStride, quantizeEvent); + } else { + StartPackQuantization(input, packed1, quantFloat, absHalf, reduceOutHalf, reduceTmpHalf, + logical, h, logicalBytes, routeStride, quantizeEvent); + } + if (slot == 0) { + FinishPackQuantization(input, packed0, quantFloat, reduceOutHalf, logical, logicalFloat, + logicalInt, h, blockCount, repeats, quantizeEvent, scalarToVectorEvent); + } else { + FinishPackQuantization(input, packed1, quantFloat, reduceOutHalf, logical, logicalFloat, + logicalInt, h, blockCount, repeats, quantizeEvent, scalarToVectorEvent); + } + inputQueue.FreeTensor(input); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, quantizeHeadStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, 0, static_cast(h)); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, 1, static_cast(blockCount)); + + GM_ADDR txRouteAddr = workspaceGM + txDataOffset + route * routeStride; + const uint64_t dataSubmitStart = ProfileBegin(perfTrace); + if (slot == 0) { + StartPackPublish(txRouteAddr, routeStride, packed0, eventId); + } else { + StartPackPublish(txRouteAddr, routeStride, packed1, eventId); + } + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TX_DATA_SUBMIT, dataSubmitStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_DATA_SUBMIT, 0, + static_cast(routeStride)); + if (copyPending) { + FinishPackCopy(workspaceGM + txDataOffset + pendingRoute * routeStride, + routeStride, PipelineEvent(pendingSlot), perfTrace, perfStats); + if (readyBatchCount == 0) { + readyBatchFirstRoute = pendingRoute; + } + ++readyBatchCount; + if (readyBatchCount == 1) { + const uint64_t readyPublishStart = ProfileBegin(perfTrace); + PublishTxReadyBatch(workspaceGM, txReadyOffset, + readyBatchFirstRoute, readyBatchCount, readyValue); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, readyPublishStart); + const uint64_t publishedLineCount = static_cast(readyBatchCount); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 1, + publishedLineCount); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 2, + static_cast(readyBatchCount)); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 3, + publishedLineCount * TileXREp::kEpUrmaCombineCacheLineBytes); + if (!firstReadyPublished) { + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, firstReadyStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, 0, + static_cast(readyBatchFirstRoute)); + firstReadyPublished = true; + } + readyBatchCount = 0; + } + } + copyPending = true; + pendingRoute = route; + pendingSlot = slot; + if (route + 1 == end) { + FinishPackCopy(txRouteAddr, routeStride, eventId, perfTrace, perfStats); + if (readyBatchCount == 0) { + readyBatchFirstRoute = route; + } + ++readyBatchCount; + copyPending = false; + const uint64_t readyPublishStart = ProfileBegin(perfTrace); + PublishTxReadyBatch(workspaceGM, txReadyOffset, + readyBatchFirstRoute, readyBatchCount, readyValue); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, readyPublishStart); + const uint64_t publishedLineCount = static_cast(readyBatchCount); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 1, + publishedLineCount); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 2, + static_cast(readyBatchCount)); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 3, + publishedLineCount * TileXREp::kEpUrmaCombineCacheLineBytes); + if (!firstReadyPublished) { + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, firstReadyStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, 0, + static_cast(readyBatchFirstRoute)); + firstReadyPublished = true; + } + readyBatchCount = 0; + } + } + pipe.ReleaseEventID(quantizeEvent); + pipe.ReleaseEventID(scalarToVectorEvent); +} + +__aicore__ inline void StartRouteReadyCopy( + GM_ADDR routeAddr, int64_t blockCount, AscendC::LocalTensor &flags) +{ + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + routeAddr + TileXREp::kEpUrmaCombinePayloadBytes)); + SetNoCacheRead(src); + AscendC::DataCopyExtParams copyParams { + static_cast(blockCount), static_cast(sizeof(float)), + static_cast(TileXREp::kEpUrmaCombineDataBlockBytes - sizeof(float)), 0, 0}; + AscendC::DataCopyPadExtParams padParams { + true, 0, static_cast(TileXR::DATA_AS_FLAG_FLAG_FLOATS - 1U), + 0.0f}; + AscendC::DataCopyPad(flags, src, copyParams, padParams); +} + +__aicore__ inline bool CheckRouteReady( + GM_ADDR routeAddr, int64_t blockCount, AscendC::LocalTensor &flags) +{ + StartRouteReadyCopy(routeAddr, blockCount, flags); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + for (int64_t block = 0; block < blockCount; ++block) { + if (flags.GetValue(static_cast( + block * TileXR::DATA_AS_FLAG_FLAG_FLOATS)) != TileXR::DATA_AS_FLAG_READY_VALUE) { + return false; + } + } + return true; +} + + +__aicore__ inline void StartUnpackRoute( + GM_ADDR routeAddr, int64_t blockCount, + AscendC::LocalTensor &logical, event_t eventId) +{ + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(routeAddr)); + SetNoCacheRead(src); + AscendC::DataCopyExtParams copyParams { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombinePayloadBytes), + static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0, 0}; + AscendC::DataCopyPadExtParams padParams {false, 0, 0, 0}; + AscendC::DataCopyPad(logical, src, copyParams, padParams); + AscendC::SetFlag(eventId); + AscendC::SetFlag(eventId); +} + +__aicore__ inline void WaitUnpackRoute(event_t eventId) +{ + AscendC::WaitFlag(eventId); + AscendC::WaitFlag(eventId); +} + +__aicore__ inline void StartClearRouteFlags( + GM_ADDR routeAddr, int64_t blockCount, AscendC::LocalTensor &zeroFlags) +{ + AscendC::GlobalTensor dst; + dst.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + routeAddr + TileXREp::kEpUrmaCombinePayloadBytes)); + SetNoCacheWrite(dst); + AscendC::DataCopyExtParams copyParams { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0, + static_cast(TileXREp::kEpUrmaCombinePayloadBytes), 0}; + AscendC::DataCopyPad(dst, zeroFlags, copyParams); +} + +__aicore__ inline void FinishTokenRouteFlagClears() +{ + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + +__aicore__ inline void ReceiveTokens( + GM_ADDR topKWeightsGM, GM_ADDR yOutGM, GM_ADDR workspaceGM, int64_t bs, + int64_t h, int64_t topK, int64_t magic, int64_t rxWindowOffset, int64_t blockCount, int64_t routeStride, + int64_t rxLaneDoneOffset, int64_t errorStatusOffset, int64_t laneId, AscendC::TPipe &pipe, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const int64_t logicalBytes = blockCount * TileXREp::kEpUrmaCombinePayloadBytes; + const int64_t flagBytes = blockCount * TileXREp::kEpUrmaCombineFlagBytes; + const int64_t accumBytes = AlignUpInt64(h * static_cast(sizeof(float)), kCursorAlignment); + const int64_t outputBytes = AlignUpInt64(h * static_cast(sizeof(half)), kCursorAlignment); + const int64_t dequantBytes = AlignUpInt64(h * static_cast(sizeof(float)), kCursorAlignment); + const int64_t weightBytes = AlignUpInt64( + topK * static_cast(sizeof(float)), kCursorAlignment); + + AscendC::TBuf logicalBuf; + AscendC::TBuf flagBuf; + AscendC::TBuf clearFlagBuf; + AscendC::TBuf accumBuf; + AscendC::TBuf outputBuf; + AscendC::TBuf dequantBuf; + AscendC::TBuf weightBuf; + pipe.InitBuffer(logicalBuf, static_cast(kPipelineBufferCount * logicalBytes)); + pipe.InitBuffer(flagBuf, static_cast(flagBytes)); + pipe.InitBuffer(clearFlagBuf, static_cast(flagBytes)); + pipe.InitBuffer(accumBuf, static_cast(accumBytes)); + pipe.InitBuffer(outputBuf, static_cast(outputBytes)); + pipe.InitBuffer(dequantBuf, static_cast(dequantBytes)); + pipe.InitBuffer(weightBuf, static_cast(weightBytes)); + AscendC::LocalTensor logical0 = logicalBuf.GetWithOffset( + static_cast(logicalBytes), 0); + AscendC::LocalTensor logical1 = logicalBuf.GetWithOffset( + static_cast(logicalBytes), static_cast(logicalBytes)); + AscendC::LocalTensor logicalFloat0 = logical0.template ReinterpretCast(); + AscendC::LocalTensor logicalFloat1 = logical1.template ReinterpretCast(); + AscendC::LocalTensor logicalInt0 = logical0.template ReinterpretCast(); + AscendC::LocalTensor logicalInt1 = logical1.template ReinterpretCast(); + AscendC::LocalTensor logicalInt80 = logical0.template ReinterpretCast(); + AscendC::LocalTensor logicalInt81 = logical1.template ReinterpretCast(); + AscendC::LocalTensor flags = flagBuf.Get(); + AscendC::LocalTensor zeroFlags = clearFlagBuf.Get(); + AscendC::LocalTensor accum = accumBuf.Get(); + AscendC::LocalTensor output = outputBuf.Get(); + AscendC::LocalTensor dequant = dequantBuf.Get(); + AscendC::LocalTensor weight = weightBuf.Get(); + AscendC::GlobalTensor yOut; + yOut.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(yOutGM), bs * h); + + const uint32_t flagCount = static_cast( + blockCount * TileXR::DATA_AS_FLAG_FLAG_FLOATS); + AscendC::Duplicate(zeroFlags, 0.0f, flagCount); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + const int64_t tokenBegin = bs * laneId / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t tokenEnd = bs * (laneId + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const bool profileFine = TileXR::TileXRPerfTraceEnabled(perfTrace); + int64_t tokensRemaining = tokenEnd - tokenBegin; + int64_t activeToken[kRxTokenScheduleWindow] = {}; + bool tokenActive[kRxTokenScheduleWindow] = {}; + uint32_t routeReadyMask[kRxTokenScheduleWindow] = {}; + int64_t nextToken = tokenBegin; + for (int64_t slot = 0; slot < kRxTokenScheduleWindow && nextToken < tokenEnd; ++slot) { + activeToken[slot] = nextToken++; + routeReadyMask[slot] = 0; + tokenActive[slot] = true; + } + int64_t nextCandidate = 0; + const uint32_t allRoutesReadyMask = (1U << static_cast(topK)) - 1U; + while (tokensRemaining > 0) { + int64_t token = 0; + int64_t selectedSlot = -1; + uint64_t bypassedTokens = 0; + const uint64_t flagPollStart = ProfileBegin(perfTrace); + uint64_t pollPasses = 0; + uint64_t routeChecks = 0; + uint64_t readyMisses = 0; + while (selectedSlot < 0) { + for (int64_t offset = 0; offset < kRxTokenScheduleWindow; ++offset) { + int64_t candidate = nextCandidate + offset; + if (candidate >= kRxTokenScheduleWindow) { + candidate -= kRxTokenScheduleWindow; + } + if (!tokenActive[candidate]) { + continue; + } + if (profileFine) { + ++pollPasses; + } + const int64_t candidateToken = activeToken[candidate]; + uint32_t readyMask = routeReadyMask[candidate]; + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const uint32_t routeBit = 1U << static_cast(topKId); + if ((readyMask & routeBit) != 0) { + continue; + } + const int64_t routeIndex = candidateToken * topK + topKId; + GM_ADDR routeAddr = workspaceGM + rxWindowOffset + routeIndex * routeStride; + if (profileFine) { + ++routeChecks; + } + if (CheckRouteReady(routeAddr, blockCount, flags)) { + readyMask |= routeBit; + } else { + if (profileFine) { + ++readyMisses; + } + } + } + routeReadyMask[candidate] = readyMask; + if (readyMask == allRoutesReadyMask) { + selectedSlot = candidate; + break; + } + } + } + token = activeToken[selectedSlot]; + for (int64_t slot = 0; slot < kRxTokenScheduleWindow; ++slot) { + if (tokenActive[slot] && activeToken[slot] < token) { + ++bypassedTokens; + } + } + nextCandidate = selectedSlot + 1; + if (nextCandidate >= kRxTokenScheduleWindow) { + nextCandidate = 0; + } + CachelessAcquireBarrier(); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, flagPollStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 0, pollPasses); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 1, routeChecks); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 2, readyMisses); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 3, + routeChecks * static_cast(blockCount)); + + LoadTopKWeights(topKWeightsGM, token * topK, topK, weight); + if (topK > 0) { + GM_ADDR firstRouteAddr = workspaceGM + rxWindowOffset + token * topK * routeStride; + StartUnpackRoute(firstRouteAddr, blockCount, logical0, PipelineEvent(0)); + } + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const int64_t slot = topKId % kPipelineBufferCount; + const event_t eventId = PipelineEvent(slot); + AscendC::LocalTensor &routeLogicalFloat = slot == 0 ? logicalFloat0 : logicalFloat1; + AscendC::LocalTensor &routeLogicalInt = slot == 0 ? logicalInt0 : logicalInt1; + AscendC::LocalTensor &routeLogicalInt8 = slot == 0 ? logicalInt80 : logicalInt81; + const int64_t routeIndex = token * topK + topKId; + GM_ADDR routeAddr = workspaceGM + rxWindowOffset + routeIndex * routeStride; + const uint64_t unpackWaitStart = ProfileBegin(perfTrace); + WaitUnpackRoute(eventId); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_UNPACK_WAIT, unpackWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_UNPACK_WAIT, 0, + static_cast(blockCount * TileXREp::kEpUrmaCombinePayloadBytes)); + + const uint64_t dequantStart = ProfileBegin(perfTrace); + const float scale = routeLogicalFloat.GetValue(0); + const int32_t quantMode = routeLogicalInt.GetValue(1); + if (!(scale > 0.0f) || quantMode != TileXREp::kEpUrmaCombineQuantModeInt8PerRoute) { + StoreError(workspaceGM, errorStatusOffset, magic, + TileXREp::kEpUrmaCombineStatusInvalidQuantHeader); + } + const float effectiveScale = scale > 0.0f ? scale : 1.0f; + const float topKWeight = weight.GetValue(static_cast(topKId)); + AscendC::Cast(output, + routeLogicalInt8[TileXREp::kEpUrmaCombineQuantHeaderBytes], AscendC::RoundMode::CAST_NONE, + static_cast(h)); + AscendC::PipeBarrier(); + if (topKId + 1 < topK) { + const int64_t nextSlot = (slot + 1) % kPipelineBufferCount; + GM_ADDR nextRouteAddr = routeAddr + routeStride; + AscendC::LocalTensor &nextLogical = nextSlot == 0 ? logical0 : logical1; + StartUnpackRoute(nextRouteAddr, blockCount, + nextLogical, PipelineEvent(nextSlot)); + } + AscendC::Cast(dequant, output, AscendC::RoundMode::CAST_NONE, + static_cast(h)); + AscendC::PipeBarrier(); + if (topKId == 0) { + AscendC::Muls(accum, dequant, effectiveScale * topKWeight, + static_cast(h)); + } else { + AscendC::Axpy(accum, dequant, effectiveScale * topKWeight, + static_cast(h)); + } + AscendC::PipeBarrier(); + StartClearRouteFlags(routeAddr, blockCount, zeroFlags); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, dequantStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, 0, + static_cast(h)); + ProfileAux(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, 1, + static_cast(blockCount)); + } + + const uint64_t clearFenceStart = ProfileBegin(perfTrace); + FinishTokenRouteFlagClears(); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, clearFenceStart); + + const uint64_t outputStart = ProfileBegin(perfTrace); + AscendC::Cast(output, accum, AscendC::RoundMode::CAST_ROUND, + static_cast(h)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyExtParams outputParams { + 1, static_cast(h * static_cast(sizeof(half))), 0, 0, 0}; + AscendC::DataCopyPad(yOut[token * h], output, outputParams); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_OUTPUT, outputStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_OUTPUT, 0, + static_cast(h * static_cast(sizeof(half)))); + ProfileAux(perfTrace, perfStats, PerfStage::RX_OUTPUT, 1, bypassedTokens); + ProfileAux(perfTrace, perfStats, PerfStage::RX_OUTPUT, 2, bypassedTokens == 0 ? 0U : 1U); + + --tokensRemaining; + routeReadyMask[selectedSlot] = 0; + if (nextToken < tokenEnd) { + activeToken[selectedSlot] = nextToken++; + tokenActive[selectedSlot] = true; + } else { + tokenActive[selectedSlot] = false; + } + } + + StoreControlValue(workspaceGM + rxLaneDoneOffset + + laneId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone)); +} + +__aicore__ inline void CopySelfRoute( + GM_ADDR srcAddr, GM_ADDR dstAddr, int64_t blockCount, AscendC::LocalTensor &payload, + AscendC::LocalTensor &readyFlags) +{ + AscendC::GlobalTensor src; + AscendC::GlobalTensor dst; + src.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(srcAddr)); + dst.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(dstAddr)); + SetNoCacheRead(src); + SetNoCacheWrite(dst); + AscendC::DataCopyExtParams payloadIn { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombinePayloadBytes), + static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0, 0}; + AscendC::DataCopyPadExtParams noPad {false, 0, 0, 0}; + AscendC::DataCopyPad(payload, src, payloadIn, noPad); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::DataCopyExtParams payloadOut { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombinePayloadBytes), + 0, static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0}; + AscendC::DataCopyPad(dst, payload, payloadOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstFlags; + dstFlags.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + dstAddr + TileXREp::kEpUrmaCombinePayloadBytes)); + SetNoCacheWrite(dstFlags); + AscendC::DataCopyExtParams flagOut { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombineFlagBytes), + 0, static_cast(TileXREp::kEpUrmaCombinePayloadBytes), 0}; + AscendC::DataCopyPad(dstFlags, readyFlags, flagOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + +__aicore__ inline bool SendRoutes(const __gm__ TileXR::CommArgs *args, GM_ADDR assistInfoGM, + GM_ADDR workspaceGM, int64_t selfSendCnt, int64_t bs, int64_t topK, int64_t magic, int64_t rxWindowOffset, + int64_t blockCount, int64_t routeStride, int64_t txReadyOffset, int64_t txDataOffset, + int64_t senderDoneOffset, int64_t errorStatusOffset, int64_t senderId, AscendC::TPipe &pipe, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const int64_t cursorBytes = AlignUpInt64( + TileXREp::kEpUrmaCombinePackLaneCount * static_cast(sizeof(uint32_t)), kCursorAlignment); + const int64_t usedPeerBytes = AlignUpInt64(args->rankSize, kCursorAlignment); + const int64_t selfCopyPayloadBytes = blockCount * TileXREp::kEpUrmaCombinePayloadBytes; + const int64_t selfCopyFlagBytes = blockCount * TileXREp::kEpUrmaCombineFlagBytes; + AscendC::TBuf cursorBuf; + AscendC::TBuf usedPeerBuf; + AscendC::TBuf selfCopyBuf; + pipe.InitBuffer(cursorBuf, static_cast(cursorBytes)); + pipe.InitBuffer(usedPeerBuf, static_cast(usedPeerBytes)); + pipe.InitBuffer(selfCopyBuf, + static_cast(selfCopyPayloadBytes + selfCopyFlagBytes)); + AscendC::LocalTensor cursor = cursorBuf.Get(); + AscendC::LocalTensor usedPeer = usedPeerBuf.Get(); + AscendC::LocalTensor selfCopyPayload = selfCopyBuf.GetWithOffset( + static_cast(selfCopyPayloadBytes), 0); + AscendC::LocalTensor selfCopyReady = selfCopyBuf.GetWithOffset( + static_cast(selfCopyFlagBytes / sizeof(float)), + static_cast(selfCopyPayloadBytes)); + AscendC::Duplicate(selfCopyReady, TileXR::DATA_AS_FLAG_READY_VALUE, + static_cast(selfCopyFlagBytes / sizeof(float))); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + + for (int64_t lane = 0; lane < TileXREp::kEpUrmaCombinePackLaneCount; ++lane) { + const int64_t begin = selfSendCnt * lane / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t end = selfSendCnt * (lane + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t laneLength = end - begin; + const int64_t remainder = laneLength % TileXREp::kEpUrmaCombineSendLaneCount; + const bool reverse = (lane & 1) != 0; + + // Preserve the globally balanced residue set, but reverse odd lanes so no sender + // consistently waits for the last route produced by every Pack core. + int64_t rotation = lane % TileXREp::kEpUrmaCombineSendLaneCount; + if (remainder != 0) { + rotation = begin % TileXREp::kEpUrmaCombineSendLaneCount; + if (reverse) { + rotation = (rotation + remainder - 1) % TileXREp::kEpUrmaCombineSendLaneCount; + } + } + const int64_t signedDelta = reverse ? rotation - senderId : senderId - rotation; + const int64_t delta = (signedDelta + TileXREp::kEpUrmaCombineSendLaneCount) % + TileXREp::kEpUrmaCombineSendLaneCount; + cursor.SetValue(static_cast(lane), static_cast(begin + delta)); + } + + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + for (int32_t peer = 0; peer < rankSize; ++peer) { + usedPeer.SetValue(static_cast(peer), 0); + } + const uint64_t readyValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineTxRouteReady); + uint64_t remotePutCount = 0; + bool allScanned = false; + while (!allScanned) { + allScanned = true; + for (int64_t lane = 0; lane < TileXREp::kEpUrmaCombinePackLaneCount; ++lane) { + const int64_t begin = selfSendCnt * lane / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t end = selfSendCnt * (lane + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t route = cursor.GetValue(static_cast(lane)); + if (route >= end) { + continue; + } + allScanned = false; + GM_ADDR readyAddr = workspaceGM + txReadyOffset + + route * TileXREp::kEpUrmaCombineCacheLineBytes; + const uint64_t readyPollStart = ProfileBegin(perfTrace); + const bool routeReady = LoadControlValue(readyAddr) == readyValue; + ProfileEnd(perfTrace, perfStats, PerfStage::TX_READY_POLL, readyPollStart); + if (!routeReady) { + ProfileAux(perfTrace, perfStats, PerfStage::TX_READY_POLL, 0, 1); + continue; + } + CachelessAcquireBarrier(); + ProfileAux(perfTrace, perfStats, PerfStage::TX_READY_POLL, 1, 1); + + const uint64_t metaScanStart = ProfileBegin(perfTrace); + const TileXREp::EpAssistTuple tuple = LoadRouteMeta(assistInfoGM, route); + if (!RouteMetaValid(tuple, rankSize, bs, topK)) { + ProfileEnd(perfTrace, perfStats, PerfStage::TX_META_SCAN, metaScanStart); + ProfileAux(perfTrace, perfStats, PerfStage::TX_META_SCAN, 0, 1); + StoreError(workspaceGM, errorStatusOffset, magic, + TileXREp::kEpUrmaCombineStatusInvalidRoute); + cursor.SetValue(static_cast(lane), static_cast( + route + TileXREp::kEpUrmaCombineSendLaneCount)); + continue; + } + ProfileEnd(perfTrace, perfStats, PerfStage::TX_META_SCAN, metaScanStart); + ProfileAux(perfTrace, perfStats, PerfStage::TX_META_SCAN, + tuple.srcRank == rank ? 1U : 2U, 1); + + const int64_t routeIndex = static_cast(tuple.tokenId) * topK + tuple.topKId; + const int64_t rxRouteOffset = rxWindowOffset + routeIndex * routeStride; + GM_ADDR txRouteAddr = workspaceGM + txDataOffset + route * routeStride; + if (tuple.srcRank == rank) { + const uint64_t selfCopyStart = ProfileBegin(perfTrace); + CopySelfRoute(txRouteAddr, workspaceGM + rxRouteOffset, blockCount, selfCopyPayload, + selfCopyReady); + ProfileEnd(perfTrace, perfStats, PerfStage::SELF_COPY, selfCopyStart); + ProfileAux(perfTrace, perfStats, PerfStage::SELF_COPY, 0, + static_cast(routeStride)); + } else { + const uint64_t postStart = ProfileBegin(perfTrace); + TileXREp::EpUrmaUDMAPutNbi(args, tuple.srcRank, + reinterpret_cast<__gm__ uint8_t *>(txRouteAddr), + static_cast(rxRouteOffset), static_cast(routeStride), + static_cast(senderId)); + ++remotePutCount; + usedPeer.SetValue(static_cast(tuple.srcRank), 1); + ProfileEnd(perfTrace, perfStats, PerfStage::UDMA_POST, postStart); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 0, + static_cast(routeStride)); + } + cursor.SetValue(static_cast(lane), static_cast( + route + TileXREp::kEpUrmaCombineSendLaneCount)); + } + } + + + const uint64_t doorbellCommitCount = remotePutCount; + const uint64_t activeSendSqCount = 0; + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 1, + doorbellCommitCount); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 2, 1); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 3, + activeSendSqCount); + + __gm__ TileXR::UDMAInfo *udmaInfo = TileXR::GetUDMAInfo(args); + for (int32_t peer = 0; peer < rankSize; ++peer) { + if (usedPeer.GetValue(static_cast(peer)) == 0) { + continue; + } + const uint64_t wqeCntAddr = TileXR::UDMAGetWQCtx( + udmaInfo, static_cast(peer), static_cast(senderId))->wqeCntAddr; + bool queueAlreadyDrained = false; + for (int32_t previous = 0; previous < peer; ++previous) { + if (usedPeer.GetValue(static_cast(previous)) != 0 && + TileXR::UDMAGetWQCtx(udmaInfo, static_cast(previous), + static_cast(senderId))->wqeCntAddr == wqeCntAddr) { + queueAlreadyDrained = true; + break; + } + } + if (queueAlreadyDrained) { + continue; + } + const uint64_t quietStart = ProfileBegin(perfTrace); + (void)TileXREp::EpUrmaUDMAQuiet(args, peer, static_cast(senderId)); + ProfileEnd(perfTrace, perfStats, PerfStage::UDMA_QUIET, quietStart); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_QUIET, 0, 1); + } + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineSenderDone)); + return true; +} + + +__aicore__ inline bool WaitLocalLines( + GM_ADDR base, int64_t lineCount, uint64_t expected) +{ + for (int64_t line = 0; line < lineCount; ++line) { + GM_ADDR lineAddr = base + line * TileXREp::kEpUrmaCombineCacheLineBytes; + while (LoadControlValue(lineAddr) != expected) { + } + } + CachelessAcquireBarrier(); + return true; +} + +__aicore__ inline bool ReleaseGenerationReached( + uint64_t observed, uint32_t expectedGeneration) +{ + const uint32_t observedStep = static_cast(observed); + const uint32_t observedGeneration = static_cast(observed >> 32U); + return observedStep == TileXREp::kEpUrmaCombineRxBufferReleased && + static_cast(observedGeneration - expectedGeneration) >= 0; +} + +__aicore__ inline bool WaitDeferredRoundCredit( + const __gm__ TileXR::CommArgs *args, GM_ADDR workspaceGM, int64_t magic, + int64_t roundDoneOffset, int64_t senderDoneOffset, int64_t roundCreditOffset, + int64_t senderId, + GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t globalWaitStart = ProfileBegin(perfTrace); + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + const int64_t creditLaneCount = rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + rankSize : TileXREp::kEpUrmaCombineSendLaneCount; + uint64_t pollLoads = 0; + uint32_t expectedGeneration = 0; + if (senderId < creditLaneCount) { + const uint64_t previousRelease = LoadControlValue(workspaceGM + roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes); + const uint32_t previousStep = static_cast(previousRelease); + expectedGeneration = static_cast(previousRelease >> 32U); + if (previousRelease != 0 && previousStep != TileXREp::kEpUrmaCombineRxBufferReleased) { + return false; + } + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(creditLaneCount)) { + if (peer == rank || previousRelease == 0) { + continue; + } + // The same-parity roundDone line is the previous release source. Complete its + // QP before this launch can eventually overwrite that line with the next release. + (void)TileXREp::EpUrmaUDMAQuiet( + args, peer, static_cast(senderId)); + GM_ADDR peerRelease = workspaceGM + roundDoneOffset + + peer * TileXREp::kEpUrmaCombineCacheLineBytes; + while (true) { + ++pollLoads; + if (ReleaseGenerationReached( + LoadControlValue(peerRelease), expectedGeneration)) { + break; + } + } + } + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditShardDone)); + } + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + senderDoneOffset, creditLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditShardDone))) { + return false; + } + StoreControlValue(workspaceGM + roundCreditOffset, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditRun)); + } + if (!WaitLocalLines(workspaceGM + roundCreditOffset, 1, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditRun))) { + return false; + } + ProfileEnd(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, globalWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, 0, pollLoads); + ProfileAux(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, 1, expectedGeneration); + return true; +} + +__aicore__ inline bool WaitLocalLineShard( + GM_ADDR base, int64_t lineCount, int64_t firstLine, int64_t lineStride, + uint64_t expected) +{ + for (int64_t line = firstLine; line < lineCount; line += lineStride) { + GM_ADDR lineAddr = base + line * TileXREp::kEpUrmaCombineCacheLineBytes; + while (LoadControlValue(lineAddr) != expected) { + } + } + CachelessAcquireBarrier(); + return true; +} + +__aicore__ inline bool WaitForSynchronizedStart( + const __gm__ TileXR::CommArgs *args, + GM_ADDR workspaceGM, int64_t magic, int64_t rxLaneDoneOffset, int64_t senderDoneOffset, + int64_t roundPublishOffset, int64_t startGateOffset, int64_t blockIdx, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t gateStart = ProfileBegin(perfTrace); + const uint64_t localReady = EncodeControlValue( + magic, TileXREp::kEpUrmaCombineStartLocalReady); + const bool isPackReceive = blockIdx < TileXREp::kEpUrmaCombinePackLaneCount; + if (isPackReceive) { + StoreControlValue(workspaceGM + rxLaneDoneOffset + + blockIdx * TileXREp::kEpUrmaCombineCacheLineBytes, + localReady); + } else { + const int64_t senderId = blockIdx - TileXREp::kEpUrmaCombinePackLaneCount; + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + localReady); + } + + const uint64_t startRun = EncodeControlValue(magic, TileXREp::kEpUrmaCombineStartRun); + if (isPackReceive) { + if (!WaitLocalLines(workspaceGM + roundPublishOffset, 1, startRun)) { + return false; + } + ProfileEnd(perfTrace, perfStats, PerfStage::START_GATE, gateStart); + return true; + } + + const int64_t senderId = blockIdx - TileXREp::kEpUrmaCombinePackLaneCount; + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + const uint64_t rankReady = EncodeControlValue( + magic, TileXREp::kEpUrmaCombineStartRankReady); + GM_ADDR rankReadyAddr = workspaceGM + startGateOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes; + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + rxLaneDoneOffset, TileXREp::kEpUrmaCombinePackLaneCount, + localReady) || + !WaitLocalLines(workspaceGM + senderDoneOffset, TileXREp::kEpUrmaCombineSendLaneCount, + localReady)) { + return false; + } + StoreControlValue(rankReadyAddr, rankReady); + StoreControlValue(workspaceGM + roundPublishOffset, rankReady); + } + if (!WaitLocalLines(workspaceGM + roundPublishOffset, 1, rankReady)) { + return false; + } + + const uint32_t qpIdx = static_cast(senderId); + const uint64_t remoteOffset = static_cast(startGateOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes); + uint64_t publishCount = 0; + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { + if (peer == rank) { + continue; + } + TileXREp::EpUrmaUDMAPutNbi(args, peer, + reinterpret_cast<__gm__ uint8_t *>(rankReadyAddr), remoteOffset, + static_cast(TileXREp::kEpUrmaCombineCacheLineBytes), qpIdx); + ++publishCount; + } + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { + if (peer == rank) { + continue; + } + (void)TileXREp::EpUrmaUDMAQuiet(args, peer, qpIdx); + } + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 0, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 1, + publishCount * TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 2, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 3, publishCount); + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineStartPublishDone)); + + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + senderDoneOffset, TileXREp::kEpUrmaCombineSendLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineStartPublishDone)) || + !WaitLocalLines(workspaceGM + startGateOffset, rankSize, rankReady)) { + return false; + } + StoreControlValue(workspaceGM + roundPublishOffset, startRun); + } + if (!WaitLocalLines(workspaceGM + roundPublishOffset, 1, startRun)) { + return false; + } + ProfileEnd(perfTrace, perfStats, PerfStage::START_GATE, gateStart); + return true; +} + +__aicore__ inline void RecordRoundPublishCounters( + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats, uint64_t publishCount) +{ + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 0, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 1, + publishCount * TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 2, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 3, publishCount); +} + +__aicore__ inline uint64_t StartParallelRoundPublish( + const __gm__ TileXR::CommArgs *args, + GM_ADDR workspaceGM, int64_t magic, int64_t roundDoneOffset, GM_ADDR perfTrace) +{ + const uint64_t publishStart = ProfileBegin(perfTrace); + const uint64_t roundValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxBufferReleased); + StoreControlValue(workspaceGM + roundDoneOffset + + args->rank * TileXREp::kEpUrmaCombineCacheLineBytes, + roundValue); + return publishStart; +} +__aicore__ inline uint64_t PublishRoundShard( + const __gm__ TileXR::CommArgs *args, + GM_ADDR workspaceGM, int64_t magic, int64_t roundDoneOffset, int64_t senderId, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + const uint64_t roundValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxBufferReleased); + GM_ADDR publishAddr = workspaceGM + roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes; + if (!WaitLocalLines(publishAddr, 1, roundValue)) { + return ~0ULL; + } + + const uint64_t shardStart = senderId == 0 ? 0 : ProfileBegin(perfTrace); + const uint32_t qpIdx = static_cast(senderId); + const uint64_t remoteOffset = static_cast(roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes); + uint64_t publishCount = 0; + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { + if (peer == rank) { + continue; + } + TileXREp::EpUrmaUDMAPutNbi(args, peer, + reinterpret_cast<__gm__ uint8_t *>(publishAddr), remoteOffset, + static_cast(TileXREp::kEpUrmaCombineCacheLineBytes), qpIdx); + ++publishCount; + } + if (senderId != 0) { + ProfileEnd(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, shardStart); + RecordRoundPublishCounters(perfTrace, perfStats, publishCount); + } + return publishCount; +} + +__aicore__ inline void FinishParallelRoundPublish( + uint64_t publishStart, uint64_t sender0PublishCount, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + ProfileEnd(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, publishStart); + RecordRoundPublishCounters(perfTrace, perfStats, sender0PublishCount); +} + +} // namespace + +extern "C" __global__ __aicore__ void tilexr_ep_urma_combine_kernel(GM_ADDR commArgsGM, + GM_ADDR expertOutGM, GM_ADDR assistInfoForCombineGM, GM_ADDR topKWeightsGM, GM_ADDR yOutGM, + GM_ADDR workspaceGM, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, + int64_t workspaceBytes, int64_t magic, int64_t commBytes, int64_t blockCount, int64_t routeStride, + int64_t rxWindowBytes, int64_t rxWindowOffset0, int64_t rxWindowOffset1, int64_t roundDoneOffset0, + int64_t roundDoneOffset1, int64_t rxLaneDoneOffset, int64_t senderDoneOffset, + int64_t roundPublishOffset, int64_t roundCreditOffset, int64_t startGateOffset, + int64_t runStartGate, int64_t errorStatusOffset, + int64_t txReadyOffset, int64_t txDataOffset, GM_ADDR perfTrace, int64_t perfTraceBytes, + GM_ADDR strictKernelCycles) +{ + if constexpr (g_coreType == AscendC::AIV) { + if (commArgsGM == nullptr || topKWeightsGM == nullptr || yOutGM == nullptr || workspaceGM == nullptr || + (selfSendCnt > 0 && (expertOutGM == nullptr || assistInfoForCombineGM == nullptr)) || + selfSendCnt < 0 || bs <= 0 || h <= 0 || h > TileXREp::kEpUrmaCombineMaxHidden || topK <= 0 || + topK > TileXREp::kEpUrmaCombineMaxTopK || magic <= 0 || commBytes != + TileXREp::kEpUrmaCombineQuantHeaderBytes + AlignUpInt64(h, kCursorAlignment) || blockCount <= 0 || + blockCount != (commBytes + TileXREp::kEpUrmaCombinePayloadBytes - 1) / + TileXREp::kEpUrmaCombinePayloadBytes || + blockCount > TileXREp::kEpUrmaCombineMaxBlocksPerRoute || routeStride <= 0 || + routeStride != blockCount * TileXREp::kEpUrmaCombineDataBlockBytes || + (runStartGate != 0 && runStartGate != 1) || + roundCreditOffset <= 0 || + roundCreditOffset > workspaceBytes - TileXREp::kEpUrmaCombineCacheLineBytes || + rxWindowBytes <= 0 || txReadyOffset < 0 || txDataOffset < txReadyOffset || + txDataOffset > workspaceBytes || selfSendCnt > + (txDataOffset - txReadyOffset) / TileXREp::kEpUrmaCombineCacheLineBytes || + selfSendCnt > (workspaceBytes - txDataOffset) / routeStride) { + return; + } + auto args = reinterpret_cast<__gm__ TileXR::CommArgs *>(commArgsGM); + if (args->rankSize <= 0 || args->rankSize > TileXR::TILEXR_MAX_RANK_SIZE || args->rank < 0 || + args->rank >= args->rankSize || (args->rankSize > 1 && !TileXR::UDMARegistryEnabled(args))) { + return; + } + if (args->rankSize > 1) { + __gm__ TileXR::UDMAInfo *udmaInfo = TileXR::GetUDMAInfo(args); + if (udmaInfo == nullptr || udmaInfo->qpNum < + static_cast(TileXREp::kEpUrmaCombineRequiredQpCount)) { + return; + } + } + if (startGateOffset < 0 || startGateOffset > workspaceBytes || + args->rankSize > (workspaceBytes - startGateOffset) / + TileXREp::kEpUrmaCombineCacheLineBytes) { + return; + } + if (strictKernelCycles != nullptr && perfTrace != nullptr) { + return; + } + if (!ProfileBufferValid(perfTrace, perfTraceBytes, args->rank, args->rankSize)) { + perfTrace = nullptr; + } + GM_ADDR kernelPerfTrace = perfTrace; + + const int64_t blockIdx = AscendC::GetBlockIdx(); + if (blockIdx < 0 || blockIdx >= TileXREp::kEpUrmaCombineAivCount) { + return; + } + const uint32_t perfRank = static_cast(args->rank); + const uint32_t perfCore = static_cast(blockIdx); + __ubuf__ PerfStats *perfStats = reinterpret_cast<__ubuf__ PerfStats *>( + TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET); + TileXR::TileXRPerfLocalStatsInit(kernelPerfTrace, perfStats, perfRank, perfCore, + TileXREp::kEpUrmaCombinePerfStageCount); + GM_ADDR finePerfTrace = nullptr; + if (kernelPerfTrace != nullptr) { + const uint32_t profileDetail = + reinterpret_cast<__gm__ TileXR::TileXRPerfTraceHeader *>(kernelPerfTrace)->flags; + if (profileDetail >= 2) { + finePerfTrace = kernelPerfTrace; + } + if (profileDetail == 0) { + perfTrace = nullptr; + } + } + if (runStartGate != 0) { + if (!WaitForSynchronizedStart(args, workspaceGM, magic, rxLaneDoneOffset, senderDoneOffset, + roundPublishOffset, startGateOffset, blockIdx, perfTrace, perfStats)) { + return; + } + } + const uint64_t kernelStart = ProfileKernelTimingBegin(kernelPerfTrace); + const uint64_t strictKernelStart = StrictKernelTimingBegin(strictKernelCycles); + const int64_t rxBufferIndex = static_cast(magic) & 1U; + const int64_t rxWindowOffset = rxBufferIndex == 0 ? rxWindowOffset0 : rxWindowOffset1; + const int64_t roundDoneOffset = rxBufferIndex == 0 ? roundDoneOffset0 : roundDoneOffset1; + AscendC::TPipe pipe; + if (blockIdx < TileXREp::kEpUrmaCombinePackLaneCount) { + const uint64_t packStart = ProfileBegin(perfTrace); + PackRoutes(expertOutGM, workspaceGM, selfSendCnt, h, magic, blockCount, routeStride, + txReadyOffset, txDataOffset, blockIdx, pipe, finePerfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TOTAL, packStart); + pipe.Reset(); + const uint64_t receiveStart = ProfileBegin(perfTrace); + ReceiveTokens(topKWeightsGM, yOutGM, workspaceGM, bs, h, topK, magic, rxWindowOffset, + blockCount, routeStride, rxLaneDoneOffset, errorStatusOffset, blockIdx, pipe, + finePerfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::RECEIVE_TOTAL, receiveStart); + StrictKernelTimingFinish(strictKernelCycles, perfCore, strictKernelStart); + ProfileFinish(kernelPerfTrace, perfRank, perfCore, perfStats, kernelStart); + return; + } + + const int64_t senderId = blockIdx - TileXREp::kEpUrmaCombinePackLaneCount; + if (senderId >= TileXREp::kEpUrmaCombineSendLaneCount) { + return; + } + if (!WaitDeferredRoundCredit(args, workspaceGM, magic, roundDoneOffset, + senderDoneOffset, roundCreditOffset, senderId, + perfTrace, finePerfTrace, perfStats)) { + return; + } + const uint64_t sendStart = ProfileBegin(perfTrace); + if (!SendRoutes(args, assistInfoForCombineGM, workspaceGM, selfSendCnt, bs, topK, magic, + rxWindowOffset, blockCount, routeStride, txReadyOffset, txDataOffset, senderDoneOffset, + errorStatusOffset, senderId, pipe, finePerfTrace, perfStats)) { + return; + } + ProfileEnd(perfTrace, perfStats, PerfStage::SEND_TOTAL, sendStart); + const int64_t publisherCount = args->rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + args->rankSize : TileXREp::kEpUrmaCombineSendLaneCount; + uint64_t publishStart = 0; + if (senderId < publisherCount) { + // Only release publishers participate, reducing sender0's serial fan-in while + // preserving complete coverage of every RX lane. + const uint64_t receiveWaitStart = ProfileBegin(perfTrace); + if (!WaitLocalLineShard(workspaceGM + rxLaneDoneOffset, + TileXREp::kEpUrmaCombinePackLaneCount, senderId, publisherCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone))) { + return; + } + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone)); + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + senderDoneOffset, publisherCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone))) { + return; + } + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, 0, + static_cast(TileXREp::kEpUrmaCombinePackLaneCount)); + publishStart = StartParallelRoundPublish(args, workspaceGM, magic, + roundDoneOffset, perfTrace); + } else { + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + } + } + + uint64_t publishCount = 0; + if (senderId < publisherCount) { + publishCount = PublishRoundShard(args, workspaceGM, magic, + roundDoneOffset, senderId, perfTrace, perfStats); + if (publishCount == ~0ULL) { + return; + } + } + if (senderId == 0) { + FinishParallelRoundPublish(publishStart, publishCount, perfTrace, perfStats); + } + StrictKernelTimingFinish(strictKernelCycles, perfCore, strictKernelStart); + ProfileFinish(kernelPerfTrace, perfRank, perfCore, perfStats, kernelStart); + } +} + + +void launch_tilexr_ep_urma_combine_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, + GM_ADDR expertOut, GM_ADDR assistInfoForCombine, GM_ADDR topKWeights, GM_ADDR yOut, GM_ADDR workspace, + int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, int64_t workspaceBytes, int64_t magic, + int64_t commBytes, int64_t blockCount, int64_t routeStride, int64_t rxWindowBytes, int64_t rxWindowOffset0, + int64_t rxWindowOffset1, int64_t roundDoneOffset0, int64_t roundDoneOffset1, int64_t rxLaneDoneOffset, + int64_t senderDoneOffset, int64_t roundPublishOffset, int64_t roundCreditOffset, + int64_t startGateOffset, int64_t runStartGate, + int64_t errorStatusOffset, int64_t txReadyOffset, int64_t txDataOffset, GM_ADDR perfTrace, int64_t perfTraceBytes, + GM_ADDR strictKernelCycles) +{ + tilexr_ep_urma_combine_kernel<<>>(commArgs, expertOut, assistInfoForCombine, + topKWeights, yOut, workspace, selfSendCnt, bs, h, topK, workspaceBytes, magic, commBytes, + blockCount, routeStride, rxWindowBytes, rxWindowOffset0, rxWindowOffset1, + roundDoneOffset0, roundDoneOffset1, rxLaneDoneOffset, senderDoneOffset, roundPublishOffset, + roundCreditOffset, startGateOffset, runStartGate, errorStatusOffset, txReadyOffset, txDataOffset, + perfTrace, perfTraceBytes, + strictKernelCycles); +} diff --git a/src/include/tilexr_ep.h b/src/include/tilexr_ep.h index 82c79cde..83540c9a 100644 --- a/src/include/tilexr_ep.h +++ b/src/include/tilexr_ep.h @@ -25,6 +25,24 @@ int TileXRMoeEpCombineV2(void *expertOut, int32_t *assistInfoForCombine, int32_t TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, void *yOut, void *workspace, TileXR::TileXRDataType dtype, aclrtStream stream); +// workspace must be zero-initialized before its first use and registered once +// as the communicator's UDMA region. Reuse is supported for serialized launches. +int TileXRMoeEpCombineUrmaGetWorkspaceSize(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCapacity, int64_t *workspaceBytes); + +int TileXRMoeEpCombineUrmaGetProfileSize(int64_t rankSize, int64_t *profileBytes); + +int TileXRMoeEpCombineUrma(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, TileXR::TileXRDataType dtype, aclrtStream stream); + +// perfTrace points to a device TileXRPerfTrace buffer initialized by the caller. +// Use TileXRMoeEpCombineUrmaGetProfileSize to size the allocation. +int TileXRMoeEpCombineUrmaProfile(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, void *perfTrace, int64_t perfTraceBytes, + TileXR::TileXRDataType dtype, aclrtStream stream); + int TileXRMoeEpDispatchV2(void *x, int32_t *expertIds, void *scales, bool *xActiveMask, void *expertScales, TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, int64_t epWorldSize, int64_t epRankId, int64_t tpWorldSize, int64_t tpRankId, int64_t expertShardType, int64_t sharedExpertNum, diff --git a/tests/comm/unit/test_tilexr_source_guards.cpp b/tests/comm/unit/test_tilexr_source_guards.cpp index 4c7e70e0..51ae6a0e 100644 --- a/tests/comm/unit/test_tilexr_source_guards.cpp +++ b/tests/comm/unit/test_tilexr_source_guards.cpp @@ -183,6 +183,25 @@ void TestCommBuildIncludesProfilingHeaders() CheckContains(commPath, commText, "${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/profiling/"); } +void TestUdmaQpProvisioningIsOperatorDriven() +{ + const std::string rootPath = "CMakeLists.txt"; + const std::string commPath = "src/comm/CMakeLists.txt"; + const auto rootText = ReadFile(rootPath); + const auto commText = ReadFile(commPath); + + CheckContains(rootPath, rootText, + "option(TILEXR_EP_BUILD_URMA_COMBINE \"Build the Ascend950 URMA EP combine operator\" OFF)"); + CheckContains(rootPath, rootText, + "set(TILEXR_UDMA_QP_COUNT \"1\" CACHE STRING"); + CheckContains(rootPath, rootText, + "TILEXR_UDMA_EFFECTIVE_QP_COUNT LESS TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT"); + CheckContains(commPath, commText, + "TILEXR_UDMA_QP_COUNT_VALUE=${TILEXR_UDMA_EFFECTIVE_QP_COUNT}"); + CheckNotContains(commPath, commText, "ascend950"); + CheckNotContains(commPath, commText, "_tilexr_udma_qp_count"); +} + void TestChipMapCoversObservedAscend950Variants() { const std::string path = "src/comm/tilexr_internal.cpp"; @@ -205,6 +224,7 @@ int main() TestRuntimeEnvDoesNotPrependCannDevlib(); TestRootCMakeRespectsAscendDriverOverride(); TestCommBuildIncludesProfilingHeaders(); + TestUdmaQpProvisioningIsOperatorDriven(); TestChipMapCoversObservedAscend950Variants(); if (g_failures != 0) { diff --git a/tests/ep/CMakeLists.txt b/tests/ep/CMakeLists.txt index e505f1d8..9c6fe7ca 100644 --- a/tests/ep/CMakeLists.txt +++ b/tests/ep/CMakeLists.txt @@ -77,6 +77,7 @@ endif() add_executable(test_tilexr_ep_layout unit/test_tilexr_ep_layout.cpp ${TILEXR_ROOT}/src/ep/host/ep_layout.cpp + ${TILEXR_ROOT}/src/ep/host/ep_urma_combine_layout.cpp ) add_executable(test_tilexr_ep_api_sources @@ -85,6 +86,10 @@ add_executable(test_tilexr_ep_api_sources add_executable(test_tilexr_ep_kernel_sources unit/test_tilexr_ep_kernel_sources.cpp) +add_executable(test_tilexr_ep_start_gate_window + unit/test_tilexr_ep_start_gate_window.cpp +) + add_executable(test_tilexr_ep_host_validation unit/test_tilexr_ep_host_validation.cpp ${TILEXR_ROOT}/src/ep/host/ep_layout.cpp @@ -95,6 +100,10 @@ target_include_directories(test_tilexr_ep_layout PRIVATE ${TILEXR_EP_TEST_INCLUDE_DIRS} ) +target_include_directories(test_tilexr_ep_start_gate_window PRIVATE + ${TILEXR_EP_TEST_INCLUDE_DIRS} +) + target_include_directories(test_tilexr_ep_host_validation PRIVATE ${TILEXR_EP_TEST_INCLUDE_DIRS} ) @@ -105,12 +114,14 @@ target_compile_definitions(test_tilexr_ep_kernel_sources PRIVATE TILEXR_SOURCE_R add_test(NAME test_tilexr_ep_layout COMMAND test_tilexr_ep_layout) add_test(NAME test_tilexr_ep_api_sources COMMAND test_tilexr_ep_api_sources) add_test(NAME test_tilexr_ep_kernel_sources COMMAND test_tilexr_ep_kernel_sources) +add_test(NAME test_tilexr_ep_start_gate_window COMMAND test_tilexr_ep_start_gate_window) add_test(NAME test_tilexr_ep_host_validation COMMAND test_tilexr_ep_host_validation) install(TARGETS test_tilexr_ep_layout test_tilexr_ep_api_sources test_tilexr_ep_kernel_sources + test_tilexr_ep_start_gate_window test_tilexr_ep_host_validation RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} ) diff --git a/tests/ep/unit/test_tilexr_ep_api_sources.cpp b/tests/ep/unit/test_tilexr_ep_api_sources.cpp index 7f942804..eb161cf9 100644 --- a/tests/ep/unit/test_tilexr_ep_api_sources.cpp +++ b/tests/ep/unit/test_tilexr_ep_api_sources.cpp @@ -91,6 +91,10 @@ void TestBuildPlacement() if (ReadFile("CMakeLists.txt", &rootCmake)) { CheckContains("CMakeLists.txt", rootCmake, "option(TILEXR_BUILD_EP \"Build TileXR EP communication library\" OFF)"); + CheckContains("CMakeLists.txt", rootCmake, + "option(TILEXR_EP_BUILD_URMA_COMBINE \"Build the Ascend950 URMA EP combine operator\" OFF)"); + CheckContains("CMakeLists.txt", rootCmake, + "TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT \"22\""); CheckContains("CMakeLists.txt", rootCmake, "add_subdirectory(src/ep)"); } @@ -99,6 +103,8 @@ void TestBuildPlacement() CheckContains("src/ep/CMakeLists.txt", epCmake, "add_library(tilexr-ep SHARED"); CheckContains("src/ep/CMakeLists.txt", epCmake, "tile-comm"); CheckContains("src/ep/CMakeLists.txt", epCmake, "libtilexr_ep_combine_kernel.so"); + CheckContains("src/ep/CMakeLists.txt", epCmake, + "TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT=${TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT}"); CheckContains("src/ep/CMakeLists.txt", epCmake, "tilexr_ep.h"); CheckContains("src/ep/CMakeLists.txt", epCmake, "install(TARGETS tilexr-ep"); } @@ -257,6 +263,35 @@ void TestDispatchDemoUsesHostBarrierBeforeValidation() CheckContains("tests/ep/demo/tilexr_ep_dispatch_demo.cpp", demo, "dispatch synchronized"); } +void TestUrmaCombineApiAndBuildIntegration() +{ + const std::string apiPath = "src/include/tilexr_ep.h"; + std::string api; + if (ReadFile(apiPath, &api)) { + CheckContains(apiPath, api, "TileXRMoeEpCombineUrmaGetWorkspaceSize"); + CheckContains(apiPath, api, "TileXRMoeEpCombineUrma("); + CheckContains(apiPath, api, "TileXRMoeEpCombineUrmaProfile("); + } + + const std::string hostPath = "src/ep/host/ep_urma_combine_host.cpp"; + std::string host; + if (ReadFile(hostPath, &host)) { + CheckContains(hostPath, host, "TileXREpPrepareUrmaCombineLaunchContext"); + CheckContains(hostPath, host, "launch_tilexr_ep_urma_combine_kernel"); + CheckNotContains(hostPath, host, "TILEXR_EP_URMA_OPERATOR_LAUNCH"); + } + + const std::string cmakePath = "src/ep/CMakeLists.txt"; + std::string cmake; + if (ReadFile(cmakePath, &cmake)) { + CheckContains(cmakePath, cmake, "tilexr_ep_urma_combine_kernel.cpp"); + CheckContains(cmakePath, cmake, "Building TileXR EP URMA combine kernel (Ascend 950 S22, O2)"); + CheckContains(cmakePath, cmake, "-O2"); + CheckNotContains(cmakePath, cmake, "DIAGNOSTIC"); + CheckNotContains(cmakePath, cmake, "KERNEL_OPT_LEVEL"); + } +} + void TestNoForbiddenDependencies() { const std::vector paths = { @@ -301,6 +336,7 @@ int main() TestDemoRunnerUsesLibAndLib64Paths(); TestDispatchDemoRegistersAlignedUdmaWorkspace(); TestDispatchDemoUsesHostBarrierBeforeValidation(); + TestUrmaCombineApiAndBuildIntegration(); TestNoForbiddenDependencies(); return g_failures == 0 ? 0 : 1; } diff --git a/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp b/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp index a8100561..b88c962e 100644 --- a/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp +++ b/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp @@ -417,6 +417,39 @@ void TestClearLocalWindowDoesNotPreclearSlotHeaders() " }"); } +void TestUrmaCombineUsesCleanS22ProductionPath() +{ + const std::string configPath = "src/ep/common/ep_urma_combine_config.h"; + std::string config; + if (ReadFile(configPath, &config)) { + CheckContains(configPath, config, "TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT 22"); + CheckNotContains(configPath, config, "#define TILEXR_EP_URMA_QDC_VERSION"); + CheckNotContains(configPath, config, "#define TILEXR_EP_URMA_TX_READY_"); + CheckNotContains(configPath, config, "#define TILEXR_EP_URMA_RX_READY_"); + CheckNotContains(configPath, config, "kEpUrmaCombineQdcVersion"); + CheckNotContains(configPath, config, "kEpUrmaCombineTxMetaPrefetchFull"); + CheckNotContains(configPath, config, "kEpUrmaCombineRxReadyBatchVector"); + } + + const std::string kernelPath = "src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp"; + std::string kernel; + if (ReadFile(kernelPath, &kernel)) { + CheckContains(kernelPath, kernel, "const bool reverse = (lane & 1) != 0;"); + CheckContains(kernelPath, kernel, "WaitDeferredRoundCredit"); + CheckContains(kernelPath, kernel, "StartParallelRoundPublish"); + CheckContains(kernelPath, kernel, "FinishParallelRoundPublish"); + CheckNotContains(kernelPath, kernel, "TILEXR_EP_URMA_DIAGNOSTIC_"); + CheckNotContains(kernelPath, kernel, "TILEXR_EP_URMA_"); + CheckNotContains(kernelPath, kernel, "\n#if"); + } + + const std::string udmaPath = "src/include/tilexr_udma.h"; + std::string udma; + if (ReadFile(udmaPath, &udma)) { + CheckNotContains(udmaPath, udma, "TILEXR_EP_URMA_DIAGNOSTIC_"); + } +} + void TestNoForbiddenDependencies() { const std::vector paths = { @@ -468,6 +501,7 @@ int main() TestKernelForwardsStaticQuantConfig(); TestKernelForwardsPerTokenDynamicQuantConfig(); TestClearLocalWindowDoesNotPreclearSlotHeaders(); + TestUrmaCombineUsesCleanS22ProductionPath(); TestNoForbiddenDependencies(); if (g_failures != 0) { std::cerr << g_failures << " TileXR EP kernel source checks failed" << std::endl; diff --git a/tests/ep/unit/test_tilexr_ep_layout.cpp b/tests/ep/unit/test_tilexr_ep_layout.cpp index 69c0e7b9..65daceb0 100644 --- a/tests/ep/unit/test_tilexr_ep_layout.cpp +++ b/tests/ep/unit/test_tilexr_ep_layout.cpp @@ -1,8 +1,11 @@ #include #include +#include #include "comm_args.h" #include "ep_layout.h" +#include "ep_urma_combine.h" +#include "ep_urma_combine_layout.h" #include "tilexr_types.h" namespace { @@ -100,6 +103,151 @@ void TestRejectsInvalidConfig() TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); } +void TestUrmaCombineWorkspaceConfig() +{ + TileXREp::EpUrmaCombineWorkspaceConfig config {}; + const int ret = TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 4, 8, 2, 3, &config); + CheckInt("URMA combine valid config", ret, TileXR::TILEXR_SUCCESS); + CheckInt64("URMA combine quant bytes", config.quantDataBytes, 32); + CheckInt64("URMA combine comm bytes", config.commBytes, 64); + CheckInt64("URMA combine blocks", config.blockCount, 1); + CheckInt64("URMA combine route stride", config.routeStride, 512); + CheckInt64("URMA combine route count", config.routeCount, 8); + CheckInt64("URMA combine rx bytes", config.rxWindowBytes, 4096); + CheckInt64("URMA combine rx0", config.rxWindowOffsets[0], 512); + CheckInt64("URMA combine rx1", config.rxWindowOffsets[1], 4608); + CheckInt64("URMA combine round0", config.roundDoneOffsets[0], 8704); + CheckInt64("URMA combine round1", config.roundDoneOffsets[1], 8832); + CheckInt64("URMA combine rx lane done", config.rxLaneDoneOffset, 8960); + CheckInt64("URMA combine sender done", config.senderDoneOffset, + config.rxLaneDoneOffset + TileXREp::kEpUrmaCombinePackLaneCount * + TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA combine round publish", config.roundPublishOffset, 13056); + CheckInt64("URMA combine round credit", config.roundCreditOffset, 13120); + CheckInt64("URMA combine start gate", config.startGateOffset, 13184); + CheckInt64("URMA combine error", config.errorStatusOffset, 13312); + CheckInt64("URMA combine fixed bytes", config.fixedBytes, 13376); + CheckInt64("URMA combine fixed control boundary", config.fixedBytes, + config.errorStatusOffset + TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA combine tx ready", config.txReadyOffset, + config.fixedBytes); + CheckInt64("URMA combine tx data", config.txDataOffset, 13824); + CheckInt64("URMA combine required bytes", config.requiredBytes, 15360); +} + +void TestUrmaCombineStartGateLayout() +{ + TileXREp::EpUrmaCombineWorkspaceConfig config {}; + CheckInt("URMA start gate config", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig( + 64, 128, 7168, 8, 16, &config), TileXR::TILEXR_SUCCESS); + CheckInt64("URMA start gate follows credit", config.startGateOffset, + config.roundCreditOffset + TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA start gate rank lines", config.errorStatusOffset, + config.startGateOffset + config.rankSize * TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA start gate required QPs", TileXREp::kEpUrmaCombineRequiredQpCount, + TileXREp::kEpUrmaCombineSendLaneCount); +} + +void TestParallelRoundPublishPeerShards() +{ + constexpr int64_t kParallelSendLanes = TileXREp::kEpUrmaCombineSendLaneCount; + const int64_t rankSizes[] = {1, 2, 8, 15, 16, 17, 22, 23, 64, 128}; + for (int64_t rankSize : rankSizes) { + for (int64_t rank = 0; rank < rankSize; ++rank) { + std::vector visits(static_cast(rankSize), 0); + int64_t publishCount = 0; + for (int64_t senderId = 0; senderId < kParallelSendLanes; ++senderId) { + for (int64_t peer = senderId; peer < rankSize; peer += kParallelSendLanes) { + if (peer == rank) { + continue; + } + ++visits[static_cast(peer)]; + ++publishCount; + } + } + CheckInt64("parallel publish excludes self", visits[static_cast(rank)], 0); + CheckInt64("parallel publish peer count", publishCount, rankSize - 1); + for (int64_t peer = 0; peer < rankSize; ++peer) { + if (peer != rank) { + CheckInt("parallel publish peer owned once", + visits[static_cast(peer)], 1); + } + } + } + } + CheckInt("parallel publish control step follows release", + static_cast(TileXREp::kEpUrmaCombinePublishDone), + static_cast(TileXREp::kEpUrmaCombineRxBufferReleased + 1)); +} + +void TestDeferredRoundCreditRxCompletionShards() +{ + const int64_t rankSizes[] = {1, 2, 8, 15, 16, 17, 22, 23, 64, 128}; + for (int64_t rankSize : rankSizes) { + const int64_t publisherCount = rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + rankSize : TileXREp::kEpUrmaCombineSendLaneCount; + std::vector visits( + static_cast(TileXREp::kEpUrmaCombinePackLaneCount), 0); + for (int64_t senderId = 0; senderId < publisherCount; ++senderId) { + for (int64_t lane = senderId; lane < TileXREp::kEpUrmaCombinePackLaneCount; + lane += publisherCount) { + ++visits[static_cast(lane)]; + } + } + CheckInt64("RX completion publisher count", publisherCount, + rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + rankSize : TileXREp::kEpUrmaCombineSendLaneCount); + for (int visitsForLane : visits) { + CheckInt("RX completion lane owned once", visitsForLane, 1); + } + } +} + +void TestUrmaCombineWorkspaceRejectsInvalidConfig() +{ + TileXREp::EpUrmaCombineWorkspaceConfig config {}; + CheckInt("URMA combine null out", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 4, 8, 2, 3, nullptr), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("URMA combine negative sends", + TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 4, 8, 2, -1, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("URMA combine oversized hidden", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig( + 2, 4, TileXREp::kEpUrmaCombineMaxHidden + 1, 2, 3, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); +} + +void TestUrmaCombineDataAsFlagPayloadBoundary() +{ + TileXREp::EpUrmaCombineWorkspaceConfig exact {}; + TileXREp::EpUrmaCombineWorkspaceConfig spill {}; + CheckInt("URMA exact payload boundary", + TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 2, 448, 1, 2, &exact), + TileXR::TILEXR_SUCCESS); + CheckInt("URMA payload spill", + TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 2, 449, 1, 2, &spill), + TileXR::TILEXR_SUCCESS); + CheckInt64("URMA exact comm bytes", exact.commBytes, 480); + CheckInt64("URMA exact block count", exact.blockCount, 1); + CheckInt64("URMA spill comm bytes", spill.commBytes, 512); + CheckInt64("URMA spill block count", spill.blockCount, 2); + CheckInt64("URMA spill route stride", spill.routeStride, 1024); +} + +void TestUrmaCombineDynamicTxDoesNotMoveRemoteRegions() +{ + TileXREp::EpUrmaCombineWorkspaceConfig small {}; + TileXREp::EpUrmaCombineWorkspaceConfig large {}; + CheckInt("URMA small dynamic tx", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(4, 8, 7168, 8, 3, &small), + TileXR::TILEXR_SUCCESS); + CheckInt("URMA large dynamic tx", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(4, 8, 7168, 8, 99, &large), + TileXR::TILEXR_SUCCESS); + CheckInt64("URMA dynamic tx keeps rx0", small.rxWindowOffsets[0], large.rxWindowOffsets[0]); + CheckInt64("URMA dynamic tx keeps rx1", small.rxWindowOffsets[1], large.rxWindowOffsets[1]); + CheckInt64("URMA dynamic tx keeps round0", small.roundDoneOffsets[0], large.roundDoneOffsets[0]); + CheckInt64("URMA dynamic tx keeps fixed bytes", small.fixedBytes, large.fixedBytes); + CheckBool("URMA dynamic tx grows required bytes", large.requiredBytes > small.requiredBytes, true); +} + } // namespace int main() @@ -108,5 +256,12 @@ int main() TestDataTypes(); TestWindowConfig(); TestRejectsInvalidConfig(); + TestUrmaCombineWorkspaceConfig(); + TestUrmaCombineStartGateLayout(); + TestParallelRoundPublishPeerShards(); + TestDeferredRoundCreditRxCompletionShards(); + TestUrmaCombineWorkspaceRejectsInvalidConfig(); + TestUrmaCombineDataAsFlagPayloadBoundary(); + TestUrmaCombineDynamicTxDoesNotMoveRemoteRegions(); return g_failures == 0 ? 0 : 1; } diff --git a/tests/ep/unit/test_tilexr_ep_start_gate_window.cpp b/tests/ep/unit/test_tilexr_ep_start_gate_window.cpp new file mode 100644 index 00000000..4c3e07a1 --- /dev/null +++ b/tests/ep/unit/test_tilexr_ep_start_gate_window.cpp @@ -0,0 +1,45 @@ +#include "ep_urma_combine_start_gate_window.h" + +#include + +namespace { + +int g_failures = 0; + +void Check(bool condition, const char *message) +{ + if (!condition) { + std::cerr << message << std::endl; + ++g_failures; + } +} + +void TestOnlyFirstLaunchAfterSynchronizationUsesGate() +{ + TileXREp::EpUrmaCombineStartGateWindow window; + Check(window.BeginLaunch(), "the first launch must use the start gate"); + Check(!window.BeginLaunch(), "the second launch in a window must skip the start gate"); + Check(!window.BeginLaunch(), "later launches in a window must skip the start gate"); + + window.StreamSynchronized(); + Check(window.BeginLaunch(), "the first launch after synchronization must use the start gate"); + Check(!window.BeginLaunch(), "only one launch after synchronization may use the start gate"); +} + +void TestWindowOneUsesGateForEveryLaunch() +{ + TileXREp::EpUrmaCombineStartGateWindow window; + for (int launch = 0; launch < 4; ++launch) { + Check(window.BeginLaunch(), "window=1 must use the start gate for every launch"); + window.StreamSynchronized(); + } +} + +} // namespace + +int main() +{ + TestOnlyFirstLaunchAfterSynchronizationUsesGate(); + TestWindowOneUsesGateForEveryLaunch(); + return g_failures == 0 ? 0 : 1; +}