From 868969d283f6c5083e9b23d33353493b6ce7a74c Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Thu, 23 Jul 2026 17:31:38 +0000 Subject: [PATCH 1/4] Fix sliding-window sequence length mask detection Validate and reconcile sequence lengths so unsupported masks are not folded and clipping semantics are preserved. Co-authored-by: Cursor --- mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 85 +++++++++++- ...-to-rock-attention-sliding-window-neg.mlir | 131 ++++++++++++++++++ 2 files changed, 209 insertions(+), 7 deletions(-) create mode 100644 mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp index e271892efadb..af7b6163eb68 100644 --- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp +++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp @@ -1976,6 +1976,17 @@ struct AttentionRewritePattern : public OpRewritePattern { isa(maybeBlockArg.value()); } + // Returns true when both values resolve to the same currentSeqLen block + // argument after skipping reshape/broadcast ops. + bool sameSeqLenBlockArg(Value a, Value b, + const DenseSet &seqLenSkip) const { + FailureOr resolvedA = getValueSkipping(a, seqLenSkip); + FailureOr resolvedB = getValueSkipping(b, seqLenSkip); + return succeeded(resolvedA) && succeeded(resolvedB) && + isa(resolvedA.value()) && + resolvedA.value() == resolvedB.value(); + } + // Helper function to detect select-based causal mask pattern: // - true branch is a splat -inf constant // - false branch is the tensor value that we want to return @@ -2104,6 +2115,11 @@ struct AttentionRewritePattern : public OpRewritePattern { // Clip bounds detected on currentSeqLen during KV-cache pattern matching. std::optional seqLenClipMin; std::optional seqLenClipMax; + // The currentSeqLen block argument and clip referenced by the + // sliding-window mask. + Value slidingWindowSeqLen; + std::optional slidingWindowClipMin; + std::optional slidingWindowClipMax; }; // Helper to try detecting prefix causal pattern: add(row_indices, offset) @@ -2281,15 +2297,28 @@ struct AttentionRewritePattern : public OpRewritePattern { return ClipBounds{*clipMin, *clipMax}; } + // Result of sliding-window pattern detection. + struct SlidingWindowResult { + int64_t windowSize; + Value seqLen; + std::optional clipMin; + std::optional clipMax; + }; + // Helper to try detecting sliding window pattern: // greater(add(seqLen, negative_const_offset) * broadcast, col_indices) - // Returns the window size if successful. - FailureOr + // Returns the window size and validated currentSeqLen operand if successful. + FailureOr trySlidingWindowPattern(Value input, const DenseSet &seqLenSkip) const { DenseSet expandAndCollapse{ tensor::CollapseShapeOp::getOperationName(), tensor::ExpandShapeOp::getOperationName()}; + DenseSet expandCollapseMinMax{ + tensor::CollapseShapeOp::getOperationName(), + tensor::ExpandShapeOp::getOperationName(), + tosa::MaximumOp::getOperationName(), + tosa::MinimumOp::getOperationName()}; // Trace through broadcast multiplication (mul by 1) FailureOr maybeNonOne = mulBroadcast(input); @@ -2304,8 +2333,8 @@ struct AttentionRewritePattern : public OpRewritePattern { auto add = maybeAdd.value(); - // One operand of the add is currentSeqLen (already tracked by KV-cache), - // the other is a negative constant (-windowSize). Try both operands. + // One operand of the add is currentSeqLen, the other is a negative constant + // (-windowSize). Try both operands. Value seqLenOperand; auto tryExtractNegativeConst = [&](Value candidate, Value other) -> FailureOr { @@ -2334,7 +2363,24 @@ struct AttentionRewritePattern : public OpRewritePattern { if (failed(maybeWindowSize)) return failure(); - return maybeWindowSize.value(); + std::optional clipMin; + std::optional clipMax; + auto maybeClip = tryClipPattern(seqLenOperand); + if (succeeded(maybeClip)) { + clipMin = maybeClip->clipMin; + clipMax = maybeClip->clipMax; + } + + // An unrelated greater(x - const, col) is not a sliding-window mask. The + // non-constant operand must resolve to an i32 currentSeqLen block argument. + FailureOr maybeSeqLen = + getValueSkipping(seqLenOperand, expandCollapseMinMax); + Value seqLen = succeeded(maybeSeqLen) ? maybeSeqLen.value() : seqLenOperand; + if (!isI32BlockArgument(seqLen, seqLenSkip)) + return failure(); + + return SlidingWindowResult{maybeWindowSize.value(), seqLen, clipMin, + clipMax}; } /* @@ -2541,7 +2587,11 @@ struct AttentionRewritePattern : public OpRewritePattern { if (!result.slidingWindowSize) { auto maybeSlidingWindow = trySlidingWindowPattern(input1, seqLenSkip); if (succeeded(maybeSlidingWindow)) { - result.slidingWindowSize = maybeSlidingWindow.value(); + auto slidingWindow = maybeSlidingWindow.value(); + result.slidingWindowSize = slidingWindow.windowSize; + result.slidingWindowSeqLen = slidingWindow.seqLen; + result.slidingWindowClipMin = slidingWindow.clipMin; + result.slidingWindowClipMax = slidingWindow.clipMax; } } return; @@ -2570,7 +2620,8 @@ struct AttentionRewritePattern : public OpRewritePattern { Value inputToContinue = select.getInput3(); SeqLenMaskResult currentResult{inputToContinue, nullptr, nullptr, - std::nullopt, std::nullopt, std::nullopt}; + std::nullopt, std::nullopt, std::nullopt, + nullptr, std::nullopt, std::nullopt}; // Analyze the first (outer) select analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip); @@ -2603,6 +2654,26 @@ struct AttentionRewritePattern : public OpRewritePattern { } } + // Sliding-window masking is defined relative to currentSeqLen. Reconcile + // the validated operand after all masks have been analyzed so the result is + // independent of the select nesting order. + if (currentResult.slidingWindowSize) { + if (currentResult.seqLen) { + if (!sameSeqLenBlockArg(currentResult.seqLen, + currentResult.slidingWindowSeqLen, seqLenSkip)) + return failure(); + // A single attention op cannot represent different clamps for the + // KV-cache and sliding-window masks. + if (currentResult.seqLenClipMin != currentResult.slidingWindowClipMin || + currentResult.seqLenClipMax != currentResult.slidingWindowClipMax) + return failure(); + } else { + currentResult.seqLen = currentResult.slidingWindowSeqLen; + currentResult.seqLenClipMin = currentResult.slidingWindowClipMin; + currentResult.seqLenClipMax = currentResult.slidingWindowClipMax; + } + } + // We need at least one pattern to be detected if (!currentResult.seqLen && !currentResult.prefixOffset && !currentResult.slidingWindowSize) diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir new file mode 100644 index 000000000000..fb170274b25c --- /dev/null +++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir @@ -0,0 +1,131 @@ +// RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-opt --tosa-to-rock -split-input-file -verify-diagnostics -o -| FileCheck %s + +// A sliding-window mask without a separate KV-cache mask must adopt its +// validated seq-len operand as currentSeqLen and preserve its clip. +// CHECK-LABEL: func @sliding_window_no_kvcache +// CHECK: %[[MAX:.*]] = tosa.maximum +// CHECK: %[[CLIP:.*]] = tosa.minimum %[[MAX]] +// CHECK: rock.attention +// CHECK: currentSeqLen = (%[[CLIP]] +// CHECK: slidingWindowSize = 3 +func.func @sliding_window_no_kvcache(%arg0: tensor<1xi32>, %arg1: tensor<12xf16>, %arg2: tensor<32xf16>, %arg3: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} { + %0 = "tosa.const"() <{values = dense<4> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32> + %4 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32> + %5 = "tosa.const"() <{values = dense<1> : tensor<1x2x1x8xi8>}> : () -> tensor<1x2x1x8xi8> + %7 = "tosa.const"() <{values = dense<1> : tensor<8x1x1x1xi32>}> : () -> tensor<8x1x1x1xi32> + %8 = "tosa.const"() <{values = dense<5.000000e-01> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16> + %9 = "tosa.const"() <{values = dense<0xFC00> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16> + %11 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16> + %16 = "tosa.const"() <{values = dense<-3> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32> + %17 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8> + %20 = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi32>}> : () -> tensor<8xi32> + %expanded = tensor.expand_shape %arg2 [[0, 1, 2, 3]] output_shape [1, 2, 8, 2] : tensor<32xf16> into tensor<1x2x8x2xf16> + %expanded_0 = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [1, 6, 1, 2] : tensor<12xf16> into tensor<1x6x1x2xf16> + %expanded_1 = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32> + %23 = tosa.maximum %expanded_1, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32> + %24 = tosa.minimum %23, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32> + %extracted_slice = tensor.extract_slice %expanded_0[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16> + %25 = tosa.transpose %expanded {perms = array} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16> + %collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16> + %collapsed_2 = tensor.collapse_shape %25 [[0, 1], [2], [3]] : tensor<1x2x2x8xf16> into tensor<2x2x8xf16> + %26 = tosa.matmul %collapsed, %collapsed_2, %11, %11 {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16> + %expanded_3 = tensor.expand_shape %26 [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16> + %27 = tosa.mul %expanded_3, %8, %17 : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16> + %28 = tosa.add %24, %16 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32> + %29 = tosa.mul %28, %7, %17 : (tensor<1x1x1x1xi32>, tensor<8x1x1x1xi32>, tensor<1xi8>) -> tensor<8x1x1x1xi32> + %collapsed_4 = tensor.collapse_shape %29 [[0, 1, 2, 3]] : tensor<8x1x1x1xi32> into tensor<8xi32> + %30 = tosa.greater %collapsed_4, %20 : (tensor<8xi32>, tensor<8xi32>) -> tensor<8xi1> + %31 = tosa.cast %30 : (tensor<8xi1>) -> tensor<8xi32> + %32 = tosa.cast %31 : (tensor<8xi32>) -> tensor<8xi8> + %expanded_5 = tensor.expand_shape %32 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8> + %33 = tosa.mul %expanded_5, %5, %17 : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8> + %34 = tosa.cast %33 : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1> + %35 = tosa.select %34, %9, %27 : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16> + %43 = tosa.cast %35 : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32> + %44 = tosa.reduce_max %43 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32> + %45 = tosa.mul %44, %4, %17 : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32> + %46 = tosa.sub %43, %45 : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32> + %47 = tosa.exp %46 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32> + %48 = tosa.reduce_sum %47 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32> + %49 = tosa.mul %48, %4, %17 : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32> + %50 = tosa.reciprocal %49 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32> + %51 = tosa.mul %47, %50, %17 : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32> + %52 = tosa.cast %51 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf16> + %collapsed_6 = tensor.collapse_shape %52 [[0, 1], [2], [3]] : tensor<1x2x1x8xf16> into tensor<2x1x8xf16> + %expanded_7 = tensor.expand_shape %arg3 [[0, 1, 2]] output_shape [2, 8, 2] : tensor<32xf16> into tensor<2x8x2xf16> + %53 = tosa.matmul %collapsed_6, %expanded_7, %11, %11 {acc_type = f32} : (tensor<2x1x8xf16>, tensor<2x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x2xf16> + %expanded_8 = tensor.expand_shape %53 [[0, 1], [2], [3]] output_shape [1, 2, 1, 2] : tensor<2x1x2xf16> into tensor<1x2x1x2xf16> + %54 = tosa.transpose %expanded_8 {perms = array} : (tensor<1x2x1x2xf16>) -> tensor<1x1x2x2xf16> + %collapsed_9 = tensor.collapse_shape %54 [[0, 1, 2, 3]] : tensor<1x1x2x2xf16> into tensor<4xf16> + return %collapsed_9 : tensor<4xf16> +} + +// ----- + +// A greater(x - window, col) mask whose x is not currentSeqLen must not be +// classified as sliding-window attention. +// CHECK-LABEL: func @not_sliding_window_wrong_operand +// CHECK: rock.attention +// CHECK: currentSeqLen = +// CHECK-NOT: slidingWindowSize +func.func @not_sliding_window_wrong_operand(%arg0: tensor<1xi32>, %arg1: tensor<12xf16>, %arg2: tensor<32xf16>, %arg3: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} { + %0 = "tosa.const"() <{values = dense<4> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32> + %4 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32> + %5 = "tosa.const"() <{values = dense<1> : tensor<1x2x1x8xi8>}> : () -> tensor<1x2x1x8xi8> + %7 = "tosa.const"() <{values = dense<1> : tensor<8x1x1x1xi32>}> : () -> tensor<8x1x1x1xi32> + %8 = "tosa.const"() <{values = dense<5.000000e-01> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16> + %9 = "tosa.const"() <{values = dense<0xFC00> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16> + %11 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16> + %16 = "tosa.const"() <{values = dense<-3> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32> + %17 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8> + %18 = "tosa.const"() <{values = dense<1> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32> + %20 = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi32>}> : () -> tensor<8xi32> + %cst = arith.constant dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32> + %expanded = tensor.expand_shape %arg2 [[0, 1, 2, 3]] output_shape [1, 2, 8, 2] : tensor<32xf16> into tensor<1x2x8x2xf16> + %expanded_0 = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [1, 6, 1, 2] : tensor<12xf16> into tensor<1x6x1x2xf16> + %expanded_1 = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32> + %23 = tosa.maximum %expanded_1, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32> + %24 = tosa.minimum %23, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32> + %extracted_slice = tensor.extract_slice %expanded_0[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16> + %25 = tosa.transpose %expanded {perms = array} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16> + %collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16> + %collapsed_2 = tensor.collapse_shape %25 [[0, 1], [2], [3]] : tensor<1x2x2x8xf16> into tensor<2x2x8xf16> + %26 = tosa.matmul %collapsed, %collapsed_2, %11, %11 {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16> + %expanded_3 = tensor.expand_shape %26 [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16> + %27 = tosa.mul %expanded_3, %8, %17 : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16> + // Subtract the window size from a constant, not from currentSeqLen. + %28 = tosa.add %0, %16 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32> + %29 = tosa.mul %28, %7, %17 : (tensor<1x1x1x1xi32>, tensor<8x1x1x1xi32>, tensor<1xi8>) -> tensor<8x1x1x1xi32> + %collapsed_4 = tensor.collapse_shape %29 [[0, 1, 2, 3]] : tensor<8x1x1x1xi32> into tensor<8xi32> + %30 = tosa.greater %collapsed_4, %20 : (tensor<8xi32>, tensor<8xi32>) -> tensor<8xi1> + %31 = tosa.cast %30 : (tensor<8xi1>) -> tensor<8xi32> + %32 = tosa.cast %31 : (tensor<8xi32>) -> tensor<8xi8> + %expanded_5 = tensor.expand_shape %32 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8> + %33 = tosa.mul %expanded_5, %5, %17 : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8> + %34 = tosa.cast %33 : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1> + %35 = tosa.select %34, %9, %27 : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16> + %36 = tosa.mul %24, %18, %17 : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32> + %37 = tosa.greater %cst, %36 : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1> + %38 = tosa.cast %37 : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32> + %39 = tosa.cast %38 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8> + %40 = tosa.mul %39, %5, %17 : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8> + %41 = tosa.cast %40 : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1> + %42 = tosa.select %41, %9, %35 : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16> + %43 = tosa.cast %42 : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32> + %44 = tosa.reduce_max %43 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32> + %45 = tosa.mul %44, %4, %17 : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32> + %46 = tosa.sub %43, %45 : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32> + %47 = tosa.exp %46 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32> + %48 = tosa.reduce_sum %47 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32> + %49 = tosa.mul %48, %4, %17 : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32> + %50 = tosa.reciprocal %49 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32> + %51 = tosa.mul %47, %50, %17 : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32> + %52 = tosa.cast %51 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf16> + %collapsed_6 = tensor.collapse_shape %52 [[0, 1], [2], [3]] : tensor<1x2x1x8xf16> into tensor<2x1x8xf16> + %expanded_7 = tensor.expand_shape %arg3 [[0, 1, 2]] output_shape [2, 8, 2] : tensor<32xf16> into tensor<2x8x2xf16> + %53 = tosa.matmul %collapsed_6, %expanded_7, %11, %11 {acc_type = f32} : (tensor<2x1x8xf16>, tensor<2x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x2xf16> + %expanded_8 = tensor.expand_shape %53 [[0, 1], [2], [3]] output_shape [1, 2, 1, 2] : tensor<2x1x2xf16> into tensor<1x2x1x2xf16> + %54 = tosa.transpose %expanded_8 {perms = array} : (tensor<1x2x1x2xf16>) -> tensor<1x1x2x2xf16> + %collapsed_9 = tensor.collapse_shape %54 [[0, 1, 2, 3]] : tensor<1x1x2x2xf16> into tensor<4xf16> + return %collapsed_9 : tensor<4xf16> +} From ba15f65c08baf95e16390a26f678d46169a49035 Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Thu, 23 Jul 2026 18:27:55 +0000 Subject: [PATCH 2/4] Fix attention mask traversal for three nested selects Iteratively peel prefix-causal, KV-cache, and sliding-window masks so the third select does not remain in the elementwise region and trigger "Cannot trace first gemm index for linalg.generic op". Add conversion and nightly E2E regression coverage. Co-authored-by: Cursor --- mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 51 +++++----- .../tosa-to-rock-attention-three-mask.mlir | 94 +++++++++++++++++++ ...-sliding-window-kvcache-prefix-causal.mlir | 90 ++++++++++++++++++ 3 files changed, 209 insertions(+), 26 deletions(-) create mode 100644 mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-three-mask.mlir create mode 100644 mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp index af7b6163eb68..09f3448ae452 100644 --- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp +++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp @@ -2626,32 +2626,31 @@ struct AttentionRewritePattern : public OpRewritePattern { // Analyze the first (outer) select analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip); - // Check if the inputToContinue (input3) is another chained select with - // -inf. This handles cases where multiple mask patterns (KVCache, prefix - // causal, sliding window) use separate selects. - bool haveSeqLen = currentResult.seqLen != nullptr; - bool havePrefixOffset = currentResult.prefixOffset != nullptr; - bool haveSlidingWindow = currentResult.slidingWindowSize.has_value(); - - // Try chaining if we found at least one pattern but not all - bool foundAny = haveSeqLen || havePrefixOffset || haveSlidingWindow; - bool foundAll = haveSeqLen && havePrefixOffset && haveSlidingWindow; - if (foundAny && !foundAll) { - auto maybeChainedSelect = getSelectWithNegInf(inputToContinue); - if (succeeded(maybeChainedSelect)) { - auto chainedSelect = maybeChainedSelect.value(); - // Try to analyze the chained select for the missing pattern - analyzeSelectForSeqLenMask(chainedSelect, currentResult, opsToSkip, - seqLenSkip); - // Only update inputToContinue if we found a complementary pattern - bool foundComplementary = - (!haveSeqLen && currentResult.seqLen) || - (!havePrefixOffset && currentResult.prefixOffset) || - (!haveSlidingWindow && currentResult.slidingWindowSize.has_value()); - if (foundComplementary) { - currentResult.inputToContinue = chainedSelect.getInput3(); - } - } + // Iteratively peel chained select(mask, -inf, scores) ops to detect + // separately nested KV-cache, prefix-causal, and sliding-window masks. + // Use prefixOffset as the recognition marker for a prefix-causal select + // (col > row + prefixOffset). A standard causal select (col > row) has no + // prefixOffset, so it remains in inputToContinue for getCausal() to handle + // after the sequence-length masks have been peeled. + auto recognizedMaskCount = [](const SeqLenMaskResult &result) { + return (result.seqLen ? 1 : 0) + (result.prefixOffset ? 1 : 0) + + (result.slidingWindowSize.has_value() ? 1 : 0); + }; + while (recognizedMaskCount(currentResult) > 0 && + recognizedMaskCount(currentResult) < 3) { + auto maybeChainedSelect = + getSelectWithNegInf(currentResult.inputToContinue); + if (failed(maybeChainedSelect)) + break; + + auto chainedSelect = maybeChainedSelect.value(); + int before = recognizedMaskCount(currentResult); + analyzeSelectForSeqLenMask(chainedSelect, currentResult, opsToSkip, + seqLenSkip); + // Leave an unrecognized or duplicate mask in the elementwise region. + if (recognizedMaskCount(currentResult) == before) + break; + currentResult.inputToContinue = chainedSelect.getInput3(); } // Sliding-window masking is defined relative to currentSeqLen. Reconcile diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-three-mask.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-three-mask.mlir new file mode 100644 index 000000000000..682234ada8c1 --- /dev/null +++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-three-mask.mlir @@ -0,0 +1,94 @@ +// RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-opt --tosa-to-rock -verify-diagnostics | FileCheck %s + +// Three nested select ops must all be folded into a single rock.attention. +// From inner to outer, the masks are prefix-causal, KV-cache, and sliding +// window. All four corresponding attention properties must be preserved. +// CHECK-LABEL: func @attention_three_mask +// CHECK: rock.attention +// CHECK-DAG: currentSeqLen = (%{{.*}} +// CHECK-DAG: prefixOffset = (%{{.*}} +// CHECK-DAG: causal +// CHECK-DAG: slidingWindowSize = 3 +// CHECK: qk = elementwise { +// CHECK-NOT: tosa.select +// CHECK: rock.yield + +module { + func.func @attention_three_mask(%arg0: tensor<1xi32>, %arg1: tensor<9216xf16>, %arg2: tensor<2048xf16>, %arg3: tensor<2048xf16>, %arg4: tensor<2xi32>) -> tensor<7168xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} { + %0 = "tosa.const"() <{values = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>}> : () -> tensor<1x8xi32> + %1 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x14x4x8xf32>}> : () -> tensor<2x14x4x8xf32> + %2 = "tosa.const"() <{values = dense<0xFC00> : tensor<2x14x4x8xf16>}> : () -> tensor<2x14x4x8xf16> + %3 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16> + %4 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x2x7x64x8xf16>}> : () -> tensor<2x2x7x64x8xf16> + %5 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x2x7x8x64xf16>}> : () -> tensor<2x2x7x8x64xf16> + %6 = "tosa.const"() <{values = dense<1> : tensor<2x14x4x8xi8>}> : () -> tensor<2x14x4x8xi8> + %7 = "tosa.const"() <{values = dense<1.250000e-01> : tensor<2x14x4x8xf16>}> : () -> tensor<2x14x4x8xf16> + %8 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8> + %9 = "tosa.const"() <{values = dense<1> : tensor<4x8xi32>}> : () -> tensor<4x8xi32> + %10 = "tosa.const"() <{values = dense<1> : tensor<4x1xi32>}> : () -> tensor<4x1xi32> + %11 = "tosa.const"() <{values = dense<1> : tensor<2x8xi32>}> : () -> tensor<2x8xi32> + %12 = "tosa.const"() <{values = dense<[[0], [1], [2], [3]]> : tensor<4x1xi32>}> : () -> tensor<4x1xi32> + %swoff = "tosa.const"() <{values = dense<-3> : tensor<1x1xi32>}> : () -> tensor<1x1xi32> + %expanded = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [2, 4, 18, 64] : tensor<9216xf16> into tensor<2x4x18x64xf16> + %13 = tosa.transpose %expanded {perms = array} : (tensor<2x4x18x64xf16>) -> tensor<2x18x4x64xf16> + %expanded_0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [1, 1] : tensor<1xi32> into tensor<1x1xi32> + %14 = tosa.mul %expanded_0, %10, %8 : (tensor<1x1xi32>, tensor<4x1xi32>, tensor<1xi8>) -> tensor<4x1xi32> + %15 = tosa.add %14, %12 : (tensor<4x1xi32>, tensor<4x1xi32>) -> tensor<4x1xi32> + %16 = tosa.mul %15, %9, %8 : (tensor<4x1xi32>, tensor<4x8xi32>, tensor<1xi8>) -> tensor<4x8xi32> + %17 = tosa.mul %0, %9, %8 : (tensor<1x8xi32>, tensor<4x8xi32>, tensor<1xi8>) -> tensor<4x8xi32> + %18 = tosa.greater %17, %16 : (tensor<4x8xi32>, tensor<4x8xi32>) -> tensor<4x8xi1> + %19 = tosa.cast %18 : (tensor<4x8xi1>) -> tensor<4x8xi32> + %20 = tosa.cast %19 : (tensor<4x8xi32>) -> tensor<4x8xi8> + %expanded_1 = tensor.expand_shape %20 [[0, 1, 2], [3]] output_shape [1, 1, 4, 8] : tensor<4x8xi8> into tensor<1x1x4x8xi8> + %21 = tosa.mul %expanded_1, %6, %8 : (tensor<1x1x4x8xi8>, tensor<2x14x4x8xi8>, tensor<1xi8>) -> tensor<2x14x4x8xi8> + %22 = tosa.mul %0, %11, %8 : (tensor<1x8xi32>, tensor<2x8xi32>, tensor<1xi8>) -> tensor<2x8xi32> + %expanded_2 = tensor.expand_shape %arg4 [[0, 1]] output_shape [2, 1] : tensor<2xi32> into tensor<2x1xi32> + %23 = tosa.mul %expanded_2, %11, %8 : (tensor<2x1xi32>, tensor<2x8xi32>, tensor<1xi8>) -> tensor<2x8xi32> + %24 = tosa.greater %22, %23 : (tensor<2x8xi32>, tensor<2x8xi32>) -> tensor<2x8xi1> + %25 = tosa.cast %24 : (tensor<2x8xi1>) -> tensor<2x8xi32> + %26 = tosa.cast %25 : (tensor<2x8xi32>) -> tensor<2x8xi8> + %expanded_3 = tensor.expand_shape %26 [[0, 1, 2], [3]] output_shape [2, 1, 1, 8] : tensor<2x8xi8> into tensor<2x1x1x8xi8> + %27 = tosa.mul %expanded_3, %6, %8 : (tensor<2x1x1x8xi8>, tensor<2x14x4x8xi8>, tensor<1xi8>) -> tensor<2x14x4x8xi8> + %sw0 = tosa.add %expanded_2, %swoff : (tensor<2x1xi32>, tensor<1x1xi32>) -> tensor<2x1xi32> + %sw1 = tosa.mul %sw0, %11, %8 : (tensor<2x1xi32>, tensor<2x8xi32>, tensor<1xi8>) -> tensor<2x8xi32> + %sw2 = tosa.greater %sw1, %22 : (tensor<2x8xi32>, tensor<2x8xi32>) -> tensor<2x8xi1> + %sw3 = tosa.cast %sw2 : (tensor<2x8xi1>) -> tensor<2x8xi32> + %sw4 = tosa.cast %sw3 : (tensor<2x8xi32>) -> tensor<2x8xi8> + %expanded_sw = tensor.expand_shape %sw4 [[0, 1, 2], [3]] output_shape [2, 1, 1, 8] : tensor<2x8xi8> into tensor<2x1x1x8xi8> + %sw5 = tosa.mul %expanded_sw, %6, %8 : (tensor<2x1x1x8xi8>, tensor<2x14x4x8xi8>, tensor<1xi8>) -> tensor<2x14x4x8xi8> + %extracted_slice = tensor.extract_slice %13[0, 0, 0, 0] [2, 14, 4, 64] [1, 1, 1, 1] : tensor<2x18x4x64xf16> to tensor<2x14x4x64xf16> + %expanded_4 = tensor.expand_shape %arg2 [[0, 1, 2, 3, 4]] output_shape [2, 2, 1, 8, 64] : tensor<2048xf16> into tensor<2x2x1x8x64xf16> + %28 = tosa.mul %expanded_4, %5, %8 : (tensor<2x2x1x8x64xf16>, tensor<2x2x7x8x64xf16>, tensor<1xi8>) -> tensor<2x2x7x8x64xf16> + %expanded_5 = tensor.expand_shape %arg3 [[0, 1, 2, 3, 4]] output_shape [2, 2, 1, 8, 64] : tensor<2048xf16> into tensor<2x2x1x8x64xf16> + %29 = tosa.transpose %expanded_5 {perms = array} : (tensor<2x2x1x8x64xf16>) -> tensor<2x2x1x64x8xf16> + %30 = tosa.mul %29, %4, %8 : (tensor<2x2x1x64x8xf16>, tensor<2x2x7x64x8xf16>, tensor<1xi8>) -> tensor<2x2x7x64x8xf16> + %collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2], [3]] : tensor<2x14x4x64xf16> into tensor<28x4x64xf16> + %collapsed_6 = tensor.collapse_shape %30 [[0, 1, 2], [3], [4]] : tensor<2x2x7x64x8xf16> into tensor<28x64x8xf16> + %31 = tosa.matmul %collapsed, %collapsed_6, %3, %3 {acc_type = f32} : (tensor<28x4x64xf16>, tensor<28x64x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<28x4x8xf16> + %expanded_7 = tensor.expand_shape %31 [[0, 1], [2], [3]] output_shape [2, 14, 4, 8] : tensor<28x4x8xf16> into tensor<2x14x4x8xf16> + %32 = tosa.mul %expanded_7, %7, %8 : (tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>, tensor<1xi8>) -> tensor<2x14x4x8xf16> + %33 = tosa.cast %21 : (tensor<2x14x4x8xi8>) -> tensor<2x14x4x8xi1> + %34 = tosa.select %33, %2, %32 : (tensor<2x14x4x8xi1>, tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf16> + %35 = tosa.cast %27 : (tensor<2x14x4x8xi8>) -> tensor<2x14x4x8xi1> + %36 = tosa.select %35, %2, %34 : (tensor<2x14x4x8xi1>, tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf16> + %sw6 = tosa.cast %sw5 : (tensor<2x14x4x8xi8>) -> tensor<2x14x4x8xi1> + %sw7 = tosa.select %sw6, %2, %36 : (tensor<2x14x4x8xi1>, tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf16> + %37 = tosa.cast %sw7 : (tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf32> + %38 = tosa.reduce_max %37 {axis = 3 : i32} : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x1xf32> + %39 = tosa.mul %38, %1, %8 : (tensor<2x14x4x1xf32>, tensor<2x14x4x8xf32>, tensor<1xi8>) -> tensor<2x14x4x8xf32> + %40 = tosa.sub %37, %39 : (tensor<2x14x4x8xf32>, tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf32> + %41 = tosa.exp %40 : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf32> + %42 = tosa.reduce_sum %41 {axis = 3 : i32} : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x1xf32> + %43 = tosa.mul %42, %1, %8 : (tensor<2x14x4x1xf32>, tensor<2x14x4x8xf32>, tensor<1xi8>) -> tensor<2x14x4x8xf32> + %44 = tosa.reciprocal %43 : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf32> + %45 = tosa.mul %41, %44, %8 : (tensor<2x14x4x8xf32>, tensor<2x14x4x8xf32>, tensor<1xi8>) -> tensor<2x14x4x8xf32> + %46 = tosa.cast %45 : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf16> + %collapsed_8 = tensor.collapse_shape %46 [[0, 1], [2], [3]] : tensor<2x14x4x8xf16> into tensor<28x4x8xf16> + %collapsed_9 = tensor.collapse_shape %28 [[0, 1, 2], [3], [4]] : tensor<2x2x7x8x64xf16> into tensor<28x8x64xf16> + %47 = tosa.matmul %collapsed_8, %collapsed_9, %3, %3 {acc_type = f32} : (tensor<28x4x8xf16>, tensor<28x8x64xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<28x4x64xf16> + %expanded_10 = tensor.expand_shape %47 [[0, 1], [2], [3]] output_shape [2, 14, 4, 64] : tensor<28x4x64xf16> into tensor<2x14x4x64xf16> + %48 = tosa.transpose %expanded_10 {perms = array} : (tensor<2x14x4x64xf16>) -> tensor<2x4x14x64xf16> + %collapsed_11 = tensor.collapse_shape %48 [[0, 1, 2, 3]] : tensor<2x4x14x64xf16> into tensor<7168xf16> + return %collapsed_11 : tensor<7168xf16> + } +} diff --git a/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir b/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir new file mode 100644 index 000000000000..1eb7dea06a62 --- /dev/null +++ b/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir @@ -0,0 +1,90 @@ +// RUN: rocmlir-gen -fut mlir_attention --arch %arch --clone-harness %s | rocmlir-driver -kernel-pipeline=migraphx,highlevel -host-pipeline=migraphx,highlevel | FileCheck %s --check-prefix=FOLD +// RUN: rocmlir-gen -fut mlir_attention --arch %arch --clone-harness %s | rocmlir-driver -kernel-pipeline=migraphx,highlevel -host-pipeline=migraphx,highlevel | rocmlir-gen -ph -fut mlir_attention_wrapper -relDiff_threshold 0.00001 -rand_min_int 1 -rand_max_int 2 -rand_type_int_for_inputs=2,4 --verifier clone - -pr | rocmlir-driver -host-pipeline mhal -kernel-pipeline full | xmir-runner --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext,%linalg_test_lib_dir/libmlir_c_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_async_runtime%shlibext --entry-point-result=void | FileCheck %s --check-prefix=E2E + +// Verify that all three nested selects are folded into one attention op. +// FOLD: rock.attention{ +// FOLD: currentSeqLen = ( +// FOLD: prefixOffset = ( +// FOLD: slidingWindowSize = 1 +// FOLD: causal +// FOLD: qk = elementwise { +// FOLD-NOT: tosa.select +// FOLD: rock.yield + +// E2E: [1 1 1] +// E2E-NEXT: [1 1 1] + +module { + func.func @mlir_attention(%arg0: !migraphx.shaped<2x6x2x2xf16, 24x4x2x1>, %arg1: !migraphx.shaped<2x2x4x2xf16, 16x8x2x1>, %arg2: !migraphx.shaped<2x1xsi32, 1x1>, %arg3: !migraphx.shaped<2x2x4x2xf16, 16x8x2x1>, %arg4: !migraphx.shaped<2x1xsi32, 1x1>) -> (!migraphx.shaped<2x2x2x4xf16, 16x8x4x1>, !migraphx.shaped<2x2x2x2x1xf32, 8x4x2x1x1>) attributes {rock.kernel = "mixr"} { + %0 = migraphx.literal(dense<[0, 1, 2, 3]> : tensor<4xsi32>) : <4xsi32, 1> + %1 = migraphx.literal(dense<[[0], [1]]> : tensor<2x1xsi32>) : <2x1xsi32, 1x1> + %2 = migraphx.literal(dense<1> : tensor<2x4xsi32>) : <2x4xsi32, 4x1> + %3 = migraphx.literal(dense<0xFC00> : tensor<1xf16>) : <1xf16, 1> + %4 = migraphx.literal(dense<5.000000e-01> : tensor<1xf16>) : <1xf16, 1> + %sliding_offset = migraphx.literal(dense<-1> : tensor<1xsi32>) : <1xsi32, 1> + %fixed_seq_len = migraphx.literal(dense<2> : tensor<2x1xsi32>) : <2x1xsi32, 1x1> + %seq_len = migraphx.clip %arg2, %fixed_seq_len, %fixed_seq_len : <2x1xsi32, 1x1>, <2x1xsi32, 1x1>, <2x1xsi32, 1x1> -> <2x1xsi32, 1x1> + %5 = migraphx.reshape %arg0 {dims = [2, 6, 1, 2, 2]} : <2x6x2x2xf16, 24x4x2x1> -> <2x6x1x2x2xf16, 24x4x4x2x1> + %6 = migraphx.multibroadcast %5 {out_dyn_dims = [], out_lens = [2, 6, 2, 2, 2]} : <2x6x1x2x2xf16, 24x4x4x2x1> -> <2x6x2x2x2xf16, 24x4x0x2x1> + %7 = migraphx.reshape %arg1 {dims = [2, 2, 2, 2, 2]} : <2x2x4x2xf16, 16x8x2x1> -> <2x2x2x2x2xf16, 16x8x4x2x1> + %8 = migraphx.reshape %arg3 {dims = [2, 2, 2, 2, 2]} : <2x2x4x2xf16, 16x8x2x1> -> <2x2x2x2x2xf16, 16x8x4x2x1> + %9 = migraphx.slice %6 {axes = [1], ends = [2], starts = [0]} : <2x6x2x2x2xf16, 24x4x0x2x1> -> <2x2x2x2x2xf16, 24x4x0x2x1> + %10 = migraphx.transpose %7 {permutation = [0, 1, 2, 4, 3]} : <2x2x2x2x2xf16, 16x8x4x2x1> -> <2x2x2x2x2xf16, 16x8x4x1x2> + %11 = migraphx.multibroadcast %3 {out_dyn_dims = [], out_lens = [2, 2, 2, 2, 2]} : <1xf16, 1> -> <2x2x2x2x2xf16, 0x0x0x0x0> + %12 = migraphx.multibroadcast %4 {out_dyn_dims = [], out_lens = [2, 2, 2, 2, 2]} : <1xf16, 1> -> <2x2x2x2x2xf16, 0x0x0x0x0> + %13 = migraphx.dot %9, %10 : <2x2x2x2x2xf16, 24x4x0x2x1>, <2x2x2x2x2xf16, 16x8x4x1x2> -> <2x2x2x2x2xf16, 16x8x4x2x1> + %14 = migraphx.mul %13, %12 : <2x2x2x2x2xf16, 16x8x4x2x1>, <2x2x2x2x2xf16, 0x0x0x0x0> -> <2x2x2x2x2xf16, 16x8x4x2x1> + + // Prefix-causal mask. + %15 = migraphx.multibroadcast %arg4 {out_dyn_dims = [], out_lens = [2, 1]} : <2x1xsi32, 1x1> -> <2x1xsi32, 1x0> + %16 = migraphx.add %1, %15 : <2x1xsi32, 1x1>, <2x1xsi32, 1x0> -> <2x1xsi32, 1x1> + %17 = migraphx.multibroadcast %16 {out_dyn_dims = [], out_lens = [2, 4]} : <2x1xsi32, 1x1> -> <2x4xsi32, 1x0> + %18 = migraphx.mul %17, %2 : <2x4xsi32, 1x0>, <2x4xsi32, 4x1> -> <2x4xsi32, 4x1> + %19 = migraphx.broadcast %0 {axis = 1 : i64, out_lens = [2, 4]} : <4xsi32, 1> -> <2x4xsi32, 0x1> + %20 = migraphx.mul %19, %2 : <2x4xsi32, 0x1>, <2x4xsi32, 4x1> -> <2x4xsi32, 4x1> + %21 = migraphx.greater %20, %18 : <2x4xsi32, 4x1>, <2x4xsi32, 4x1> -> <2x4xsi32, 4x1> + %22 = migraphx.convert %21 {target_type = 0 : i64} : <2x4xsi32, 4x1> to <2x4xsi8, 4x1> + %23 = migraphx.reshape %22 {dims = [1, 1, 2, 2, 2]} : <2x4xsi8, 4x1> -> <1x1x2x2x2xsi8, 8x8x4x2x1> + %24 = migraphx.multibroadcast %23 {out_dyn_dims = [], out_lens = [2, 2, 2, 2, 2]} : <1x1x2x2x2xsi8, 8x8x4x2x1> -> <2x2x2x2x2xsi8, 0x0x4x2x1> + %25 = migraphx.where %24, %11, %14 : <2x2x2x2x2xsi8, 0x0x4x2x1>, <2x2x2x2x2xf16, 0x0x0x0x0>, <2x2x2x2x2xf16, 16x8x4x2x1> -> <2x2x2x2x2xf16, 16x8x4x2x1> + + // KV-cache mask. + %26 = migraphx.broadcast %0 {axis = 1 : i64, out_lens = [2, 4]} : <4xsi32, 1> -> <2x4xsi32, 0x1> + %27 = migraphx.multibroadcast %seq_len {out_dyn_dims = [], out_lens = [2, 4]} : <2x1xsi32, 1x1> -> <2x4xsi32, 1x0> + %28 = migraphx.greater %26, %27 : <2x4xsi32, 0x1>, <2x4xsi32, 1x0> -> <2x4xsi32, 4x1> + %29 = migraphx.convert %28 {target_type = 0 : i64} : <2x4xsi32, 4x1> to <2x4xsi8, 4x1> + %30 = migraphx.reshape %29 {dims = [2, 1, 2, 1, 2]} : <2x4xsi8, 4x1> -> <2x1x2x1x2xsi8, 4x4x2x2x1> + %31 = migraphx.multibroadcast %30 {out_dyn_dims = [], out_lens = [2, 2, 2, 2, 2]} : <2x1x2x1x2xsi8, 4x4x2x2x1> -> <2x2x2x2x2xsi8, 4x0x2x0x1> + %32 = migraphx.where %31, %11, %25 : <2x2x2x2x2xsi8, 4x0x2x0x1>, <2x2x2x2x2xf16, 0x0x0x0x0>, <2x2x2x2x2xf16, 16x8x4x2x1> -> <2x2x2x2x2xf16, 16x8x4x2x1> + + // Sliding-window mask. + %sliding_offset_bcast = migraphx.multibroadcast %sliding_offset {out_dyn_dims = [], out_lens = [2, 1]} : <1xsi32, 1> -> <2x1xsi32, 0x1> + %window_start = migraphx.add %seq_len, %sliding_offset_bcast : <2x1xsi32, 1x1>, <2x1xsi32, 0x1> -> <2x1xsi32, 1x1> + %window_starts = migraphx.multibroadcast %window_start {out_dyn_dims = [], out_lens = [2, 4]} : <2x1xsi32, 1x1> -> <2x4xsi32, 1x0> + %window_pred = migraphx.greater %window_starts, %26 : <2x4xsi32, 1x0>, <2x4xsi32, 0x1> -> <2x4xsi32, 4x1> + %window_i8 = migraphx.convert %window_pred {target_type = 0 : i64} : <2x4xsi32, 4x1> to <2x4xsi8, 4x1> + %window_reshaped = migraphx.reshape %window_i8 {dims = [2, 1, 2, 1, 2]} : <2x4xsi8, 4x1> -> <2x1x2x1x2xsi8, 4x4x2x2x1> + %window_mask = migraphx.multibroadcast %window_reshaped {out_dyn_dims = [], out_lens = [2, 2, 2, 2, 2]} : <2x1x2x1x2xsi8, 4x4x2x2x1> -> <2x2x2x2x2xsi8, 4x0x2x0x1> + %window_masked = migraphx.where %window_mask, %11, %32 : <2x2x2x2x2xsi8, 4x0x2x0x1>, <2x2x2x2x2xf16, 0x0x0x0x0>, <2x2x2x2x2xf16, 16x8x4x2x1> -> <2x2x2x2x2xf16, 16x8x4x2x1> + + %33 = migraphx.convert %window_masked {target_type = 2 : i64} : <2x2x2x2x2xf16, 16x8x4x2x1> to <2x2x2x2x2xf32, 16x8x4x2x1> + %34 = migraphx.reshape %33 {dims = [2, 2, 2, 2, 2]} : <2x2x2x2x2xf32, 16x8x4x2x1> -> <2x2x2x2x2xf32, 16x8x4x2x1> + %35 = migraphx.reduce_max %34 {axes = [4]} : <2x2x2x2x2xf32, 16x8x4x2x1> -> <2x2x2x2x1xf32, 8x4x2x1x1> + %36 = migraphx.reshape %35 {dims = [2, 2, 2, 2, 1]} : <2x2x2x2x1xf32, 8x4x2x1x1> -> <2x2x2x2x1xf32, 8x4x2x1x1> + %37 = migraphx.multibroadcast %36 {out_dyn_dims = [], out_lens = [2, 2, 2, 2, 2]} : <2x2x2x2x1xf32, 8x4x2x1x1> -> <2x2x2x2x2xf32, 8x4x2x1x0> + %38 = migraphx.sub %33, %37 : <2x2x2x2x2xf32, 16x8x4x2x1>, <2x2x2x2x2xf32, 8x4x2x1x0> -> <2x2x2x2x2xf32, 16x8x4x2x1> + %39 = migraphx.exp %38 : <2x2x2x2x2xf32, 16x8x4x2x1> -> <2x2x2x2x2xf32, 16x8x4x2x1> + %40 = migraphx.reshape %39 {dims = [2, 2, 2, 2, 2]} : <2x2x2x2x2xf32, 16x8x4x2x1> -> <2x2x2x2x2xf32, 16x8x4x2x1> + %41 = migraphx.reduce_sum %40 {axes = [4]} : <2x2x2x2x2xf32, 16x8x4x2x1> -> <2x2x2x2x1xf32, 8x4x2x1x1> + %42 = migraphx.reshape %41 {dims = [2, 2, 2, 2, 1]} : <2x2x2x2x1xf32, 8x4x2x1x1> -> <2x2x2x2x1xf32, 8x4x2x1x1> + %43 = migraphx.multibroadcast %42 {out_dyn_dims = [], out_lens = [2, 2, 2, 2, 2]} : <2x2x2x2x1xf32, 8x4x2x1x1> -> <2x2x2x2x2xf32, 8x4x2x1x0> + %44 = migraphx.div %39, %43 : <2x2x2x2x2xf32, 16x8x4x2x1>, <2x2x2x2x2xf32, 8x4x2x1x0> -> <2x2x2x2x2xf32, 16x8x4x2x1> + %45 = migraphx.convert %44 {target_type = 1 : i64} : <2x2x2x2x2xf32, 16x8x4x2x1> to <2x2x2x2x2xf16, 16x8x4x2x1> + %46 = migraphx.dot %45, %8 : <2x2x2x2x2xf16, 16x8x4x2x1>, <2x2x2x2x2xf16, 16x8x4x2x1> -> <2x2x2x2x2xf16, 16x8x4x2x1> + %47 = migraphx.transpose %46 {permutation = [0, 2, 3, 1, 4]} : <2x2x2x2x2xf16, 16x8x4x2x1> -> <2x2x2x2x2xf16, 16x4x2x8x1> + %48 = migraphx.reshape %47 {dims = [2, 2, 2, 4]} : <2x2x2x2x2xf16, 16x4x2x8x1> -> <2x2x2x4xf16, 16x8x4x1> + %49 = migraphx.log %42 : <2x2x2x2x1xf32, 8x4x2x1x1> -> <2x2x2x2x1xf32, 8x4x2x1x1> + %50 = migraphx.add %36, %49 : <2x2x2x2x1xf32, 8x4x2x1x1>, <2x2x2x2x1xf32, 8x4x2x1x1> -> <2x2x2x2x1xf32, 8x4x2x1x1> + return %48, %50 : !migraphx.shaped<2x2x2x4xf16, 16x8x4x1>, !migraphx.shaped<2x2x2x2x1xf32, 8x4x2x1x1> + } +} From 381912b5d80c42c70eb76ca53af4830224a408dc Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Thu, 23 Jul 2026 19:37:48 +0000 Subject: [PATCH 3/4] [NFC] Make three-mask folding checks order-independent Use FileCheck DAG assertions so optional attention property ordering does not create brittle test failures. Co-authored-by: Cursor --- ...xr-attention-sliding-window-kvcache-prefix-causal.mlir | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir b/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir index 1eb7dea06a62..1fd16f9b75da 100644 --- a/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir +++ b/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir @@ -3,10 +3,10 @@ // Verify that all three nested selects are folded into one attention op. // FOLD: rock.attention{ -// FOLD: currentSeqLen = ( -// FOLD: prefixOffset = ( -// FOLD: slidingWindowSize = 1 -// FOLD: causal +// FOLD-DAG: currentSeqLen = ( +// FOLD-DAG: prefixOffset = ( +// FOLD-DAG: slidingWindowSize = 1 +// FOLD-DAG: causal // FOLD: qk = elementwise { // FOLD-NOT: tosa.select // FOLD: rock.yield From 5c4f12bb874fcbbaff1471d67aa45fcca50ab044 Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Wed, 29 Jul 2026 16:31:42 +0000 Subject: [PATCH 4/4] [NFC] Remove duplicate sliding-window mask reconciliation Eliminate a branch-update artifact that repeated the same validation block. Co-authored-by: Cursor --- mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp index 4ea64588f676..09f3448ae452 100644 --- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp +++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp @@ -2673,26 +2673,6 @@ struct AttentionRewritePattern : public OpRewritePattern { } } - // Sliding-window masking is defined relative to currentSeqLen. Reconcile - // the validated operand after all masks have been analyzed so the result is - // independent of the select nesting order. - if (currentResult.slidingWindowSize) { - if (currentResult.seqLen) { - if (!sameSeqLenBlockArg(currentResult.seqLen, - currentResult.slidingWindowSeqLen, seqLenSkip)) - return failure(); - // A single attention op cannot represent different clamps for the - // KV-cache and sliding-window masks. - if (currentResult.seqLenClipMin != currentResult.slidingWindowClipMin || - currentResult.seqLenClipMax != currentResult.slidingWindowClipMax) - return failure(); - } else { - currentResult.seqLen = currentResult.slidingWindowSeqLen; - currentResult.seqLenClipMin = currentResult.slidingWindowClipMin; - currentResult.seqLenClipMax = currentResult.slidingWindowClipMax; - } - } - // We need at least one pattern to be detected if (!currentResult.seqLen && !currentResult.prefixOffset && !currentResult.slidingWindowSize)