[QST] TiledMMA.make_fragment_A fails with non-swizzled smem layout

Description

When using TiledMMA (GMMA instructions on Hopper) with a plain non-swizzled shared memory layout (e.g., created via cute.make_ordered_layout), calling thr_mma.make_fragment_A(tCsA) raises an MLIR legalization error:

failed to legalize operation 'cute_nvgpu.make_gmma_smem_desc' that was explicitly marked illegal

However, if the shared memory tensor is allocated with a swizzled layout (e.g., SmemLayoutAtomKind.K_SW128), the code compiles and runs correctly.

Reproduction

Env

GPU: Hopper (SM90)
Python 3.11.14
torch==2.8.0
nvidia-cutlass-dsl==4.3.2
nvcc==cuda_12.9.r12.9/compiler.35813241_0

Minimal reproducible example

import torch
import cutlass
from cutlass import const_expr, Float16, Float32, Int32, Int64, Boolean, Int8
import cutlass.cute as cute
import cutlass.utils.hopper_helpers as sm90_utils
from cutlass.cute.nvgpu.warpgroup import (
    OperandMajorMode,
    make_smem_layout_atom,
)

@cute.kernel
def mykernel(
    tiled_mma: cute.TiledMma
):
    tid, _, _ = cute.arch.thread_idx()
    smem = cutlass.utils.SmemAllocator()
    # smem_layout_atom = make_smem_layout_atom(
    #     cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128,
    #     Float16,
    # )
    # layout = cute.tile_to_shape(smem_layout_atom, (64,128), order=(0,1))
    # sA = smem.allocate_tensor(Float16, layout.outer, swizzle=layout.inner)
    # sB = smem.allocate_tensor(Float16, layout.outer, swizzle=layout.inner)
    sA = smem.allocate_tensor(Float16, cute.make_ordered_layout((64,128), order=(1,0)))
    sB = smem.allocate_tensor(Float16, cute.make_ordered_layout((64,128), order=(1,0)))
    thr_mma = tiled_mma.get_slice(tid)
    tCsA = thr_mma.partition_A(sA)
    tCsB = thr_mma.partition_B(sB)
    print(f"sA={sA} tCsA={tCsA}")
    tCrA = thr_mma.make_fragment_A(tCsA)
    tCrB = thr_mma.make_fragment_B(tCsB)
    tCrC = thr_mma.make_fragment_C(thr_mma.partition_shape_C((64, 64)))
    cute.gemm(tiled_mma, tCrC, tCrA, tCrB, tCrC)

@cute.jit
def launcher():
    tiled_mma = sm90_utils.make_trivial_tiled_mma(
        Float16,
        Float16,
        OperandMajorMode.K,
        OperandMajorMode.K,
        Float32,
        atom_layout_mnk=(1,1,1),
        tiler_mn=(64, 64),
    )
    mykernel(tiled_mma).launch(
        grid=[1,1,1],
        block=[128, 1, 1],
        cluster=[1, 1, 1],
    )

launcher()

Error message

base_dsl/compiler.py", line 186, in compile_and_jit
    self.compile(

base_dsl/compiler.py", line 148, in compile
    pm.run(module.operation)

cutlass._mlir._mlir_libs._site_initialize.<locals>.MLIRError: Failure while executing pass pipeline:
error: "tCrA = thr_mma.make_fragment_A(tCsA)"(...): failed to legalize operation 'cute_nvgpu.make_gmma_smem_desc' that was explicitly marked illegal
 note: see current operation: %25 = "cute_nvgpu.make_gmma_smem_desc"(%23) 
       <{layout = #cute.layout<"((64,16),1,8):((128,1),0,16)">, major = #cute_nvgpu.major<k>}> 
       : (!cute.ptr<f16, smem, align<1024>>) -> !cute_nvgpu.smem_desc

Questions

1. Is using a swizzled smem layout a hard requirement for GMMA TiledMMA in the Python DSL? If so, should the DSL raise an explicit, user-friendly error earlier (e.g., at allocate_tensor or make_fragment_A call time)?
2. Is there any supported path to use GMMA with a non-swizzled layout (e.g., for debugging or correctness testing purposes)?

Thanks for any help!

Start with the minimal Hopper reproduction and trace thr_mma.make_fragment_A through the cute_nvgpu.make_gmma_smem_desc legalization failure shown in base_dsl/compiler.py. Compare the non-swizzled and K_SW128 allocation paths, then determine whether non-swizzled layouts are supported and what behavior or diagnostic should result. Done means the supported-layout behavior and any required user-facing error are established and tested.