Skip to content

[Enh] Improve type closure for primitive func#552

Open
sjfeng1999 wants to merge 22 commits into
mainfrom
pr/enh-type-closure
Open

[Enh] Improve type closure for primitive func#552
sjfeng1999 wants to merge 22 commits into
mainfrom
pr/enh-type-closure

Conversation

@sjfeng1999

@sjfeng1999 sjfeng1999 commented May 21, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Copilot AI review requested due to automatic review settings May 21, 2026 10:43
@sjfeng1999

sjfeng1999 commented May 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

may need update the same aiter kernels simultaneously.

Copilot AI reviewed May 21, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR updates FlyDSL’s Python expression-layer wrappers to better preserve (“close over”) DSL types (e.g., Numeric, Vector) when calling primitive ops, and adjusts affected tests/kernels to use the new return-value behavior.

Changes:

  • Wrap primitive-op scalar results back into Numeric types (and propagate this through helpers like get_scalar, get_leaves, ptr_load, memref_load).
  • Centralize memref_load_vec to return a Vector with shape/dtype metadata, and simplify Tensor.load() accordingly.
  • Update unit tests and a few kernel utilities to align with the revised scalar/vector return types and pipeline string formatting.

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_static_vs_dynamic.py Adjusts dynamic-layout tests to return i32 scalars directly and simplifies pipeline string formatting.
tests/unit/test_layout_algebra.py Updates dynamic test functions to return i32 scalars from fx.get_scalar(...).ir_value() and simplifies pipeline string formatting.
python/flydsl/expr/typing.py Updates IntTuple reconstruction and makes Tensor.load() rely on the updated memref_load_vec wrapper.
python/flydsl/expr/primitive.py Introduces numeric re-wrapping helper and applies it across several primitive ops; moves vector wrapping into memref_load_vec.
python/flydsl/expr/math.py Extends traced math-op wrapping to preserve DSL closure for both Numeric and Vector inputs.
kernels/silu_and_mul_fq.py Simplifies scale-offset computation by relying on the updated fx.get_scalar behavior.
kernels/mfma_preshuffle_pipeline.py Updates crd2idx helper to unwrap int-tuples and cast to index type using the new scalar typing behavior.
kernels/layout_utils.py Updates dynamic-layout crd2idx fallback to unwrap/cast through fx.get_scalar(...).ir_value().

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread kernels/mfma_preshuffle_pipeline.py
Comment thread python/flydsl/expr/primitive.py Outdated
@sjfeng1999 sjfeng1999 changed the title [Enh] Ensure type closure for primitive func [Enh] Improve type closure for primitive func May 25, 2026
@coderfeli

coderfeli commented May 27, 2026

Copy link
Copy Markdown
Collaborator

CI failed. @sjfeng1999

@sjfeng1999 sjfeng1999 force-pushed the pr/enh-type-closure branch from 205ec38 to abb308a Compare June 2, 2026 08:56
@sjfeng1999 sjfeng1999 requested a review from Copilot June 2, 2026 08:59
Copilot AI reviewed Jun 2, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 25 out of 25 changed files in this pull request and generated no new comments.

xudoyuan and others added 12 commits June 5, 2026 03:39
@xudoyuan
[Refactor]: rm IR value
8813c9a
@xudoyuan
[Bugfix]: import __all__
a22c8ec
@xudoyuan
Merge branch 'main' into pr/enh-type-closure
e1c82ce
@xudoyuan @claude
Fix DSL type compatibility in layout_utils.idx2crd
9587ccd 
Unwrap DSL types (Int32, etc.) to raw ir.Value at the entry of
idx2crd, so that downstream arith ops (shrui, andi) receive
proper ir.Value operands instead of Numeric wrappers.

Also replace fragile hasattr/str-based type check with isinstance.

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
@xudoyuan @claude
Replace arith.constant(index=True) with fx.Int32() in mixed_moe_gemm_…
0750258 
…2stage

Eliminate all index-typed arith.constant and arith.index calls in
mixed_moe_gemm_2stage.py, replacing them with fx.Int32() DSL types.
This prevents type mismatch errors (i32 vs index) when DSL-typed
values (from fx.get, fx.idx2crd) participate in arithmetic with
index-typed constants.

