mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
cd02a30c46
* [Feature] Extend t.copy to support TMA multicast and SM-to-SM cluster copy Add multicast copy (tma_load_multicast) and shared-memory-to-shared-memory cluster copy (tma_store_cluster / ptx_cluster_store) to the t.copy operator. Multicast copy: - A single CTA issues tma_load_multicast to broadcast a tile to multiple CTAs in the cluster simultaneously; other masked CTAs receive passively. - New `cluster_mask` parameter on T.copy() controls which CTAs participate. SM-to-SM cluster copy: - Fast path via cp.async.bulk.shared::cluster with mbarrier synchronisation (new `remote_barrier` parameter on T.copy()). - Slow path via cooperative_groups::map_shared_rank for element-by-element fallback (new `dst_block` parameter on T.copy()). Supporting changes: - New builtins: tma_load_multicast, tma_store_cluster, ptx_cluster_store, cluster_sync, get_cluster_id, get_cluster_block_rank, get_cluster_block_nums, mbarrier_arrive (cluster-scope). - Codegen (codegen_cuda.cc) and device templates (barrier.h, copy_sm90.h) for all new intrinsics; unified mbarrier API via tl::mbarrier_* free functions. - inject_tma_barrier: handle tma_load_multicast; distinguish thread guards from cluster-rank conditions for correct expect_tx injection. - lower_hopper_intrin: migrate barrier allocation to builtin::create_barriers; hoist user ptx_init_barrier_thread_count calls alongside compiler barriers. - warp_specialized_rewriter: support tma_load_multicast as producer op; UserBarrierInitExtractor to separate user barrier inits from body; ragged (dynamic-extent) pipeline prefix counter; fix nested IfThenElse state corruption in shuffle-elect optimisation. - multi_version_buffer_rewriter: ragged pipeline support with runtime prefix counter for correct ping-pong buffer versioning. - thread_storage_sync: reserve fixed barrier IDs (kProducer/kConsumer) for 256-thread CTA warp-specialised split to prevent deadlocks. - pipeline_planning: defensive handling of non-BufferLoad mbarrier args. - Shared memory alignment bumped to 128 bytes for TMA. Co-authored-by: Guangda Sun <2012661711@qq.com> * [Docs] Add programming guide for Cluster TMA features Add docs/programming_guides/cluster_tma.md covering the two new T.copy extensions introduced in the t_copy_extend feature branch: - TMA multicast (cluster_mask): explains how a single TMA transaction broadcasts one global tile to multiple CTAs in a cluster simultaneously, with API usage, a per-rank behaviour table, and a complete code example. - SM-to-SM cluster copy (dst_block / remote_barrier): documents the fast path (cp.async.bulk.shared::cluster + mbarrier, single-thread async DMA) and the slow path (map_shared_rank element-by-element SIMT fallback), including the synchronisation contract for source and destination CTAs. Also covers cluster helper builtins (T.get_cluster_block_rank, T.cluster_sync, T.mbarrier_init/arrive/wait_parity, etc.) and a split-K sketch combining both features end-to-end. * change cluster_size -> cluster_dims * fix merge conflict * fix TILELANG CHECK bug * unify cluster function * fix pre-commit errors * docs(cluster_tma): clarify multicast non-issuer behavior in lowering note * fix(copy): gate dst_block cluster lowering on target capability * fix(cuda/cluster): hard-trap when cluster intrinsics are unavailable instead of assert(false)+dummy return. * fix inject_tma & lower_hopper & lower_tile bugs * [warp specialize]Fix mbarrier retarget for tma_load_multicast in warp-specialized producer path * fix(tma-barrier): correct arrive thread count under equality-guarded threadIdx * fix: overlapped TMA barrier IDs * fix: multicast checks * minor fix * minor fix * fix: tma_load_multicast() in the stage-local mbarrier rewrite * minor fix * format * minor fix * fix: Track allocation from every hoisted barrier init. * fix: Scope the 128-byte alignment to kernels that actually use TMA. * fix: Assert the tl::tma_store_cluster lowering, not just the output. * format fix * fix: Don't assume both if arms transfer the same number of bytes. * fix: Don't wrap the kept pipelined loop in a SeqStmt before the For post-processing runs. * rm mbarrier_init and use alloc_cluster_barrier * fix: Don't make cluster copy a global exemption for unrelated TMA loads. * fix: Don't drop dependency modeling for handle-based mbarrier_wait_parity(). * fix: Don't hoist extracted barrier-init if statements to the block root. * fix(copy): honor remapped dst layout in cluster-copy slow path * fix(copy): gate cluster TMA fast path on provable contiguity * fix(copy): add barrier completion for SIMT cluster-copy fallback * fix: remove dup cluster_sync * fix: remove not used _get_mbarrier * fix: remove mbarrier related code in codegen_cuda.cc and barrier.h * fix: codegen_cuda.cc: reuse functions in cluster.h and remove dependency on cg * fix testting * fix(transform): prevent ragged_prefix free var from breaking MakePackedAPI * minor fix * format * fix inject tma barrier: fix crash and correct mbarrier thread-count injection * revert inject tma * fix(inject_tma_barrier): avoid SIGSEGV and correctly infer barrier arrive thread counts * fix: add tma_load_multicast same as tma_load * T.copy_cluster * fix: renaming cluster mbarrier_arrive to ptx_arrive_cluster_barerier * fix: remove mbarrier_arrive in builtin.cc * test_tma_dsmem: fallback test * [Feature] Multi-TMA fallback for non-contiguous T.copy_cluster regions Previously, a non-contiguous copy region (e.g. a 2-D slice where the inner dim does not span the full buffer width) with a remote_barrier fell back to element-wise SIMT stores, losing TMA throughput. This commit introduces a "multi-TMA" middle path in LowerClusterCopy: the ND copy region is recursively decomposed into individual contiguous rows, each emitting one tma_store_cluster call. The mbarrier arrive_count is updated to N (the number of rows) so the destination CTA's wait() still works correctly without any API change. Key changes: - src/op/copy.cc: add static MakeTMARows() that recursively splits an ND region into contiguous rows; static extents are unrolled at compile time, symbolic extents produce TIR For loops. - src/op/operator.h: add UpdateBarrierArriveCallback (Var → PrimExpr) to LowerArgs so LowerClusterCopy can propagate the new arrive_count. - src/transform/lower_tile_op.cc: collect arrive_count overrides in barrier_arrive_updates_ and apply them to the barrier_init block annotation after visiting the block body, before LowerSharedBarrier consumes it. - testing/python/cuda/test_tma_dsmem.py: replace the old SIMT-barrier test with test_store_cluster_multi_tma_barrier (2-D, M rows) and add test_store_cluster_3d_multi_tma (3-D, D×M rows) to cover the two-level recursive decomposition. * bug fix * fix: unified continuous check * fix doc * fix: remote redundant annotation converting * fix: remove unused code * fix: remove unused code * rm unused comments * minor fix * fix: rm unused helper func * minor * minor * Clean up copy cluster merge formatting * Reduce copy cluster PR churn --------- Co-authored-by: Guangda Sun <2012661711@qq.com> Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>