NVIDIA / NVIDIA/cutlass

[BUG] Explicit 2D block indexing breaks predicate simplification compared to implicit 1D indexing

Open
#3,353 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug CuTe DSL inactive-30d
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
CuteDSL fails to simplify and consolidate symbolic expressions for out-of-bound predicates when a 2D block tensor is explicitly indexed via a tuple coordinate (bidy, bidx).

  • Using a 1D coordinate bidx (after transposing the thread block mapping in row-major order) allows the compiler to successfully consolidate predicates.

  • Even though the explicit 2D coordinate (bidy, bidx) is semantically and mathematically identical to the 1D variant, it generates scattered, element-wise predicates instead.

This will lead to branch divergence and performance degradation.

Steps/Code to reproduce bug

import cutlass
import cutlass.cute as cute
from cutlass.cute import KeepPTX

# =====================================================================
# Variant 1: 1D Indexing
# =====================================================================
@cute.kernel
def add_one_copy_atom_pred_kernel_1d(
    x: cute.Tensor, crd: cute.Tensor, shape: cute.Shape,
    thr_layout: cute.Layout, val_layout: cute.Layout
):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    blkx = x[(None, None), bidx]
    blkCrd = crd[(None, None), bidx]

    copy_atom = cute.make_copy_atom(
        cute.nvgpu.CopyUniversalOp(), x.element_type
    )
    tile_copy: cute.TiledCopy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
    thr_copy = tile_copy.get_slice(tidx)
    thrx = thr_copy.partition_S(blkx)
    thrCrd = thr_copy.partition_S(blkCrd)
    frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
    for i in range(cute.size(frgPred)):
        val = cute.elem_less(thrCrd[i], shape)
        frgPred[i] = val

    frgx = cute.make_fragment_like(thrx)
    cute.copy(copy_atom, thrx, frgx, pred=frgPred)
    
    result = frgx.load() + cutlass.BFloat16(1)
    frgx.store(result)
    
    cute.copy(copy_atom, frgx, thrx, pred=frgPred)
    
@cute.jit
def add_one_copy_atom_pred_1d(x: cute.Tensor):
    tile_m = 8
    tile_n = 256
    num_threads_per_block = 256

    elems_per_load = 128 // x.element_type.width
    threads_per_row = tile_n // elems_per_load
    rows_per_load = num_threads_per_block // threads_per_row
    iters = tile_m // rows_per_load
    thr_layout = cute.make_ordered_layout((rows_per_load, threads_per_row), order=(1, 0))
    val_layout = cute.make_ordered_layout((iters, elems_per_load), order=(1, 0))
    tiler_mn, _ = cute.make_layout_tv(thr_layout, val_layout)
    
    gx = cute.zipped_divide(x, tiler_mn)
    crd = cute.make_identity_tensor(x.shape)
    gcrd = cute.zipped_divide(crd, tiler_mn)

    remap_block = cute.make_ordered_layout(
        cute.select(gx.shape[1], mode=[1, 0]), order=(1, 0)
    )
    gx = cute.composition(gx, (None, remap_block))
    gcrd = cute.composition(gcrd, (None, remap_block))
    
    add_one_copy_atom_pred_kernel_1d(gx, gcrd, x.shape, thr_layout, val_layout).launch(
        grid=(cute.size(gx, mode=[1]), 1, 1),
        block=(num_threads_per_block, 1, 1)
    )

# =====================================================================
# Variant 2: Explicit 2D Indexing
# =====================================================================
@cute.kernel
def add_one_copy_atom_pred_kernel_2d(
    x: cute.Tensor, crd: cute.Tensor, shape: cute.Shape,
    thr_layout: cute.Layout, val_layout: cute.Layout
):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()
    blkx = x[(None, None), (bidy, bidx)]
    blkCrd = crd[(None, None), (bidy, bidx)]

    copy_atom = cute.make_copy_atom(
        cute.nvgpu.CopyUniversalOp(), x.element_type
    )
    tile_copy: cute.TiledCopy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
    thr_copy = tile_copy.get_slice(tidx)
    thrx = thr_copy.partition_S(blkx)
    thrCrd = thr_copy.partition_S(blkCrd)
    frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
    for i in range(cute.size(frgPred)):
        val = cute.elem_less(thrCrd[i], shape)
        frgPred[i] = val

    frgx = cute.make_fragment_like(thrx)
    cute.copy(copy_atom, thrx, frgx, pred=frgPred)
    
    result = frgx.load() + cutlass.BFloat16(1)
    frgx.store(result)
    
    cute.copy(copy_atom, frgx, thrx, pred=frgPred)
    
