Skip to content

[CuTeDSL] Support dynamic shapes in right_inverse - #3582

Open
kzos wants to merge 2 commits into
NVIDIA:mainfrom
kzos:codex/issue-3471-dynamic-right-inverse
Open

[CuTeDSL] Support dynamic shapes in right_inverse#3582
kzos wants to merge 2 commits into
NVIDIA:mainfrom
kzos:codex/issue-3471-dynamic-right-inverse

Conversation

@kzos

@kzos kzos commented Sep 4, 2026

Copy link
Copy Markdown

Summary

  • allow right_inverse to accept dynamic shapes when every stride is a fully static integer
  • retain only the maximal inverse prefix whose stride continuity is statically provable, and carry dynamic shape values through CuTe expansion and base lowering
  • preserve the existing static-shape behavior, including dynamic strides and supported scaled-basis forms
  • add verifier, expansion, full-pipeline, and cutegen regression coverage for continuous, truncated, and zero-stride cases

For (16, N):(1, 16), this returns the functionally equivalent uncoalesced (16, N):(1, 16) layout demonstrated by PyCuTe in the issue discussion. Producing the canonical 16*N:1 representation depends on dynamic coalescing, which is tracked separately in #3469. Consequently, callers that require a coalesced rank-1 inverse, including the blackwell_helpers.py TMA-shape path, still depend on #3469.

Testing

  • cmake --build build-cutlass-compiler --target check-cute -j 8 (236/236 passed)
  • cmake --build build-cutlass-compiler --target check-cutegen-unittests -j 8 (220/220 passed across 25 test binaries)
  • focused post-review MLIR suite (5/5 passed)
  • ctest --test-dir build-cutlass-compiler -R 'cg_right_inverse_test' --output-on-failure (1/1 passed)
  • git diff --check

The compiler and tests were built CPU-only against the LLVM revision pinned by cutlass_compiler/LLVM_COMMIT.

Contribution disclosure

AI assistance was used for repository research, implementation, independent review, and running the checks above.

Fixes #3471

Sync the dialect reference with the dynamic-shape contract and cover the
runtime prefix-product and scalar-order lowering paths. Also fix a dangling
comment in the neighboring left-inverse test.

Test Plan:

```bash
build-cutlass-compiler/llvm-project/bin/llvm-lit -v build-cutlass-compiler/cute_ir/test/Conversion/CuteToBase/pipeline_full.mlir build-cutlass-compiler/cute_ir/test/Conversion/CuteExpandOps/LayoutAlgebra/left_inverse.mlir build-cutlass-compiler/cute_ir/test/Conversion/CuteExpandOps/LayoutAlgebra/right_inverse.mlir build-cutlass-compiler/cute_ir/test/Dialect/Cute/LayoutAlgebra/right_inverse.mlir build-cutlass-compiler/cute_ir/test/Dialect/Cute/LayoutAlgebra/right_inverse_errors.mlir
ctest --test-dir build-cutlass-compiler -R 'cg_right_inverse_test' --output-on-failure
git diff --check
```
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] [CuTeDSL] cute.right_inverse refuses layouts with dynamic sizes

1 participant