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..1fd16f9b75da --- /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-DAG: currentSeqLen = ( +// FOLD-DAG: prefixOffset = ( +// FOLD-DAG: slidingWindowSize = 1 +// FOLD-DAG: 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> + } +}