Skip literals that are not scalars or iotas when fusing attention - #5211
Open
ahsan-ca wants to merge 2 commits into
Open
Skip literals that are not scalars or iotas when fusing attention#5211ahsan-ca wants to merge 2 commits into
ahsan-ca wants to merge 2 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR refines the fuse_attention pass’s constant-capture behavior so attention fusion only pulls in “inlinable” literal constants (scalars or iota/range literals), avoiding problematic large/non-iota literals while still enabling MLIR causal-mask detection.
Changes:
- Add a literal classifier in
src/fuse_attention.cppto restrict which literals get pulled into the attention subgraph during fusion. - Update existing attention-fusion tests to pass relevant literals as group inputs (instead of recreating them inside the grouped module).
- Add a new regression test covering a packed-but-nonstandard-stride bias literal that should not be inlined as an iota/range.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/fuse_attention.cpp |
Adds “range literal” detection and uses it to gate which evaluable constants are pulled into attention fusion. |
test/fuse_attention.cpp |
Updates expected fused graphs to pass literals via group inputs and adds a regression test for non-iota bias literals. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+118
to
+127
| bool is_range_literal(const literal& l) | ||
| { | ||
| bool result = false; | ||
| l.visit([&](auto x) { | ||
| result = std::adjacent_find(x.begin(), x.end(), [](auto cur, auto next) { | ||
| return not float_equal(next - cur, 1); | ||
| }) == x.end(); | ||
| }); | ||
| return result; | ||
| } |
Comment on lines
+122
to
+124
| result = std::adjacent_find(x.begin(), x.end(), [](auto cur, auto next) { | ||
| return not float_equal(next - cur, 1); | ||
| }) == x.end(); |
Comment on lines
+131
to
+133
| if(ins->name() != "@literal") | ||
| return true; | ||
| return ins->get_shape().elements() == 1 or is_range_literal(ins->get_literal()); |
Comment on lines
+361
to
+362
| for(std::size_t i = 0; i < bias_vec.size(); i++) | ||
| bias_vec[i] = 0.5f * i; |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
src/fuse_attention.cpp:126
next - curcan overflow for signed integer literals even after the monotonicity check (for example,{INT64_MIN, INT64_MAX}), causing undefined behavior while running the fusion pass. Compare againstcur + 1instead;next > curguarantees that increment is safe and also expresses the iota condition directly.
result = std::adjacent_find(x.begin(), x.end(), [](auto cur, auto next) {
return next <= cur or not float_equal(next - cur, 1);
| return 0; | ||
| } | ||
|
|
||
| bool is_range_literal(const literal& l) |
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.
Motivation
Technical Details
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot ApplicableFollow the LLVM AI Tool Use Policy for contributions using AI.