NVIDIA / NVIDIA/cutlass

[BUG] [CuTeDSL] cute.coalesce doesn't do anything on dynamic shapes

Open
#3,469 11 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug CuTe DSL
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
A mode with a dynamic shape doesn't get coalesced into previous modes, even though the stride is static and the previous modes are fully static too (and that's all that's needed to check whether the modes can be coalesced).

Steps/Code to reproduce bug

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


@cute.jit
def repro(tensor: cute.Tensor):
    layout = cute.make_layout(
        (16, tensor.shape[1] // 128),
        stride=(8, 128),
    )
    expected = cute.make_layout(
        (16 * (tensor.shape[1] // 128),),
        stride=(8,),
    )

    print("input:    ", layout)
    print("coalesced:", cute.coalesce(layout))
    print("expected: ", expected)


tensor = from_dlpack(torch.empty((4096, 8192), device="cuda"))
tensor = tensor.mark_compact_shape_dynamic(mode=1)
repro(tensor)

Output:

input:     (16,?):(8,128)
coalesced: (16,?):(8,128)
expected:  (?{div=16}):(8)

Environment details (please complete the following information):
Version 4.7.0 of CuTeDSL

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 with cute.coalesce and reproduce the issue using the repro function and dynamic tensor shape shown in the report. Compare the current coalesced layout with the expected layout, then verify that the dynamic mode is coalesced when its stride and preceding modes allow it.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.