[BUG][CuTeDSL] cute.copy(..., tma_desc_ptr=TensorMap.get_ptr()) is silently ignored on tiled TMA loads #3670

import hashlib, re, sys
import torch, cutlass, cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync
from cutlass.cute.runtime import from_dlpack
from cutlass.experimental.cuda import TensorMap, create_tensor_map_tiled_from_view
from cutlass.utils import SmemAllocator, TensorMapManager, TensorMapUpdateMode
from cutlass._mlir import ir
from cutlass._mlir.dialects import llvm

M, N = 64, 64


@cute.kernel
def load_tile(atom: cute.CopyAtom, mX: cute.Tensor, y: cute.Tensor,
              desc: cutlass.GridConstant[TensorMap], mode: 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))
    if tidx == 0:
        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 (elect.sync)
        if cutlass.const_expr(mode == "baked"):
            cute.copy(atom, tXgX[(None, 0, 0)], tXsX[(None, 0)], tma_bar_ptr=bar)
        elif cutlass.const_expr(mode == "get_ptr"):
            p = desc.get_ptr()
            print(f"  get_ptr() -> {type(p).__module__}.{type(p).__qualname__}, "
                  f"isinstance(p, cute.Pointer) = {isinstance(p, cute.Pointer)}")
            cute.copy(atom, tXgX[(None, 0, 0)], tXsX[(None, 0)], tma_bar_ptr=bar, tma_desc_ptr=p)
        else:
            # same address as get_ptr(), re-typed as a generic cute pointer to a tensormap; the manager
            # asserts its second argument as the pointer alignment, and the descriptor param is .align 64 in PTX
            addr = cutlass.Int64(llvm.ptrtoint(ir.IntegerType.get_signless(64), desc.get_ptr()))
            p = TensorMapManager(TensorMapUpdateMode.GMEM, 64).get_tensormap_ptr(
                cute.make_ptr(cutlass.Int64, addr, cute.AddressSpace.gmem, assumed_align=64),
                cute.AddressSpace.generic)
            cute.copy(atom, tXgX[(None, 0, 0)], tXsX[(None, 0)], tma_bar_ptr=bar, tma_desc_ptr=p)
    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, z: cute.Tensor, y: cute.Tensor, mode: cutlass.Constexpr):
    atom, mX = cpasync.make_tiled_tma_atom(
        cpasync.CopyBulkTensorTileG2SOp(), x, cute.make_layout((M, N), stride=(N, 1)), (M, N))
    desc = create_tensor_map_tiled_from_view(z, box_dims=(M, N))  # same box as the atom, base = z
    load_tile(atom, mX, y, desc, mode).launch(grid=(1, 1, 1), block=(N, 1, 1), smem=M * N * 4 + 1024)


def tma_lines(ptx):
    return [l.strip() for l in ptx.splitlines() if "cp.async.bulk.tensor" in l]


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)
z = -x - 1
fx, fz = from_dlpack(x, assumed_align=16), from_dlpack(z, assumed_align=16)
bodies = {}
for mode in ("baked", "get_ptr", "retyped"):
    print(f"[{mode}]")
    y = torch.zeros_like(x)
    fy = from_dlpack(y, assumed_align=16)
    compiled = cute.compile(run, fx, fz, fy, mode, options="--keep-ptx")
    compiled(fx, fz, fy)
    torch.cuda.synchronize()
    ptx = compiled.__ptx__ or ""
    body = re.sub(r"\w*load_tile\w*", "K", ptx)
    bodies[mode] = hashlib.md5(body.encode()).hexdigest()[:12]
    for l in tma_lines(ptx):
        print("  ptx:", l)
    print(f"  y == x: {torch.equal(y, x)}   y == z: {torch.equal(y, z)}   ptx md5 (name-normalised): {bodies[mode]}")
print("get_ptr PTX identical to baked PTX:", bodies["get_ptr"] == bodies["baked"])