[Fix][DSL] Fix while-loop lowering, nested control flow, and add fp32 GEMM output#232
Open
ChuanLi1101 wants to merge 3 commits intomainfrom
Open
[Fix][DSL] Fix while-loop lowering, nested control flow, and add fp32 GEMM output#232ChuanLi1101 wants to merge 3 commits intomainfrom
ChuanLi1101 wants to merge 3 commits intomainfrom
Conversation
Implement a single-query-token Flash Attention kernel targeting the decode phase of autoregressive LLM inference. Uses online softmax with warp-level xor-shuffle reductions on AMD wave64. Includes correctness and performance tests against PyTorch SDPA reference. Made-with: Cursor
… GEMM output - Fix CanonicalizeWhile: clone condition op in the before region with block arguments so the loop condition is re-evaluated each iteration instead of using stale values from outside the WhileOp (#211) - Fix CanonicalizeWhile: add carry-variable rebinding at body start (from after-region block args) and explicit scf.yield at body end so the after region has a proper terminator with updated values (#211) - Fix CanonicalizeWhile: rebind carry variables from WhileOp results after the loop so subsequent code sees the final loop-carried values - Fix _ASTREWRITE_MARKER: always visit children of generated helper functions so nested control flow (while inside if, for inside if) is still lowered by subsequent AST transformers (#210) - Add fp32 output support to preshuffle GEMM (direct epilog path): skip arith.trunc_f when out_dtype="fp32" and adjust buffer resource byte calculations for 4-byte elements - Add AST-level unit tests for the while-loop transformation and fp32 output parameter validation Made-with: Cursor
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
1 participant
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.
Summary
Fix while-loop SCF lowering ([Issue]: While loop counter update issue #211, [Issue]: _ASTREWRITE_MARKER prevents later AST rewriter passes from visiting generated function bodies #210): The
CanonicalizeWhileAST rewriter had three critical bugs: (1) thebeforeregion used original operands instead of block arguments, making the condition evaluate to a constant; (2) theafterregion had noscf.yieldterminator; (3) carry variables were never rebound from block arguments or WhileOp results. This PR clones the condition op with remapped block arguments, adds carry-variable rebinding at body start, explicitscf_while_yield_at body end, and post-loop rebinding from WhileOp results.Fix nested control flow rewriting ([Issue]: _ASTREWRITE_MARKER prevents later AST rewriter passes from visiting generated function bodies #210): The
_ASTREWRITE_MARKERcheck inTransformer.visit_FunctionDefprevented all transformers from visiting children of generated helper functions (__then_X/__else_X), so awhileorforloop nested inside anif-branch was never lowered to SCF ops.Add fp32 output to preshuffle GEMM: The direct epilog path now supports
out_dtype="fp32"by skipping thearith.trunc_f(accumulator is already f32) and adjusting buffer resource byte calculations. The cshuffle epilog path rejects fp32 with a clear error.Test plan
test_ast_rewriter.pypasses: AST-level checks for while-loop transformation, nested control flow, and fp32 parameter validationtest_preshuffle_gemm.pytests still pass (no regression)out_dtype="fp32"for correctnesswhileloop with dynamic condition