[FEA] Support `grid_constant` annotation in CuTe DSL

[Gin-Sin]

Which component requires the feature?

CuTe DSL

Feature Request

Recently I've been up to megakernel, which usually involves changing TMA descriptor. Unlike grouped gemm, the TMA descriptors are probably unchanged, so it will be optimal to be stored in grid constant cache at the very beginning. However, CuTe DSL doesn't support annotating a kernel parameter as grid_constant just like the annotation __grid_constant__ in CuTe does. Typically the tensormap is constructed on the host using cuTensorMapEncodeTiled.

Essentially, what I find useful is something like this:

@cute.kernel
def kernel(
    tma_atom: cute.CopyAtom,
    @cute.grid_constant tma_desc: cute.Tensor,
):
    cute.copy(
        tma_atom,
        src,
        dst,
        tma_desc_ptr=tma_desc.iterator,
    )

[depaulmillz]

For clarification, is there something specifically limiting you from using multiple copy atoms? These will carry the grid const descriptor in them. If the atom does not match the descriptor that you are using, it could lead to some difficult to debug issues.

[Gin-Sin]

For clarification, is there something specifically limiting you from using multiple copy atoms? These will carry the grid const descriptor in them. If the atom does not match the descriptor that you are using, it could lead to some difficult to debug issues.

It's quite typical in megakernel, where I need probably tens of tensormaps. Using multiple copy atoms is feasible but ugly, because I cannot create an array of copy atoms and index it.

[depaulmillz]

In your use case the array would be known at compile time, right? In that case, you should be able to currently use lists. I was able to hack this up quickly as a reference:

import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
import cutlass.utils as utils
from cutlass.cute.nvgpu import cpasync
from typing import List


@cute.kernel
def batch_memcpy_kernel(
    tma_load_atoms: List[cute.CopyAtom],
    tma_store_atoms: List[cute.CopyAtom],
    src_tensors: List[cute.Tensor],
    dst_tensors: List[cute.Tensor],
    smem_layout: cute.Layout,
):
    """
    Kernel that performs 3 memcpy operations.

    1. TMA load: global memory -> shared memory
    2. TMA store: shared memory -> global memory
    """
    warp_idx = cute.arch.warp_idx()
    warp_idx = cute.arch.make_warp_uniform(warp_idx)
    bidx, _, _ = cute.arch.block_idx()

    # Prefetch TMA descriptors (done by first warp)
    if warp_idx == 0:
        for i in cutlass.range_constexpr(len(tma_load_atoms)):
            cpasync.prefetch_descriptor(tma_load_atoms[i])
            cpasync.prefetch_descriptor(tma_store_atoms[i])

    # Allocate shared memory
    smem = utils.SmemAllocator()
    mbar_ptr = smem.allocate_array(cutlass.Int64, 2)

    # Initialize pipeline for TMA synchronization
    producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
    consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)

    # Transaction bytes for all 3 copies
    tx_bytes = cute.size(smem_layout) * 4 * 3  # 4 bytes per float32, 3 copies

    producer, consumer = pipeline.PipelineTmaAsync.create(
        barrier_storage=mbar_ptr,
        num_stages=1,
        producer_group=producer_group,
        consumer_group=consumer_group,
        tx_count=tx_bytes,
    ).make_participants()

    # Allocate shared memory tensors for all 3 copies
    smem_tensors = [
        smem.allocate_tensor(cutlass.Float32, smem_layout, 128)
        for _ in range(len(tma_load_atoms))
    ]

    # CTA layout for single-CTA (no multicast)
    cta_coord = 0
    cta_layout = cute.make_layout(1)

    # Slice and partition tensors for this CTA
    tile_coord = ((None, None), bidx)
    
    # Partition tensors using tma_partition
    load_partitions = []
    store_partitions = []
    for i in cutlass.range_constexpr(len(tma_load_atoms)):
        src_tile = src_tensors[i][tile_coord]
        dst_tile = dst_tensors[i][tile_coord]
        
        # Partition for TMA load: gmem -> smem
        smem_load, gmem_load = cpasync.tma_partition(
            tma_load_atoms[i],
            cta_coord,
            cta_layout,
            smem_tensor=cute.group_modes(smem_tensors[i], 0, 2),
            gmem_tensor=cute.group_modes(src_tile, 0, 2),
        )
        load_partitions.append((smem_load, gmem_load))
        
        # Partition for TMA store: smem -> gmem
        smem_store, gmem_store = cpasync.tma_partition(
            tma_store_atoms[i],
            cta_coord,
            cta_layout,
            smem_tensor=cute.group_modes(smem_tensors[i], 0, 2),
            gmem_tensor=cute.group_modes(dst_tile, 0, 2),
        )
        store_partitions.append((smem_store, gmem_store))

    # === TMA LOAD: Global Memory -> Shared Memory ===
    if warp_idx == 0:
        handle = producer.acquire()
        for i in cutlass.range_constexpr(len(tma_load_atoms)):
            smem_ptn, gmem_ptn = load_partitions[i]
            cute.copy(
                tma_load_atoms[i],
                gmem_ptn,
                smem_ptn,
                tma_bar_ptr=handle.barrier,
            )
        handle.commit()

    # Wait for TMA loads to complete
    handle = consumer.wait()
    handle.release()
    cute.arch.sync_threads()

    # === TMA STORE: Shared Memory -> Global Memory ===
    if warp_idx == 0:
        for i in cutlass.range_constexpr(len(tma_store_atoms)):
            smem_ptn, gmem_ptn = store_partitions[i]
            cute.copy(tma_store_atoms[i], smem_ptn, gmem_ptn)