@cute.jit
def add_one_copy_atom_pred_2d(x: cute.Tensor):
    tile_m = 8
    tile_n = 256
    num_threads_per_block = 256

    elems_per_load = 128 // x.element_type.width
    threads_per_row = tile_n // elems_per_load
    rows_per_load = num_threads_per_block // threads_per_row
    iters = tile_m // rows_per_load
    thr_layout = cute.make_ordered_layout((rows_per_load, threads_per_row), order=(1, 0))
    val_layout = cute.make_ordered_layout((iters, elems_per_load), order=(1, 0))
    tiler_mn, _ = cute.make_layout_tv(thr_layout, val_layout)
    
    gx = cute.zipped_divide(x, tiler_mn)
    crd = cute.make_identity_tensor(x.shape)
    gcrd = cute.zipped_divide(crd, tiler_mn)
    
    add_one_copy_atom_pred_kernel_2d(gx, gcrd, x.shape, thr_layout, val_layout).launch(
        grid=(gx.shape[1][1], gx.shape[1][0], 1),
        block=(num_threads_per_block, 1, 1)
    )

import torch
from cutlass.cute.runtime import from_dlpack
x = torch.randn(16384, 16384, device="cuda", dtype=torch.bfloat16)
x_ = from_dlpack(x, assumed_align=16)
cute.compile[KeepPTX](add_one_copy_atom_pred_1d, x_)
cute.compile[KeepPTX](add_one_copy_atom_pred_2d, x_)

The issue can be clearly illustrated by the generated PTX codes, or by benchmarking two kernel. It turns out that the 2D variant is significantly slower.

Expected behavior
Both indexing variants should produce equivalent, fully simplified predicates for out-of-bound checks since their block level layouts are identical.

Environment details (please complete the following information):

  • GPU: RTX 5090
  • CUTLASS: 4.5.2
  • CUDA: 13.1.2

Additional context

Here is the comparison of the generated PTX. Note how the 1D variant successfully consolidates predicates, while the 2D variant generates divergent instructions.

1D variant (all predications are identical) 
...
setp.gt.u32 	%p1, %r1, 16383;
	@%p1 bra 	$L__BB0_2;
	ld.global.b16 	%rs1, [%rd1];
	mov.b32 	%r89, {%rs1, %rs2};
$L__BB0_2:
	setp.gt.u32 	%p2, %r1, 16383;
	@%p2 bra 	$L__BB0_4;
	ld.global.b16 	%rs3, [%rd1+2];
	{ .reg .b16 tmp; mov.b32 {%rs4, tmp}, %r89; }
	mov.b32 	%r89, {%rs4, %rs3};
$L__BB0_4:
	setp.gt.u32 	%p3, %r1, 16383;
	@%p3 bra 	$L__BB0_6;
	ld.global.b16 	%rs5, [%rd1+4];
	{ .reg .b16 tmp; mov.b32 {tmp, %rs6}, %r15; }
	mov.b32 	%r93, {%rs5, %rs6};
...
2D variant (predications differ, leading to warp divergence) 
...
setp.gt.u32 	%p7, %r81, 16383;
	setp.lt.u32 	%p8, %r81, 16384;
	setp.gt.s32 	%p9, %r82, 16383;
	setp.lt.s32 	%p10, %r82, 16383;
	and.pred  	%p1, %p8, %p10;
	setp.lt.s32 	%p11, %r82, 16382;
	and.pred  	%p2, %p8, %p11;
	setp.lt.s32 	%p12, %r82, 16380;
	and.pred  	%p3, %p8, %p12;
	setp.lt.s32 	%p13, %r82, 16379;
	and.pred  	%p4, %p8, %p13;
	setp.lt.s32 	%p14, %r82, 16378;
	and.pred  	%p5, %p8, %p14;
	setp.lt.s32 	%p15, %r82, 16377;
	and.pred  	%p6, %p8, %p15;
	or.pred  	%p16, %p7, %p9;
	@%p16 bra 	$L__BB0_2;
	ld.global.b16 	%rs9, [%rd1];
	mov.b32 	%r119, {%rs9, %rs10};