174 occurrences replaced.

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
@xudoyuan
Merge branch 'main' into pr/enh-type-closure
8874319
@xudoyuan @claude
Format mixed_moe_gemm_2stage.py with black
c957349 
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
@xudoyuan @claude
Restore index type for scf.ForOp arguments in mixed_moe_gemm_2stage
234240c 
scf.ForOp requires index-typed ir.Value for lower_bound, upper_bound,
step, and iter_args. These were incorrectly converted to fx.Int32() in
the previous commit. Restore them to arith.constant(..., index=True).

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
@xudoyuan
Revert "Restore index type for scf.ForOp arguments in mixed_moe_gemm_…
83469fe 
…2stage"

This reverts commit 234240c.
@xudoyuan
Revert "Format mixed_moe_gemm_2stage.py with black"
9cdd537 
This reverts commit c957349.
@xudoyuan
Revert "Replace arith.constant(index=True) with fx.Int32() in mixed_m…
9ef6433 
…oe_gemm_2stage"

This reverts commit 0750258.
@xudoyuan @claude
Revert global index→Int32 replacement, fix type mismatch in framework
c6ebbe2 
Revert the bulk arith.constant(index=True)→fx.Int32() replacement in
mixed_moe_gemm_2stage.py — the file mixes index and i32 type chains
that cannot be uniformly converted.

Instead, fix the root cause in _try_coerce_rhs (numeric.py): when
Numeric.__mul__ encounters an index-typed ArithValue as the rhs
operand, cast it to Int32 so the operands have matching types. This
resolves the arith.muli(i32, index) type mismatch without changing
any kernel code.

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes

@jhinpan jhinpan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found two issues that still apply on the current head. The earlier mixed_moe index/i32 arithmetic concern appears fixed in this version: sorted_m is back to an index constant path.

jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread python/flydsl/expr/__init__.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread python/flydsl/expr/arith.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread kernels/mixed_moe_gemm_2stage.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread python/flydsl/expr/numeric.py Outdated
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread kernels/mixed_moe_gemm_2stage.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread kernels/gemm_fp8fp4_gfx1250.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread python/flydsl/expr/math.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread python/flydsl/expr/primitive.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread python/flydsl/expr/primitive.py
jhinpan
jhinpan reviewed Jun 8, 2026
View reviewed changes
Comment thread python/flydsl/expr/numeric.py Outdated
xudoyuan and others added 6 commits June 8, 2026 10:49
@xudoyuan @claude
Fix Index(Index) nesting in Index.__init__ and coerce index-typed Ari…
46ed1ca 
…thValue in _try_coerce_rhs

Index.__init__ skipped unwrapping when x was already an Index instance,
causing Index(Index(ArithValue)) nesting. Fix by taking .value directly
when x is Index, avoiding the nested DSL type.

Also add IndexType handling in _try_coerce_rhs: cast index-typed
ArithValue to Int32 so Numeric binary ops can process mixed i32/index
operands without falling through to ArithValue.__rmul__.

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
@xudoyuan @claude
Fix _try_coerce_rhs: wrap index ArithValue as Index instead of Int32
212c837 
Wrapping index-typed ArithValue as Int32 caused index_cast (index→i32),
changing the IR type of the entire downstream computation chain. This
broke MLIR ops that require index-typed operands (vector.load, scf.for).

Use Index(rhs) instead to preserve the index type.

Also fix Index.__init__ to handle Index(Index) by taking .value directly,
preventing nested DSL types.

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
@sjfeng1999
minor fix
94ae353
@sjfeng1999
re-export select
b2e2af6
@xudoyuan
[Bugfix]: Replace / with //
271fecd
@xudoyuan
[Bugfix]: Replace / with /
14b3bdc
aoli26
aoli26 reviewed Jun 10, 2026
View reviewed changes

@aoli26 aoli26 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot Copilot left review comments

@xudoyuan xudoyuan xudoyuan left review comments

@aoli26 aoli26 aoli26 approved these changes

+1 more reviewer

@jhinpan jhinpan jhinpan left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@sjfeng1999 @coderfeli @jhinpan @xudoyuan @aoli26