[Feat] Complete BasisAttr support in IntTupleBuilder (#574)

[jhinpan]

Restored PR: this is a recreation of #605. The previous PR was closed accidentally when the old head repository was deleted, so GitHub lost the head repo association. This branch is restored at the same old head commit b3c06cfb5e1de02e9fc78a23febef07d438ad343.

---

Closes #574 (the implementable core — see the scope note in the issue thread).

What

Completes BasisAttr (scaled-basis / CuTe E<I>) leaf support in IntTupleBuilder<IntTupleAttr>, finishing what #195 started (it extended only mul/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 existing intSafeDiv(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) and fx.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 to make_identity_layout((4, 8)).

Two notes

• div(Basis, Basis) is reachable for rank≥3 identity layouts via complement (LayoutUtils.h: lastStride becomes basis after mul(minStride, shape), so the next div(minStride, lastStride) is basis-by-basis). It has no quotient mode, so it's rejected with a named assert rather than miscomputing a stride.
• emitOpError is out of scope here. IntTupleBuilder<IntTupleAttr> carries only an MLIRContext*, no Location, so located diagnostics aren't reachable at this layer — asserts are the deliberate invariant (the IntTupleValueAdaptor twin 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.cpp asserts vs FlyOps.cpp inference).

Verification

• bash scripts/build.sh with assertions live — clean.
• All 30 tests/mlir FileCheck tests pass, incl. the new LayoutAlgebra/basis.mlir (div + identity logical_divide inference) and equal folding appended to Transforms/layout_lowering.mlir.
• tests/unit/test_layout_algebra.py: 24 passed / 1 pre-existing skip, incl. the new fx.E surface and a full convert-fly-to-rocdl pipeline test; broader unit sweep 398 passed, 0 failed.
• clang-format (LLVM) + black/ruff clean.

🤖 Generated with Claude Code

[Copilot]

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(...) and fx.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.

[jhinpan]

Good catch — fixed. The precondition now gates the dividend explicitly to isLeafInt() || isLeafBasis() (with rhs.isLeafInt()), so getLeafAsBasis() is only reached on the basis branch and a none/other leaf trips the named assert instead of crashing:

assert((lhs.isLeafInt() || lhs.isLeafBasis()) && 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()));

[jhinpan]

Fixed. Extracted an isMarkerTruthy helper that propagates the -1 case as a nb::python_error (which captures the pending Python exception) instead of treating it as false:

static bool isMarkerTruthy(nb::handle obj) {
  int truth = PyObject_IsTrue(obj.ptr());
  if (truth < 0) {
    throw nb::python_error();
  }
  return truth == 1;
}

So a __fly_basis__ that raises in its getter now surfaces the real error rather than silently falling through to the generic int path.

[jhinpan]

Rebased onto current main and folded in the review.

Conflict. The restored branch sat on the old head's pre-squash history, so GitHub had no common ancestor with main and marked it CONFLICTING. Re-applied the change as a single commit on top of main (5d4f7727) — it's now a clean 8-file diff and MERGEABLE. The two core files (IntTupleUtils.cpp, LayoutUtils.h) had no upstream drift; FlyExtension.cpp/primitive.py did, and the basis additions slot in cleanly against the current dsl_loc_tracing surface.

Review fixes (both Copilot points, replied inline):

• div: precondition tightened to isLeafInt() || isLeafBasis() so getLeafAsBasis() is unreachable on a none leaf.
• FlyExtension: marker truthiness now propagates a -1 PyObject_IsTrue as nb::python_error instead of swallowing it.

Verify (assertions live): scripts/build.sh clean, 30/30 tests/mlir FileCheck (incl. LayoutAlgebra/basis.mlir), tests/unit/test_layout_algebra.py 26 passed / 1 pre-existing skip, test_expr_optional_rocdl green (basis surface stays target-neutral). clang-format-18 + black/ruff clean.

@sjfeng1999 this is the implementable core of #574 — div(Basis,Int)/eq/ne get real formulas, the rest stay int-only with sharpened asserts (rationale per-op in the #574 thread), and emitOpError is deferred to #583 since IntTupleBuilder carries no Location. Would appreciate your read on the scope call and the assert-vs-diagnostic boundary when you have a moment.