[Feat] Complete BasisAttr support in IntTupleBuilder (#574)#638
Open
jhinpan wants to merge 6 commits into
Open
[Feat] Complete BasisAttr support in IntTupleBuilder (#574)#638jhinpan wants to merge 6 commits into
jhinpan wants to merge 6 commits into
Conversation
Extend IntTupleBuilder<IntTupleAttr> to accept scaled-basis (BasisAttr / CuTe E<I>) stride leaves where the algebra is well-defined, and add a Python construction surface for them. - div(Basis, Int): divide the coefficient, keep modes (reuse intSafeDiv). - eq / ne: compare basis leaves by (coefficient, modes); a basis monomial never equals a plain integer leaf. - Remaining ops (mod, lt/le/gt/ge, min/max, shapeDiv, logical*, swizzle) stay integer-only -- basis is either ill-posed there or structurally unreachable -- but now carry precise assert messages instead of bare leaf-int asserts. - Python: fx.E(mode, *, value=1) and fx.make_basis_stride(value, modes), wired through the int-tuple builder via a __fly_basis__ marker. div(Basis, Basis) is reachable for rank>=3 identity layouts via complement; it has no quotient mode, so it is rejected with a named assert rather than miscomputing a stride. Located op-layer diagnostics are out of scope here (IntTupleBuilder carries no Location) and belong to the sibling issue ROCm#583. Tests: tests/mlir/LayoutAlgebra/basis.mlir (div + identity logical_divide), equal folding in tests/mlir/Transforms/layout_lowering.mlir, and Python surface + pipeline tests in tests/unit/test_layout_algebra.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ROCm#574) Resolve the inline review comments: - E()/make_basis_stride: validate value and modes as int32 in Python (operator.index, so NumPy integer scalars are accepted) and reject negative modes -- the IntTuple assembly format cannot round-trip a negative E<mode> (1E-1 fails to re-parse). Mirror the non-negative mode check in the public __fly_basis__ binding hook. - FlyExtension binding: gate the basis branch on a *truthy* __fly_basis__ (PyObject_IsTrue, not a strict Py_True/Py_False cast that throws on non-bool) so a falsy marker is not mistaken for a basis; read modes from any iterable, not only a list. - IntTupleUtils eq/ne: return an explicit getLeafStatic(0/1) leaf instead of materializeConstantLeaf, consistent with the sibling branches. - IntTupleUtils div(): correct the comment -- a basis divisor is reached via complement() of a rank>=2 identity layout, not logical_divide. - compositionImpl: a non-tiling divisor yields a 0-extent complement mode; assert before the % so it fails with a named message instead of a SIGFPE. Pre-existing and not basis-specific (a non-coalescible integer layout with the same divisor crashes identically). Add a rank-3 identity logical_divide regression with a valid tiler. - test_layout_algebra: normalize IR whitespace before matching; add an E() input-validation regression. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…OCm#605 review) The generic fallback in IntTupleAttrBuilder reported the rejected value's type via nb::type_name(args), but nb_type_name expects a *type* object -- passing an instance reinterprets it as a PyTypeObject and segfaults (and on Python <3.11 the __name__ lookup on the instance yields a NULL that feeds PyUnicode_FromFormat). Use Py_TYPE(args.ptr())->tp_name, which is always valid, so any non-stride object (e.g. an object exposing a falsy __fly_basis__ marker) raises a clean ValueError. Add a regression. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds first-class support for scaled-basis stride leaves (“CuTe E”) in the Python front-end and MLIR lowering/tests, including validation to avoid crashes/segfaults for invalid inputs.
Changes:
- Introduces
fx.E(...)andfx.make_basis_stride(...)for constructing basis stride leaves in Python. - Extends C++ int-tuple building and algebra utilities to handle basis leaves (notably
eq/ne/div) and improves assertion diagnostics. - Adds MLIR + Python unit tests and documentation examples for basis strides and related lowering behavior.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_layout_algebra.py | Adds unit tests for fx.E, basis stride layout IR, and invalid-input behavior. |
| tests/mlir/Transforms/layout_lowering.mlir | Adds FileCheck coverage for fly.equal folding on basis leaves. |
| tests/mlir/LayoutAlgebra/basis.mlir | New MLIR tests for basis leaves flowing through tuple ops and logical_divide. |
| python/flydsl/expr/primitive.py | Adds Python API (E, make_basis_stride) + int32 validation helper. |
| lib/Dialect/Fly/Utils/IntTupleUtils.cpp | Extends int-tuple ops with basis-aware behavior and better asserts. |
| lib/Bindings/Python/FlyExtension.cpp | Adds duck-typed basis-leaf ingestion and fixes type-name reporting to avoid segfault. |
| include/flydsl/Dialect/Fly/Utils/LayoutUtils.h | Improves assertion to prevent divide-by-zero for non-tiling divisors. |
| docs/layout_system_guide.md | Documents the new basis stride construction APIs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+145
to
+150
| assert(lhs.isLeaf() && rhs.isLeafInt() && | ||
| "div is undefined for a basis divisor; lower the identity layout first"); | ||
| if (lhs.isLeafInt()) { | ||
| return IntTupleAttr::get(lhs.getLeafAsInt() / rhs.getLeafAsInt()); | ||
| } | ||
| return IntTupleAttr::get(intSafeDiv(lhs.getLeafAsBasis(), rhs.getLeafAsInt())); |
Comment on lines
+74
to
+75
| } else if (nb::hasattr(args, "__fly_basis__") && | ||
| PyObject_IsTrue(nb::object(args.attr("__fly_basis__")).ptr()) == 1) { |
15 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #574 (the implementable core — see the scope note in the issue thread).
What
Completes
BasisAttr(scaled-basis / CuTeE<I>) leaf support inIntTupleBuilder<IntTupleAttr>, finishing what #195 started (it extended onlymul/safeDiv/ceilDiv), plus a Python construction surface.Implemented — basis is well-defined and reachable:
div(Basis, Int)→(value/k)·E: divide the coefficient, keep modes. Reuses the existingintSafeDiv(BasisAttr, IntAttr); no new scalar overload.eq/ne: compare basis leaves by(coefficient, modes); a basis monomial never equals a plain integer leaf.Kept integer-only, now with precise assert messages (
mod,lt/le/gt/ge,min/max,shapeDiv,logicalAnd/Or/Not,applySwizzle,applyCoordSwizzle): for each, basis is either ill-posed (no total order on free-module generators; bit-XOR on a symbol; symmetric integer divisibility) or structurally unreachable. Implementing formulas there would be untested dead code — per-op reasoning is in the #574 thread.Python:
fx.E(mode, *, value=1)andfx.make_basis_stride(value, modes), wired through the int-tuple builder via a__fly_basis__duck-type marker.fx.make_layout(fx.make_shape(4, 8), fx.make_stride(fx.E(0), fx.E(1)))is byte-identical tomake_identity_layout((4, 8)).Two notes
div(Basis, Basis)is reachable for rank≥3 identity layouts viacomplement(LayoutUtils.h:lastStridebecomes basis aftermul(minStride, shape), so the nextdiv(minStride, lastStride)is basis-by-basis). It has no quotient mode, so it's rejected with a named assert rather than miscomputing a stride.emitOpErroris out of scope here.IntTupleBuilder<IntTupleAttr>carries only anMLIRContext*, noLocation, so located diagnostics aren't reachable at this layer — asserts are the deliberate invariant (theIntTupleValueAdaptortwin asserts too). The user-facing op-layer gate belongs to the sibling issue [Compiler]: Adding verification to the layout-algebra ops — what's the right approach? #583. This PR and [Compiler]: Adding verification to the layout-algebra ops — what's the right approach? #583 touch disjoint files (IntTupleUtils.cppasserts vsFlyOps.cppinference).Verification
bash scripts/build.shwith assertions live — clean.tests/mlirFileCheck tests pass, incl. the newLayoutAlgebra/basis.mlir(div + identitylogical_divideinference) andequalfolding appended toTransforms/layout_lowering.mlir.tests/unit/test_layout_algebra.py: 24 passed / 1 pre-existing skip, incl. the newfx.Esurface and a fullconvert-fly-to-rocdlpipeline test; broader unit sweep 398 passed, 0 failed.🤖 Generated with Claude Code