[BUG][CuTeDSL] block_dim() range [1, 1024) excludes 1024, which causes wrong-code folding

Open
  • Is something wrong?
#3,794 1 comment 1 reaction 0 assignees View on GitHub

Maintainers usually reply within 4 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
1-2 days
Newbie friendliness
48/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
cpp, python
Domain
compilers

Research direction

Start by tracing how CuTe DSL supplies range information for cute.arch.block_dim() and cute.arch.grid_dim() to LLVM; the issue does not name the implementation file or a test. Reproduce with the provided kernel using a runtime Int32(1024) block size and compare behavior with the range-information pass disabled. Done when the range includes 1024 and the equality check evaluates true.

Written by the indexing model from the issue text.

Description

bug CuTe DSL
Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
The compiler specifies the wrong half-open range for the block_dim(): [1, 1024) which excludes 1024 in the register annotations passed to LLVM. Same issue with grid_dim() as well.

This causes the compiler to fold conditionals based on incorrect range by assuming that block_dim cannot be 1024. The block size must be a runtime Int32 to reproduce. With a compile-time block size, CuTe emits .reqntid 1024, and on 4.8.0 the comparison is then folded correctly, which hides the bug.

Steps/Code to reproduce bug

import cutlass
import cutlass.cute as cute
import torch
from cutlass.cute.runtime import from_dlpack

@cute.kernel
def kernel(out: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    bdim, _, _ = cute.arch.block_dim()
    if tidx == 0:
        out[0] = bdim
        out[1] = cutlass.Int32(1) if bdim == cutlass.Int32(1024) else cutlass.Int32(0)

@cute.jit
def launch(out: cute.Tensor, block: cutlass.Int32):
    kernel(out).launch(grid=(1, 1, 1), block=(block, 1, 1))

out = torch.full((2,), -1, dtype=torch.int32, device="cuda")
t_out = from_dlpack(out)
cute.compile(launch, t_out, cutlass.Int32(1024))(t_out, cutlass.Int32(1024))
torch.cuda.synchronize()
raw, eq_max = out.tolist()
assert (raw, eq_max) == (1024, 1) # bdim should be 1024, and bdim == 1024 should be True

In the above code, the compiler assumes that bdim is in the range [1,1024) (i.e. 1024 excluded) and incorrectly deduces that bdim == cutlass.Int32(1024) is always False.

The code works by disabling use-insert-range-information-pass with:

CUTE_DSL_COMPILER_OPT=use-insert-range-information-pass=false

Expected behavior
cutedsl provides the correct range to LLVM and the condition is not folded.

Environment details (please complete the following information):

  • GPUs: NVIDIA B200 (driver 580.126.20) and NVIDIA H200 (driver 580.126.09), CUDA 13.0, Ubuntu 24.04.4
  • Python 3.12.3, torch 2.14.1+cu130
  • nvidia-cutlass-dsl 4.8.0 (with libs-cu13) and 4.7.1 (with libs-cu12)

Additional context
n/a

Dominant language
C++
Stars
10.7k
Forks
2.2k
Avg merge
3d 20h
Merged PRs (30d)
4

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/cutlass

All issues in NVIDIA/cutlass

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.