Fix attention mask traversal for three nested selects

[umangyadav]

Motivation

Attention graphs can combine prefix-causal, KV-cache, and sliding-window masks as three nested select(mask, -inf, scores) operations. The matcher previously inspected only the outer select and one chained select, so the third recognized mask remained in the QK elementwise region instead of being represented by rock.attention attributes.

That incomplete folding prevents the attention graph from reaching code generation and can surface with this error signature:

error: Cannot trace first gemm index for linalg.generic op
Lowering failed.

This PR is stacked on #2436 because it relies on that PR's sliding-window sequence-length detection and clip preservation.

Summary

• Iteratively peel up to all three recognized nested masks while preserving unknown or duplicate masks in the elementwise region.
• Add a TOSA-to-Rock regression test that checks prefix offset, causal, KV-cache sequence length, and sliding-window attributes and verifies no tosa.select remains.
• Add a nightly fusion E2E test with active prefix-causal, KV-cache, and sliding-window masks, including output and LSE verification.

Test plan

• tosa-to-rock-attention-three-mask.mlir
• mixr-attention-sliding-window-kvcache-prefix-causal.mlir
• git clang-format --diff origin/users/umayadav/fix-sliding-window-mask-detection

Made with Cursor

[rocmlir-pr-reviewer]

Verdict: APPROVE -- submitted as COMMENT (automated reviews are advisory)  ·  Findings: 0 (0 Critical, 0 Major, 0 Minor)

---

Scope

Fixes TOSA-to-Rock attention mask folding so that three separately-nested select(mask, -inf, scores) ops (prefix-causal, KV-cache, sliding-window) are all recognized and represented as rock.attention attributes, instead of leaving the third mask in the QK elementwise region (which blocked codegen). Touches one C++ file plus two new tests.

Findings

No blocking issues found.

Notes

• TosaToRock.cpp:2629-2649: the new peeling loop is well-behaved. foundCount is bounded to [0,3]; every non-break iteration strictly increases it and the loop condition requires < 3, so termination is guaranteed. The foundCount == before check correctly leaves an unrecognized or duplicate mask in the elementwise region by not advancing inputToContinue. The > 0 gate preserves the prior "at least one mask required" precondition, so this is a clean generalization of the old single-step logic rather than a behavior change for the 1–2 mask cases.
• The seqLen reconciliation block (:2651-2669) is correctly kept after the loop, so the result stays independent of select nesting order.
• Test coverage is appropriate: tosa-to-rock-attention-three-mask.mlir asserts all four attributes plus CHECK-NOT: tosa.select inside the elementwise region, and the nightly E2E adds output + LSE verification. Both satisfy the Major "new optimization needs a FileCheck test" and fusion-coverage bullets.
• Minor (non-blocking, not raised inline): foundCount(currentResult) is evaluated a few times per iteration; hoisting into a local would be marginally tidier but is not a checklist item.

CI status

No non-self CI failures. The review pipeline's own checks and the Copilot reviewer are still in progress, which is expected and not a CI failure.

[Copilot]

Pull request overview

This PR improves the TOSA-to-Rock attention matcher so it can fully fold attention graphs that apply three nested select(mask, -inf, scores) masks (prefix-causal, KV-cache, sliding-window), enabling downstream codegen and avoiding lowering failures.

Changes:

• Update TosaToRock attention matching to iteratively peel and analyze up to three nested select(..., -inf, ...) masks.
• Add a TOSA-to-Rock regression test ensuring all three masks are folded into a single rock.attention (with no remaining tosa.select in the elementwise region).
• Add a nightly fusion E2E test covering the three-mask combination and verifying runtime correctness (output + LSE).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
mlir/lib/Conversion/TosaToRock/TosaToRock.cpp	Iteratively traverses chained select masks to detect KV-cache, prefix-causal, and sliding-window masks across three nest levels.
mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-three-mask.mlir	New lit regression test asserting all three masks fold into rock.attention attributes and no tosa.select remains in the QK elementwise region.
mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir	New nightly E2E test exercising the three-mask case and checking both folding and execution results.

---

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

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2437      +/-   ##
===========================================
+ Coverage    82.57%   83.71%   +1.14%     
===========================================
  Files          120      121       +1     
  Lines        42852    43181     +329     
  Branches      7110     7180      +70     
===========================================
+ Hits         35381    36145     +764     
+ Misses        4815     4482     -333     
+ Partials      2656     2554     -102     

Flag	Coverage Δ
gfx120x	83.50% <100.00%> (+0.97%)	⬆️
gfx950	83.39% <100.00%> (+1.04%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
mlir/lib/Conversion/TosaToRock/TosaToRock.cpp	79.58% <100.00%> (+1.16%)	⬆️

... and 38 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[justinrosner]

Not necessarily urgent for this PR, but in rocmlirTriton (assuming we are going to port this change there) we may run into errors with the duplicate/unrecognized masks being left in the elementwise region. The regularize* passes cannot sink transforms through non-splat constants (I've seen errors with this before when we were first testing the FusionZoo kernels with the rocmlirTriton prototype)

[justinrosner]

I'm confused here. Doesn't this same block of logic already exist below? Why do we need the same thing twice?

[justinrosner]

Is this just a merge conflict artifact from updating the branch before?

[umangyadav]

Yes, Fixed it