[BUG][CuTeDSL] cute.copy(..., tma_desc_ptr=p) with a gmem-space p passes the Python checks, then fails in libNVVM #3671

import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync
from cutlass.cute.runtime import from_dlpack
from cutlass.utils import SmemAllocator, TensorMapManager, TensorMapUpdateMode

M, N = 64, 64


@cute.kernel
def load_tile(atom: cute.CopyAtom, mX: cute.Tensor, y: cute.Tensor, ws: cute.Tensor,
              space: cutlass.Constexpr):
    tidx, _, _ = cute.arch.thread_idx()
    smem = SmemAllocator()
    sX = smem.allocate_tensor(cutlass.Float32, cute.make_layout((M, N, 1), stride=(N, 1, M * N)), 128)
    bar = smem.allocate_array(cutlass.Int64, 1)
    gX = cute.local_tile(mX, (M, N), (None, None))
    tXsX, tXgX = cpasync.tma_partition(atom, 0, cute.make_layout(1),
                                       cute.group_modes(sX, 0, 2), cute.group_modes(gX, 0, 2))
    tm = TensorMapManager(TensorMapUpdateMode.GMEM, 128)
    desc = tm.get_tensormap_ptr(ws.iterator)  # default address_space=AddressSpace.gmem
    copy_desc = desc
    if cutlass.const_expr(space == "generic"):
        copy_desc = tm.get_tensormap_ptr(desc, cute.AddressSpace.generic)
    if tidx == 0:
        cpasync.copy_tensormap(atom, desc)
        cpasync.fence_tma_desc_release()
        cute.arch.mbarrier_init(bar, 1)
        cute.arch.mbarrier_init_fence()
    cute.arch.sync_threads()
    if tidx == 0:
        cute.arch.mbarrier_arrive_and_expect_tx(bar, M * N * 4)
    if tidx < 32:  # whole warp: the TMA copy elects one lane itself
        cpasync.fence_tma_desc_acquire(desc)
        cute.copy(atom, tXgX[(None, 0, 0)], tXsX[(None, 0)], tma_bar_ptr=bar, tma_desc_ptr=copy_desc)
    cute.arch.mbarrier_wait(bar, 0)
    for m in cutlass.range_constexpr(M):
        y[m, tidx] = sX[m, tidx, 0]


@cute.jit
def run(x: cute.Tensor, y: cute.Tensor, ws: cute.Tensor, space: cutlass.Constexpr):
    atom, mX = cpasync.make_tiled_tma_atom(
        cpasync.CopyBulkTensorTileG2SOp(), x, cute.make_layout((M, N), stride=(N, 1)), (M, N))
    load_tile(atom, mX, y, ws, space).launch(grid=(1, 1, 1), block=(N, 1, 1), smem=M * N * 4 + 1024)


print("cutlass", cutlass.__version__, "|", torch.cuda.get_device_name(), "| torch", torch.__version__)
x = torch.arange(M * N, dtype=torch.float32, device="cuda").reshape(M, N)
ws = torch.zeros(16, dtype=torch.int64, device="cuda")  # 128 B, one descriptor
fx, fws = from_dlpack(x, assumed_align=16), from_dlpack(ws, assumed_align=128)
for space in ("generic", "gmem"):
    print("tma_desc_ptr address space:", space, flush=True)
    y = torch.zeros_like(x)
    fy = from_dlpack(y, assumed_align=16)
    compiled = cute.compile(run, fx, fy, fws, space)
    compiled(fx, fy, fws)
    torch.cuda.synchronize()
    print("  y == x:", torch.equal(y, x), flush=True)