[BUG] CuTe DSL seems to discard `setmaxnreg` hints

[Gin-Sin]

Which component has the problem?

CuTe DSL

Bug Report

CUDA Version: 12.9

CuTe DSL Version: 4.3.0.dev0

In hopper/dense_gemm_persistent.py, I tried to add another warp group but the following errors appeared:

error: unknown: NVPTX compiler invocation failed, error log: ptxas fatal   : (C7602) Insufficient registers (128) to compile instruction at line 1582 in function kernel_cutlass_kernel___main__HopperWgmmaGemmPersistentKernel_object_at__CopyAtom_ThrID10_TVLayoutSrc1819201_TVLayoutDst1819201_Valuetypef16_tensor000odiv16111012_CopyAtom_ThrID10_TVLayou_0. Try to compile with register target of 154 or higher.

The errors suggest that each thread in the MMA warp group has only 128 registers, but I have used setmaxnreg to configure the number of register files for each warp group as the following:

if warp_group_idx == 0:
    cute.arch.warpgroup_reg_dealloc(24)

if warp_group_idx == 1:
    cute.arch.warpgroup_reg_dealloc(24)

if warp_group_idx in (2, 3):  # Two warp groups: 2 & 3
    cute.arch.warpgroup_reg_alloc(224)

I have used 63488 register files in total. It seemed that CuTe DSL directly uses the average number of register files when compiling instead of considering setmaxnreg hints.

It doesn't seem to be a coincidence that in the above errors CuTe DSL said Insufficient registers (128), which happened to be 65536//(128*4)==128, and it said Try to compile with register target of 154 or higher, which happened to be satisfied when only 3 WGs is used, to wit, 65536//(128*3) == 170.67 > 154. That's why I think CuTe DSL directly uses the average number of register files per thread when compiling, discarding setmaxnreg hints.

A minimal reproduction can be run like this:

python3 dense_gemm_persistent_four_wgs.py --mnkl 8192,4096,4096,1 --tile_shape_mn 128,256 --cluster_shape_mn 2,1 --a_dtype Float16 --b_dtype Float16 --c_dtype Float16 --acc_dtype Float32 --a_major k --b_major n --c_major n --warmup_iterations 0 --iterations 1 2>&1 | tee dense_gemm_4_wgs.log

[shubaoyu2]

the requirement when using setmaxnreg is that the program should be compile-able with the number of registers threads are born with

[thakkarV]

Note, this is not a CuTe DSL limitation either, this is the PTX programming model for this construct.

[YouJiacheng]

Note, this is not a CuTe DSL limitation either, this is the PTX programming model for this construct.

It's more like a limitation of ptxas? It can recognize setmaxnreg and decide not to spill, but can't use the register count set by setmaxnreg to determine whether there are sufficient registers to compile an instruction.

[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.