[BUG] CuteDSL compiler hangs when calling copy with mismatched predicates
Open
Nobody has claimed this yet.
? - Needs Triage
bug
inactive-30d
inactive-90d
- Dominant language
- C++
- Stars
- 10.5k
- Forks
- 2.1k
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 7
Description
The compiler hangs with this code, where we call copy with the wrong copy atom / predicates.
The compiler still should error out instead of hanging.
Steps/Code to reproduce bug
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@cute.kernel
def hang_kernel(
gX: cute.Tensor,
cX: cute.Tensor, # coordinate tensor
shape: cute.Shape,
tv_layout: cute.Layout,
tiler_mn: cute.Shape,
):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
blkCrd = cX[(None, None), 0]
blkX = gX[(None, None), 0]
copy_atom_load_bf16 = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16)
copy_atom_load_fp32 = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float32)
thr_copy_bf16 = cute.make_tiled_copy(copy_atom_load_bf16, tv_layout, tiler_mn).get_slice(tidx)
thr_copy_fp32 = cute.make_tiled_copy(copy_atom_load_fp32, tv_layout, tiler_mn).get_slice(tidx)
thrFp32 = thr_copy_fp32.partition_S(blkX)
thrBf16 = thr_copy_bf16.partition_S(blkX)
frgX = cute.make_fragment_like(thrFp32)
thrCrd = thr_copy_bf16.partition_S(blkCrd)
frgPred = cute.make_fragment(frgX.shape, cutlass.Boolean)
for i in range(cute.size(frgPred)):
frgPred[i] = cute.elem_less(thrCrd[i], shape)
cute.copy(copy_atom_load_bf16, thrFp32, frgX, pred=frgPred)
# Does not hang if we copy without the predicate, or we call copy with copy_atom_load_fp32
@cute.jit
def hang(mX):
N = mX.shape[-1]
warpsize = 32
num_blocks_N = cute.ceil_div(N, warpsize)
tiler_mn = (1, N)
tv_layout = cute.make_layout(
((warpsize, 1), (1, num_blocks_N)),
stride=((1, 1), (1, warpsize))
)
print(f"[DSL INFO] Input Tensors:")
print(f"[DSL INFO] mX = {mX.type}")
print(f"[DSL INFO] Tiling Parameters:")
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
print(f"[DSL INFO] tv_layout = {tv_layout}")
mX_expanded_layout = cute.prepend(mX.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
mX_expanded = cute.make_tensor(mX.iterator, mX_expanded_layout)
gX = cute.zipped_divide(mX_expanded, tiler_mn) # ((TileM,TileN),(RestM,RestN))
print(f"[DSL INFO] Tiled Tensors:")
print(f"[DSL INFO] gX = {gX.type}")
shape = (4, N)
idX = cute.make_identity_tensor(shape)
cX = cute.zipped_divide(idX, tiler=tiler_mn)
print(f"[DSL INFO] coord tensor = {cX.type}")
hang_kernel(gX, cX, shape, tv_layout, tiler_mn).launch(
grid=[1, 1, 1],
block=[cute.size(tv_layout, mode=[0]), 1, 1],
)
def run(N):
if not torch.cuda.is_available():
raise RuntimeError(f"Ampere GPU is required to run this example!")
print(f"Tensor dimensions: [{N}]")
x = torch.randn(N, device=torch.device("cuda"), dtype=torch.float32)
print(f"Input tensor shapes:")
print(f"x: {x.shape}, dtype: {x.dtype}")
x_tensor = from_dlpack(x, assumed_align=16)
print("Compiling kernel with cute.compile ...")
compiled_func = cute.compile(hang, x_tensor)
print("Done compiling...")
if __name__ == "__main__":
run(1024)
Cc @thakkarV
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Run the supplied reproducer through the hang entry point with N=1024, focusing on hang_kernel and its cute.copy call with the predicate. Trace the CuteDSL compiler path for mismatched copy atoms and predicates; done means the compiler reports an error instead of hanging, while the unpredicated and matching-atom cases remain unaffected.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100