@cute.jit
def batch_memcpy(
    src0: cute.Tensor,
    src1: cute.Tensor,
    src2: cute.Tensor,
    dst0: cute.Tensor,
    dst1: cute.Tensor,
    dst2: cute.Tensor,
    stream,
):
    """
    Host function that sets up and launches the TMA batch memcpy kernel.

    Creates 3 TMA load atoms and 3 TMA store atoms, then passes them
    along with source/destination tensors to the kernel for batch processing.
    """
    # Tile configuration
    tile_m = 128
    tile_n = 64
    threads_per_block = 128
    tiler_mn = (tile_m, tile_n)

    # Shared memory layout
    smem_layout = cute.make_layout((tile_m, tile_n), stride=(tile_n, 1))

    # Bundle source and destination tensors
    srcs = [src0, src1, src2]
    dsts = [dst0, dst1, dst2]

    # Create lists for TMA atoms and tiled tensors
    tma_load_atoms = []
    tma_store_atoms = []
    src_tiled = []
    dst_tiled = []

    # Create TMA atoms for each of the 3 copies
    for i in cutlass.range_constexpr(len(srcs)):
        # TMA Load: Global -> Shared
        load_atom, load_tma_tensor = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileG2SOp(),
            srcs[i],
            smem_layout,
            tiler_mn,
        )
        tma_load_atoms.append(load_atom)
        src_tiled.append(cute.zipped_divide(load_tma_tensor, tiler_mn))

        # TMA Store: Shared -> Global
        store_atom, store_tma_tensor = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(),
            dsts[i],
            smem_layout,
            tiler_mn,
        )
        tma_store_atoms.append(store_atom)
        dst_tiled.append(cute.zipped_divide(store_tma_tensor, tiler_mn))

        print(f"[DSL INFO] Copy {i}:")
        print(f"[DSL INFO]   load_atom = {load_atom.type}")
        print(f"[DSL INFO]   store_atom = {store_atom.type}")

    # Calculate grid size
    num_tiles = cute.size(src_tiled[0], mode=[1])

    print("[DSL INFO] Launching batch_memcpy_kernel:")
    print(f"[DSL INFO]   tile_shape = ({tile_m}, {tile_n})")
    print("[DSL INFO]   num_tiles =", num_tiles)

    # Launch the kernel
    batch_memcpy_kernel(
        tma_load_atoms,
        tma_store_atoms,
        src_tiled,
        dst_tiled,
        smem_layout,
    ).launch(
        grid=(num_tiles, 1, 1),
        block=(threads_per_block, 1, 1),
        stream=stream,
    )


# Test and verification code
if __name__ == "__main__":
    import torch
    import cuda.bindings.driver as cuda
    from cutlass.cute.runtime import from_dlpack

    # Initialize CUDA context
    cutlass.cuda.initialize_cuda_context()

    # Test configuration
    M = 256
    N = 128

    print("=" * 60)
    print("CuTe DSL Batch Memcpy with TMA")
    print("=" * 60)
    print(f"Copying 3 tensors of shape ({M}, {N}) each")
    print()

    # Get current CUDA stream from PyTorch
    torch_stream = torch.cuda.current_stream()
    current_stream = cuda.CUstream(torch_stream.cuda_stream)

    # Create source and destination tensors
    src0_torch = torch.randn(M, N, dtype=torch.float32, device="cuda")
    src1_torch = torch.randn(M, N, dtype=torch.float32, device="cuda")
    src2_torch = torch.randn(M, N, dtype=torch.float32, device="cuda")

    dst0_torch = torch.zeros(M, N, dtype=torch.float32, device="cuda")
    dst1_torch = torch.zeros(M, N, dtype=torch.float32, device="cuda")
    dst2_torch = torch.zeros(M, N, dtype=torch.float32, device="cuda")

    # Convert to CuTe tensors
    src0 = from_dlpack(src0_torch, assumed_align=16).mark_layout_dynamic()
    src1 = from_dlpack(src1_torch, assumed_align=16).mark_layout_dynamic()
    src2 = from_dlpack(src2_torch, assumed_align=16).mark_layout_dynamic()
    dst0 = from_dlpack(dst0_torch, assumed_align=16).mark_layout_dynamic()
    dst1 = from_dlpack(dst1_torch, assumed_align=16).mark_layout_dynamic()
    dst2 = from_dlpack(dst2_torch, assumed_align=16).mark_layout_dynamic()

    # Compile and run
    print("Compiling TMA kernel...")
    compiled_fn = cute.compile(
        batch_memcpy,
        src0, src1, src2,
        dst0, dst1, dst2,
        current_stream,
    )

    print("Running batch memcpy kernel...")
    compiled_fn(
        src0, src1, src2,
        dst0, dst1, dst2,
        current_stream,
    )

    torch.cuda.synchronize()

    # Verify results
    print()
    print("Verifying results...")
    assert torch.allclose(src0_torch, dst0_torch), "Memcpy 0 failed!"
    assert torch.allclose(src1_torch, dst1_torch), "Memcpy 1 failed!"
    assert torch.allclose(src2_torch, dst2_torch), "Memcpy 2 failed!"

    print("PASS")

[depaulmillz]

In general, we are looking into exposure of grid constant, but it will take time.

[Gin-Sin]

Thanks for your explanation!

However, a List in CuTe DSL must be indexed with a compile-time determined value, but typically I will need to use a run-time calculated value as its index, using something like tile_coord_mnl to determine the index into the copy atom array.

FYI, I have also tried using tidx to index List but failed.

[github-actions]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.