$L__BB0_2:
	not.pred 	%p17, %p1;
	@%p17 bra 	$L__BB0_4;
	ld.global.b16 	%rs11, [%rd1+2];
	{ .reg .b16 tmp; mov.b32 {%rs12, tmp}, %r119; }
	mov.b32 	%r119, {%rs12, %rs11};
$L__BB0_4:
	not.pred 	%p18, %p2;
	@%p18 bra 	$L__BB0_6;
	ld.global.b16 	%rs13, [%rd1+4];
	{ .reg .b16 tmp; mov.b32 {tmp, %rs14}, %r14; }
	mov.b32 	%r123, {%rs13, %rs14};
...
Similar kernels of the two variants, only without predication, will generate identical PTX codes as expected, proving the issue lies strictly in the predication logic. 
@cute.kernel
def add_one_copy_atom_kernel_1d(
    x: cute.Tensor,
    thr_layout: cute.Layout, val_layout: cute.Layout
):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    blkx = x[(None, None), bidx]

    copy_atom = cute.make_copy_atom(
        cute.nvgpu.CopyUniversalOp(), x.element_type
    )
    tile_copy: cute.TiledCopy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
    thr_copy = tile_copy.get_slice(tidx)
    thrx = thr_copy.partition_S(blkx)

    frgx = cute.make_fragment_like(thrx)
    cute.copy(copy_atom, thrx, frgx)
    
    result = frgx.load() + cutlass.BFloat16(1)
    frgx.store(result)
    
    cute.copy(copy_atom, frgx, thrx)
    
@cute.jit
def add_one_copy_atom_1d(x: cute.Tensor):
    # ... (same as previous kernel)
    
    gx = cute.zipped_divide(x, tiler_mn)

    remap_block = cute.make_ordered_layout(
        cute.select(gx.shape[1], mode=[1, 0]), order=(1, 0)
    )
    gx = cute.composition(gx, (None, remap_block))
    
    add_one_copy_atom_kernel_1d(gx, thr_layout, val_layout).launch(
        grid=(cute.size(gx, mode=[1]), 1, 1),
        block=(num_threads_per_block, 1, 1)
    )

@cute.kernel
def add_one_copy_atom_kernel_2d(
    x: cute.Tensor,
    thr_layout: cute.Layout, val_layout: cute.Layout
):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()
    blkx = x[(None, None), (bidy, bidx)]

    copy_atom = cute.make_copy_atom(
        cute.nvgpu.CopyUniversalOp(), x.element_type
    )
    tile_copy: cute.TiledCopy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
    thr_copy = tile_copy.get_slice(tidx)
    thrx = thr_copy.partition_S(blkx)

    frgx = cute.make_fragment_like(thrx)
    cute.copy(copy_atom, thrx, frgx)
    
    result = frgx.load() + cutlass.BFloat16(1)
    frgx.store(result)
    
    cute.copy(copy_atom, frgx, thrx)
    
@cute.jit
def add_one_copy_atom_2d(x: cute.Tensor):
    # ... (same as previous kernel)
    
    gx = cute.zipped_divide(x, tiler_mn)
    
    add_one_copy_atom_kernel_2d(gx, thr_layout, val_layout).launch(
        grid=(gx.shape[1][1], gx.shape[1][0], 1),
        block=(num_threads_per_block, 1, 1)
    )

Contributor guide

No contributing guide indexed for this repository

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.

Research direction

Start by compiling add_one_copy_atom_pred_1d and add_one_copy_atom_pred_2d with KeepPTX, then compare the generated predicates shown in the issue. Trace the CuTe DSL predicate simplification and layout-indexing paths involved in the two kernel entry points. Done means equivalent consolidated predicates for both variants, with the reproduced performance difference removed.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.