Skip to content

[compiler] Fix Rational codegen and pipeline unroll bug for dynamic s…#1057

Merged
harsh-nod merged 1 commit intoiree-org:mainfrom
harsh-nod:32x64x256
Mar 6, 2026
Merged

[compiler] Fix Rational codegen and pipeline unroll bug for dynamic s…#1057
harsh-nod merged 1 commit intoiree-org:mainfrom
harsh-nod:32x64x256

Conversation

@harsh-nod
Copy link
Collaborator

…hapes

Handle _Rational values in gen_sympy_index for Max, Min, Piecewise, and comparison operators (StrictGreaterThan, GreaterThan, Eq, Ne) that previously threw CodegenError with non-standard block sizes.

Fix pipeline unroll interaction for dynamic shapes where the kernel loop count was not divisible by the unroll factor, causing the scf.for loop to execute extra iterations that read invalid pipeline state. Thread unroll_factor through PipelinedLoop to
build_guarded_pipeline_with_remainder so pipelined_iterations is rounded to guarantee divisibility.

Extend testScaledGemmMXFP4PreshuffleB with a new block shape (32,64,256) at shape (1024,2048,16384), wave_shape and dynamic_shapes parameters.

@harsh-nod harsh-nod requested review from panditsa and xintin March 6, 2026 00:23
@harsh-nod harsh-nod force-pushed the 32x64x256 branch 2 times, most recently from 32ebfdf to 625613a Compare March 6, 2026 02:48
… dynamic shapes

Handle _Rational values in Piecewise/select expressions by resolving
them via floor division before passing to arith.select, which cannot
operate on deferred _Rational values.

Fix pipeline unroll interaction for dynamic shapes where the kernel
loop count was not divisible by the unroll factor, causing the scf.for
loop to execute extra iterations that read invalid pipeline state.
Thread unroll_factor through PipelinedLoop to
build_guarded_pipeline_with_remainder and use lcm(num_stages,
unroll_factor) as the rounding stride for pipelined_iterations.

Add (32,64,256) block shape to testScaledGemmMXFP4PreshuffleBDynamic.

Signed-off-by: Harsh Menon <harsh.menon@amd.com>
@harsh-nod harsh-nod merged commit f84fa0c into iree-org:main Mar 6, 2026
32 of 35 checks passed
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.

2 participants