[QST][CuteDSL] Deadlock when using TMA loads data

[LRlr239]

What is your question?

kernel always hang when using TMA to copy data, ref doc &code: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#using-tma-to-transfer-multi-dimensional-arrays

demo code in cute dsl:

import torch

import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.torch as cutlass_torch
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack

@cute.struct
class SharedStorage:
    barriers: cute.struct.MemRange[cutlass.Int64, 2]

@cute.kernel
def demo_kernel(
    src: cute.Tensor,
    g2s_copy_atom: cute.CopyAtom,
    g2s_tma_coord: cute.Tensor,
    s2g_copy_atom: cute.CopyAtom,
    s2g_tma_coord: cute.Tensor,
    smem_layout: cute.Layout,

    rows_per_block: cutlass.Constexpr,
    TPB: cutlass.Constexpr,
):
    allocator = cutlass.utils.SmemAllocator()
    bar_storage = allocator.allocate(
        SharedStorage, byte_alignment=8
    ).barriers.data_ptr()
    smem = allocator.allocate_tensor(
        src.element_type,
        smem_layout,
        byte_alignment=128,
    )

    warp_idx = cute.arch.warp_idx()
    warp_idx = cute.arch.make_warp_uniform(warp_idx)
    bid_x, bid_y, _ = cute.arch.block_idx()
    tid_x, _, _ = cute.arch.thread_idx()

    if warp_idx == 0:
        cpasync.prefetch_descriptor(g2s_copy_atom)
    
    if tid_x == 0:
        cute.arch.mbarrier_init(bar_storage, cnt=TPB)
        cute.arch.mbarrier_init_fence()
    
    cute.arch.sync_threads()

    if tid_x == 0:
        g_in = cute.local_tile(
            g2s_tma_coord, (rows_per_block, TPB), (bid_x, bid_y), (1, 1)
        )
        st, gt = cpasync.tma_partition(
            atom=g2s_copy_atom,
            cta_coord=0,
            cta_layout=cute.make_layout(1),
            smem_tensor=smem,
            gmem_tensor=g_in,
        )
        print(f"st: {st}")
        print(f"gt: {gt}")
        cute.copy(
            g2s_copy_atom,
            gt,
            st,
            tma_bar_ptr=bar_storage
        )
        print(f"g_in: {g_in}")        
        cute.arch.mbarrier_arrive_and_expect_tx(
            bar_storage,
            cute.size_in_bytes(src.element_type, smem_layout)
        )
    else:
        cute.arch.mbarrier_arrive(bar_storage)
    # cute.printf("tid {} before sync\n", tid_x)
    cute.arch.mbarrier_wait(bar_storage, 0)
    
    cute.arch.sync_threads()
    if tid_x == 0:
        cute.print_tensor(smem)


@cute.jit
def demo_host_func(src: cute.Tensor, dst: cute.Tensor, rows_per_block: cutlass.Constexpr,  TPB: cutlass.Constexpr):
    print(f"src: {src}")
    smem_layout = cute.make_ordered_layout((rows_per_block, TPB), order=(1, 0))
    g2s_tma_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp()
    s2g_tma_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp()
    g2s_atom, g2s_tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom(
        g2s_tma_op,
        src,
        smem_layout,
        (rows_per_block, TPB)
    )
    s2g_atom, s2g_tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom(
        s2g_tma_op,
        src,
        smem_layout,
        (rows_per_block, TPB)
    )
    print(f"op: {g2s_tma_op}")
    print(f"g2s_atom: {g2s_atom}")
    print(f"s2g_atom: {s2g_atom}")
    print(f"g2s_tma_tensor: {g2s_tma_tensor}")
    # cute.print_tensor(g2s_tma_tensor)

    grid = (1,1,1)
    block = (TPB, 1, 1)

    demo_kernel(
        src,
        g2s_atom, g2s_tma_tensor,
        s2g_atom, s2g_tma_tensor,
        smem_layout,
        rows_per_block,
        TPB
    ).launch(
        grid=grid,
        block=block
    )

[LRlr239]

if tid_x == 0:
    cute.arch.mbarrier_init(bar_storage, cnt=1)
    cute.arch.mbarrier_init_fence()

