[BUG] [CuTeDSL] cute.coalesce doesn't do anything on dynamic shapes
Nobody has claimed this yet.
- 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
- 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
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