[BUG] Call `cute.printf` on swizzled tensor causes crash

[melonedo]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
cute.printf("{}", a_tensor_with_swizzled_layout) makes the whole script crash without any error message.

Steps/Code to reproduce bug
Run this snippet with nvidia-cutlass-dsl==4.2.1

import cutlass
import cutlass.cute as cute
import torch

# Can not remove this, maybe it is used to initialize CUDA?
torch.empty(1, device='cuda')

@cute.kernel
def main():
    sA_layout = cute.make_composed_layout(
        cute.make_swizzle(3, 3, 3), 0, cute.make_layout(1024)
    )
    smem = cutlass.utils.SmemAllocator()
    sA = smem.allocate_tensor(cutlass.Float16, sA_layout, 16)

    cute.printf("{}", sA)


@cute.jit
def launcher():
    main().launch(grid=(1, 1, 1), block=(1, 1, 1))


launcher()

The call cute.printf("{}", sA) makes the program crash with message: “python cuda/print-tensor-bug.py” terminated by signal SIGABRT (Abort).

Expected behavior
Should print the tensor, just like the unswizzled case. If that is not supported, at least crash with an error message.

Environment details (please complete the following information):

Ubuntu 20.04 with RTX 3090.

[fengxie]

Thanks for reporting. @brandon-yujie-sun for visibility.