if tid_x == 0:
    ...
    cute.arch.mbarrier_arrive(bar_storage)
    cute.arch.mbarrier_expect_tx(
        bar_storage,
        cute.size_in_bytes(src.element_type, smem_layout)
    )
    cute.copy(
        g2s_copy_atom,
        gt,
        st,
        tma_bar_ptr=bar_storage
    )

#else:
#    cute.arch.mbarrier_arrive(bar_storage)
# cute.printf("tid {} before sync\n", tid_x)
cute.arch.mbarrier_wait(bar_storage, 0)

This way won't work either.

[LRlr239]

    if warp_idx == 0:
        if tid_x == 0:
            cute.arch.mbarrier_expect_tx(
                bar_storage,
                cute.size_in_bytes(src.element_type, smem_layout)
            )
            cute.arch.mbarrier_arrive(bar_storage)

        cute.copy(
            g2s_copy_atom,
            tAgA,
            tAsA,
            tma_bar_ptr=bar_storage
        )

I found it works: operate barrier with warp0.tid0, then call cute.copy with warp0，

[Aneureka]

I think the "issue" is due to the fact that, in CuTe world, by design we want to keep consistency at the algorithm level (cute.copy in this case). That is, we do not expect, in some cases, cute.copy to be wrapped with something like if idx_x == 0, and in the else case, it is not. Instead, from cute.copy's perspective, the copy behavior is handled by the whole atom, without any implementation details such as how many threads are involved in this copy.

Following this design, during the lowering process of the tiled TMA atom, we will emit an NVVM::ElectSyncOp to select the leader lane to perform the TMA load/store (should be verifiable with the dumped IR/PTX code). This is an instruction with the .sync qualifier, so if we wrap it with if tid_x == 0:, we instead encounter a deadlock. :-)

[LRlr239]

I think the "issue" is due to the fact that, in CuTe world, by design we want to keep consistency at the algorithm level (cute.copy in this case). That is, we do not expect, in some cases, cute.copy to be wrapped with something like if idx_x == 0, and in the else case, it is not. Instead, from cute.copy's perspective, the copy behavior is handled by the whole atom, without any implementation details such as how many threads are involved in this copy.

Following this design, during the lowering process of the tiled TMA atom, we will emit an NVVM::ElectSyncOp to select the leader lane to perform the TMA load/store (should be verifiable with the dumped IR/PTX code). This is an instruction with the .sync qualifier, so if we wrap it with if tid_x == 0:, we instead encounter a deadlock. :-)

Thanks for the reply. but according to the PTX documentation, a TMA copy is issued by a single thread, and the C++ version of cute::tma_copy is invoked in a single-threaded manner， ( https://github.com/NVIDIA/cutlass/blob/main/examples/cute/tutorial/blackwell/03_mma_tma_multicast_sm100.cu#L347 ), but the python cute copy is wrapped as a warp‑level call, which internally performs an ‘elect‑one-thread’ operation. The documentation lacks an explanation of this intermediate step， and where can I find more details about it? thanks

[Aneureka]

according to the PTX documentation, a TMA copy is issued by a single thread

Correct, but again, it's the implementation details of TMA copy atom -- cute.copy algorithm should not concern about this.

the C++ version of cute::tma_copy is invoked in a single-threaded manner

Thanks for pointing out; it looks like we do have some inconsistency between C++ and DSL..

The documentation lacks an explanation of this intermediate step， and where can I find more details about it?

You may refer to the doc, which mentions "For Copy Atoms requiring single-threaded execution, thread election is managed automatically by the copy operation. External thread selection mechanisms are not necessary", but I'd agree that it's far from being sufficient and friendly 😢

[LRlr239]

according to the PTX documentation, a TMA copy is issued by a single thread

Correct, but again, it's the implementation details of TMA copy atom -- cute.copy algorithm should not concern about this.

the C++ version of cute::tma_copy is invoked in a single-threaded manner

Thanks for pointing out; it looks like we do have some inconsistency between C++ and DSL..

The documentation lacks an explanation of this intermediate step， and where can I find more details about it?

You may refer to the doc, which mentions "For Copy Atoms requiring single-threaded execution, thread election is managed automatically by the copy operation. External thread selection mechanisms are not necessary", but I'd agree that it's far from being sufficient and friendly 😢

Thanks, I can indeed see that description in the documentation.