From 8fd152817f2fd4102ebc71b8f772f148118f6cd1 Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Wed, 12 Aug 2026 18:29:19 +0000 Subject: [PATCH 1/5] [Attention] Handle fully masked reference rows Keep CPU and split-KV softmax normalization defined when every logit is masked, and report zero normalized RMS for exact all-zero matches. Co-authored-by: Cursor --- .../conv-validation-wrappers.cpp | 8 ++- mlir/test/rocmlir-driver/verify-all-zero.mlir | 36 ++++++++++ .../rocmlir-gen/attention-kernel-causal.mlir | 4 +- .../rocmlir-gen/attention-kernel-f16.mlir | 4 +- ...ention-kernel-gqa-kvcache-lse-splitkv.mlir | 4 +- .../attention-kernel-gqa-kvcache-lse.mlir | 4 +- .../attention-kernel-gqa-kvcache.mlir | 8 +-- .../rocmlir-gen/attention-kernel-gqa.mlir | 8 +-- mlir/test/rocmlir-gen/attention-kernel.mlir | 8 +-- .../rocmlir-gen/attention-sliding-window.mlir | 19 +++++ .../attention-splitkv-host-f32-combine.mlir | 10 ++- mlir/tools/rocmlir-gen/rocmlir-gen.cpp | 70 +++++++++++++++---- 12 files changed, 143 insertions(+), 40 deletions(-) create mode 100644 mlir/test/rocmlir-driver/verify-all-zero.mlir diff --git a/mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp b/mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp index 5ee12d2805e7..48db27c55b01 100644 --- a/mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp +++ b/mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp @@ -213,8 +213,12 @@ void mcpuVerify(T *gpuResults, T *validationResults, long long dataSize, } double aveAbsDiff = sumAbsDiff / static_cast(dataSize); double aveRelDiff = sumRelDiff / static_cast(dataSize); - double err_RMS = sqrt(sumDiffSq) / (static_cast(maxMag) * - sqrt(static_cast(dataSize))); + // Avoid 0/0 for identical all-zero tensors. Checking the numerator rather + // than maxMag preserves a NaN result for non-finite mismatches. + double err_RMS = sumDiffSq == 0.0 ? 0.0 + : sqrt(sumDiffSq) / + (static_cast(maxMag) * + sqrt(static_cast(dataSize))); // Check if pass based on all three metrics: RMS, maxAbsDiff, maxRelDiff int RMS_pass = (err_RMS <= thr_RMS) ? 1 : 0; int absDiff_pass = (maxAbsDiff <= thr_absDiff) ? 1 : 0; diff --git a/mlir/test/rocmlir-driver/verify-all-zero.mlir b/mlir/test/rocmlir-driver/verify-all-zero.mlir new file mode 100644 index 000000000000..b298fa20015f --- /dev/null +++ b/mlir/test/rocmlir-driver/verify-all-zero.mlir @@ -0,0 +1,36 @@ +// RUN: rocmlir-driver --host-pipeline=runner %s \ +// RUN: | mlir-runner -O2 --shared-libs=%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void \ +// RUN: | FileCheck %s + +// Identical all-zero tensors have zero absolute error and zero scale. Verify +// that normalized RMS treats this exact match as zero instead of computing 0/0. + +// CHECK: [1 1 1] + +module { + func.func @main() { + %gpu = memref.alloc() : memref<4xf32> + %reference = memref.alloc() : memref<4xf32> + %zero = arith.constant 0.0 : f32 + linalg.fill ins(%zero : f32) outs(%gpu : memref<4xf32>) + linalg.fill ins(%zero : f32) outs(%reference : memref<4xf32>) + + %gpuDynamic = memref.cast %gpu : memref<4xf32> to memref + %referenceDynamic = memref.cast %reference : memref<4xf32> to memref + %threshold = arith.constant 0.0 : f32 + %printDebug = arith.constant 0 : i8 + %isFP32 = arith.constant true + %useAbsDiffGate = arith.constant false + call @mcpuVerifyFloat( + %gpuDynamic, %referenceDynamic, %threshold, %threshold, %threshold, + %printDebug, %isFP32, %useAbsDiffGate) + : (memref, memref, f32, f32, f32, i8, i1, i1) -> () + + memref.dealloc %gpu : memref<4xf32> + memref.dealloc %reference : memref<4xf32> + return + } + + func.func private @mcpuVerifyFloat(memref, memref, f32, f32, + f32, i8, i1, i1) +} diff --git a/mlir/test/rocmlir-gen/attention-kernel-causal.mlir b/mlir/test/rocmlir-gen/attention-kernel-causal.mlir index 923fa962c843..e5c6fdf4b180 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-causal.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-causal.mlir @@ -53,10 +53,10 @@ // CHECK: %[[qkTensor:.*]] = tosa.select %[[mask3]], %[[negInf]], %[[sqkTensorCast]] : (tensor<1x512x1024xi1>, tensor<1x512x1024xf32>, tensor<1x512x1024xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : (tensor<1x512x1024xf32>) -> tensor<1x512x1xf32> -// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %[[sqkMaxs]] : (tensor<1x512x1024xf32>, tensor<1x512x1xf32>) -> tensor<1x512x1024xf32> +// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : (tensor<1x512x1024xf32>, tensor<1x512x1xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : (tensor<1x512x1024xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : (tensor<1x512x1024xf32>) -> tensor<1x512x1xf32> -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : (tensor<1x512x1xf32>) -> tensor<1x512x1xf32> +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : (tensor<1x512x1xf32>) -> tensor<1x512x1xf32> // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : (tensor<1x512x1024xf32>, tensor<1x512x1xf32>, tensor<1xi8>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : (tensor<1x512x1024xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : (tensor<1x512x1024xf32>, tensor<1x1024x32xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x512x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-f16.mlir b/mlir/test/rocmlir-gen/attention-kernel-f16.mlir index 5d253d8149c1..1ff4a51a22a7 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-f16.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-f16.mlir @@ -38,10 +38,10 @@ // CHECK-DAG: %[[sqkTensor:.*]] = tosa.mul %[[qkTensor]], %[[scaleTensor:.*]], %{{.*}} : ([[squareShape]], [[squareShape]], tensor<1xi8>) -> [[squareShape]] // CHECK-DAG: %[[sqkTensorCast:.*]] = tosa.cast %[[sqkTensor]] : ([[squareShape]]) -> [[squareShapeF32:tensor<.*>]] // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[sqkTensorCast]] {{.*}} : ([[squareShapeF32]]) -> [[reducedShape:tensor<.*>]] -// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %[[sqkMaxs]] : ([[squareShapeF32]], [[reducedShape]]) -> [[squareShapeF32]] +// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %{{.*}} : ([[squareShapeF32]], [[reducedShape]]) -> [[squareShapeF32]] // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShapeF32]]) -> [[squareShapeF32]] // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShapeF32]]) -> [[reducedShape]] -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShapeF32]], [[reducedShape]], tensor<1xi8>) -> [[squareShapeF32]] // CHECK-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShapeF32]]) -> [[squareShape]] // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} {acc_type = f32} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf16>, tensor<1xf16>) -> [[squareShape:tensor<.*>]] diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir index 602ba6c910d3..f6015ca01685 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir @@ -54,7 +54,7 @@ // CHECK: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1x1024xf32>, !tosa.shape<3>) -> tensor<4x1x1024xf32> // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] @@ -63,7 +63,7 @@ // CHECK-DAG: %[[logL:.*]] = tosa.log %[[expsSumsTensorCasted]] : (tensor<4x1x1xf32>) -> tensor<4x1x1xf32> // CHECK-DAG: %[[resultLse:.*]] = tosa.add %[[logL]], %[[sqkMaxsCasted]] : (tensor<4x1x1xf32>, tensor<4x1x1xf32>) -> tensor<4x1x1xf32> -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK-DAG: %[[softmaxTensorCasted:.*]] = tosa.cast %[[softmaxTensor]] : (tensor<4x1x1024xf32>) -> tensor<4x1x1024xf32> // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCasted]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} {acc_type = f32} : (tensor<4x1x1024xf32>, tensor<4x1024x32xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x1x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir index c06bffc852c9..8db886fb11d9 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir @@ -53,7 +53,7 @@ // CHECK-DAG: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1024x1024xf32>, !tosa.shape<3>) -> tensor<4x1024x1024xf32> // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] @@ -62,7 +62,7 @@ // CHECK-DAG: %[[logL:.*]] = tosa.log %[[expsSumsTensorCast]] : (tensor<4x1024x1xf32>) -> tensor<4x1024x1xf32> // CHECK-DAG: %[[resultLse:.*]] = tosa.add %[[logL]], %[[sqkMaxsCast]] : (tensor<4x1024x1xf32>, tensor<4x1024x1xf32>) -> tensor<4x1024x1xf32> -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> tensor<4x1024x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir index 208c00f7627a..eea4debe77b5 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir @@ -68,10 +68,10 @@ // CHECK_SCALE: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1024x1024xf32>, !tosa.shape<3>) -> tensor<4x1024x1024xf32> // CHECK_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape]], tensor<1xf32>, tensor<1xf32>) -> tensor<4x1024x32xf32> @@ -131,10 +131,10 @@ // CHECK_NO_SCALE: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1024x1024xf32>, !tosa.shape<3>) -> tensor<4x1024x1024xf32> // CHECK_NO_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> tensor<4x1024x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir index d040060e1571..55327150eab5 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir @@ -33,10 +33,10 @@ // CHECK_SCALE-DAG: %[[sqkTensor:.*]] = tosa.mul %[[qkTensor]], %[[scaleTensor:.*]], %{{.*}} : ([[squareShape]], [[squareShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkTensorCast:.*]] = tosa.cast %[[sqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[sqkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] @@ -75,10 +75,10 @@ // CHECK_NO_SCALE: %[[qkTensor:.*]] = tosa.matmul %[[queriesTensor:.*]], %[[keysTensor:.*]], %{{.*}}, %{{.*}} : ([[queriesShape:tensor<.*>]], [[keysShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[squareShape:tensor<.*>]] // CHECK_NO_SCALE: %[[qkTensorCast:.*]] = tosa.cast %[[qkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] diff --git a/mlir/test/rocmlir-gen/attention-kernel.mlir b/mlir/test/rocmlir-gen/attention-kernel.mlir index f447919f2255..a809757077ae 100644 --- a/mlir/test/rocmlir-gen/attention-kernel.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel.mlir @@ -26,10 +26,10 @@ // CHECK_SCALE-DAG: %[[sqkTensor:.*]] = tosa.mul %[[qkTensor]], %[[scaleTensor:.*]], %{{.*}} : ([[squareShape]], [[squareShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkTensorCast:.*]] = tosa.cast %[[sqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[sqkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] @@ -61,10 +61,10 @@ // CHECK_NO_SCALE: %[[qkTensor:.*]] = tosa.matmul %[[queriesTensor:.*]], %[[keysTensor:.*]], %{{.*}}, %{{.*}} : ([[queriesShape:tensor<.*>]], [[keysShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[squareShape:tensor<.*>]] // CHECK_NO_SCALE: %[[qkTensorCast:.*]] = tosa.cast %[[qkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[expsSumsTensor]] : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] diff --git a/mlir/test/rocmlir-gen/attention-sliding-window.mlir b/mlir/test/rocmlir-gen/attention-sliding-window.mlir index 58ca50463135..27348f354321 100644 --- a/mlir/test/rocmlir-gen/attention-sliding-window.mlir +++ b/mlir/test/rocmlir-gen/attention-sliding-window.mlir @@ -1,4 +1,5 @@ // RUN: rocmlir-gen --arch gfx90a:sramecc+:xnack- --operation attention -current_seq_len=33 -sliding_window_size=16 -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -pv --apply-bufferization-pipeline=false | rocmlir-opt | FileCheck %s --enable-var-scope +// RUN: rocmlir-gen --arch gfx90a:sramecc+:xnack- --operation attention -current_seq_len=2 -sliding_window_size=1 --causal -return_lse -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -pv --apply-bufferization-pipeline=false | rocmlir-opt | FileCheck %s --enable-var-scope --check-prefix=SAFE // CHECK: module attributes {mhal.arch = "[[$ARCH:.*]]"} @@ -38,3 +39,21 @@ // CHECK-DAG: tosa.reciprocal // CHECK: tosa.matmul // CHECK: return + +// A fully masked row has max=-inf and exp-sum=0. Verify the CPU reference uses +// finite normalization operands while retaining the original values for LSE. +// SAFE-LABEL: func.func @host_naive_attention +// SAFE: %[[MAX:.*]] = tosa.reduce_max +// SAFE: %[[LOWEST:.*]] = "tosa.const"() <{values = dense<-3.40282347E+38> : tensor<1x1x1xf32>}> : () -> tensor<1x1x1xf32> +// SAFE: %[[SAFE_MAX:.*]] = tosa.maximum %[[MAX]], %[[LOWEST]] +// SAFE: %[[NORMALIZED:.*]] = tosa.sub %{{.*}}, %[[SAFE_MAX]] +// SAFE: %[[EXPS:.*]] = tosa.exp %[[NORMALIZED]] +// SAFE: %[[SUM:.*]] = tosa.reduce_sum %[[EXPS]] +// SAFE: %[[MAX_FOR_LSE:.*]] = tosa.cast %[[MAX]] +// SAFE: %[[LOG_SUM:.*]] = tosa.log +// SAFE: tosa.add %[[LOG_SUM]], %[[MAX_FOR_LSE]] +// SAFE: %[[ONE:.*]] = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x1x1xf32>}> : () -> tensor<1x1x1xf32> +// SAFE: %[[SAFE_SUM:.*]] = tosa.maximum %[[SUM]], %[[ONE]] +// SAFE: tosa.reciprocal %[[SAFE_SUM]] +// SAFE: tosa.matmul +// SAFE: return diff --git a/mlir/test/rocmlir-gen/attention-splitkv-host-f32-combine.mlir b/mlir/test/rocmlir-gen/attention-splitkv-host-f32-combine.mlir index 083afcad7b42..0fc4060a3258 100644 --- a/mlir/test/rocmlir-gen/attention-splitkv-host-f32-combine.mlir +++ b/mlir/test/rocmlir-gen/attention-splitkv-host-f32-combine.mlir @@ -23,10 +23,14 @@ // LSE re-normalization across the splitKV axis (axis = 1), entirely in f32. // CHECK: %[[mx:.+]] = tosa.reduce_max %{{.+}} {axis = 1 : i32} : (tensor<4x8x1x1xf32>) -> tensor<4x1x1x1xf32> -// CHECK: tosa.sub %{{.+}}, %[[mx]] : (tensor<4x8x1x1xf32>, tensor<4x1x1x1xf32>) -> tensor<4x8x1x1xf32> +// CHECK: %[[lowest:.+]] = "tosa.const"() <{values = dense<-3.40282347E+38> : tensor<4x1x1x1xf32>}> : () -> tensor<4x1x1x1xf32> +// CHECK: %[[safeMx:.+]] = tosa.maximum %[[mx]], %[[lowest]] : (tensor<4x1x1x1xf32>, tensor<4x1x1x1xf32>) -> tensor<4x1x1x1xf32> +// CHECK: tosa.sub %{{.+}}, %[[safeMx]] : (tensor<4x8x1x1xf32>, tensor<4x1x1x1xf32>) -> tensor<4x8x1x1xf32> // CHECK: tosa.exp %{{.+}} : (tensor<4x8x1x1xf32>) -> tensor<4x8x1x1xf32> -// CHECK: tosa.reduce_sum %{{.+}} {axis = 1 : i32} : (tensor<4x8x1x1xf32>) -> tensor<4x1x1x1xf32> -// CHECK: tosa.reciprocal %{{.+}} : (tensor<4x1x1x1xf32>) -> tensor<4x1x1x1xf32> +// CHECK: %[[sum:.+]] = tosa.reduce_sum %{{.+}} {axis = 1 : i32} : (tensor<4x8x1x1xf32>) -> tensor<4x1x1x1xf32> +// CHECK: %[[one:.+]] = "tosa.const"() <{values = dense<1.000000e+00> : tensor<4x1x1x1xf32>}> : () -> tensor<4x1x1x1xf32> +// CHECK: %[[safeSum:.+]] = tosa.maximum %[[sum]], %[[one]] : (tensor<4x1x1x1xf32>, tensor<4x1x1x1xf32>) -> tensor<4x1x1x1xf32> +// CHECK: tosa.reciprocal %[[safeSum]] : (tensor<4x1x1x1xf32>) -> tensor<4x1x1x1xf32> // Weighted sum of the partial outputs across splits, still in f32. // CHECK: tosa.reduce_sum %{{.+}} {axis = 1 : i32} : (tensor<4x8x1x32xf32>) -> tensor<4x1x1x32xf32> diff --git a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp index 58ddafbad1cd..0579c6bf6e9a 100644 --- a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp +++ b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp @@ -2979,24 +2979,30 @@ static Value addTensorArgToBlock(OpBuilder &builder, Location loc, return funcArgTensor; } -static Value applyMask(OpBuilder builder, Location loc, Value inputTensor, - Value mask, float initValue) { - auto inpType = cast(inputTensor.getType()); - ArrayRef inpShape = inpType.getShape(); +static Value createFloatSplatTensor(OpBuilder builder, Location loc, + RankedTensorType type, + const APFloat &value) { + assert(isa(type.getElementType())); + DenseElementsAttr valueAttr = DenseFPElementsAttr::get(type, value); + return tosa::ConstOp::create(builder, loc, valueAttr.getType(), valueAttr); +} - // create a tensor with a single value and broadcast it - assert(isa(inpType.getElementType())); +static Value createFloatSplatTensor(OpBuilder builder, Location loc, + RankedTensorType type, float value) { std::pair floatRes = - rock::createAPFloat(inpType.getElementType(), initValue); + rock::createAPFloat(type.getElementType(), value); APFloat fpVal = floatRes.first; auto status = floatRes.second; assert(status == APFloat::opOK); - DenseElementsAttr initValueAttr = DenseFPElementsAttr::get( - RankedTensorType::get(inpShape, inpType.getElementType()), fpVal); + return createFloatSplatTensor(builder, loc, type, fpVal); +} - Value initVal = tosa::ConstOp::create(builder, loc, initValueAttr.getType(), - initValueAttr); +static Value applyMask(OpBuilder builder, Location loc, Value inputTensor, + Value mask, float initValue) { + auto inpType = cast(inputTensor.getType()); + + Value initVal = createFloatSplatTensor(builder, loc, inpType, initValue); // mask is 1 for values we want to set to "initVal" auto result = rock::tosa::createOpAndInfer( @@ -3359,15 +3365,29 @@ static Value computeFinalAttentionStage(OpBuilder builder, Location loc, auto maxSplitKV = rock::tosa::createOpAndInfer( builder, loc, computeType, lseTensor, axisAttr); + auto maxSplitKVType = cast(maxSplitKV.getType()); + APFloat lowestFinite = APFloat::getLargest( + cast(computeType).getFloatSemantics(), /*Negative=*/true); + Value lowestFiniteTensor = + createFloatSplatTensor(builder, loc, maxSplitKVType, lowestFinite); + Value maxSplitKVForNormalization = + rock::tosa::createOpAndInfer( + builder, loc, computeType, maxSplitKV, lowestFiniteTensor); + auto norm = rock::tosa::createOpAndInfer( - builder, loc, computeType, lseTensor, maxSplitKV); + builder, loc, computeType, lseTensor, maxSplitKVForNormalization); auto exp = rock::tosa::createOpAndInfer(builder, loc, computeType, norm); auto sumExpNorm = rock::tosa::createOpAndInfer( builder, loc, computeType, exp, axisAttr); + auto sumExpNormType = cast(sumExpNorm.getType()); + Value oneTensor = createFloatSplatTensor(builder, loc, sumExpNormType, 1.0f); + Value sumExpNormForNormalization = + rock::tosa::createOpAndInfer(builder, loc, computeType, + sumExpNorm, oneTensor); auto sumExpNormRecip = rock::tosa::createOpAndInfer( - builder, loc, computeType, sumExpNorm); + builder, loc, computeType, sumExpNormForNormalization); Value outExp = rock::tosa::getMulOp(builder, loc, exp, resultTensor, computeType); @@ -4552,8 +4572,20 @@ static func::FuncOp createCpuAttentionKernelWithMlir(ModuleOp module, constexpr int64_t reductionAxis = 2; auto qkMaxs = rock::tosa::createOpAndInfer( builder, loc, softmaxType, qkTensor, reductionAxis); + + // A fully masked row reduces to -inf. Clamp the normalization max to the + // lowest finite value so masked scores remain -inf instead of producing + // -inf - (-inf) = NaN. This leaves every finite row maximum unchanged. + auto qkMaxsType = cast(qkMaxs.getType()); + APFloat lowestFinite = APFloat::getLargest( + cast(softmaxType).getFloatSemantics(), /*Negative=*/true); + Value lowestFiniteTensor = + createFloatSplatTensor(builder, loc, qkMaxsType, lowestFinite); + Value qkMaxsForNormalization = rock::tosa::createOpAndInfer( + builder, loc, softmaxType, qkMaxs, lowestFiniteTensor); + auto normalizedQkTensor = rock::tosa::createOpAndInfer( - builder, loc, softmaxType, qkTensor, qkMaxs); + builder, loc, softmaxType, qkTensor, qkMaxsForNormalization); auto expsTensor = rock::tosa::createOpAndInfer( builder, loc, softmaxType, normalizedQkTensor); auto expsSums = rock::tosa::createOpAndInfer( @@ -4576,8 +4608,16 @@ static func::FuncOp createCpuAttentionKernelWithMlir(ModuleOp module, builder, loc, lseType, lseTensor, qkMaxsForLSE); } + // A valid row contains exp(max - max) = 1, so its sum is at least one. + // Fully masked rows sum to zero; use one as their denominator to produce a + // zero softmax row. Keep the original zero sum for LSE so it becomes -inf. + auto expsSumsType = cast(expsSums.getType()); + Value oneTensor = createFloatSplatTensor(builder, loc, expsSumsType, 1.0f); + Value expsSumsForNormalization = + rock::tosa::createOpAndInfer(builder, loc, softmaxType, + expsSums, oneTensor); auto invExpsSums = rock::tosa::createOpAndInfer( - builder, loc, softmaxType, expsSums); + builder, loc, softmaxType, expsSumsForNormalization); Value softmaxTensor = rock::tosa::getMulOp(builder, loc, expsTensor, invExpsSums, softmaxType); From 174c5e2ef3459d968667e5738a6031dce174630c Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Wed, 12 Aug 2026 19:35:00 +0000 Subject: [PATCH 2/5] [Attention] Strengthen masked-row regression checks Preserve the safe-softmax def-use checks, gate the verifier test on runner support, and document helper preconditions. Co-authored-by: Cursor --- mlir/test/rocmlir-driver/verify-all-zero.mlir | 1 + mlir/test/rocmlir-gen/attention-kernel-causal.mlir | 6 ++++-- mlir/test/rocmlir-gen/attention-kernel-f16.mlir | 6 ++++-- .../attention-kernel-gqa-kvcache-lse-splitkv.mlir | 6 ++++-- .../attention-kernel-gqa-kvcache-lse.mlir | 6 ++++-- .../rocmlir-gen/attention-kernel-gqa-kvcache.mlir | 12 ++++++++---- mlir/test/rocmlir-gen/attention-kernel-gqa.mlir | 12 ++++++++---- mlir/test/rocmlir-gen/attention-kernel.mlir | 12 ++++++++---- mlir/tools/rocmlir-gen/rocmlir-gen.cpp | 6 ++++-- 9 files changed, 45 insertions(+), 22 deletions(-) diff --git a/mlir/test/rocmlir-driver/verify-all-zero.mlir b/mlir/test/rocmlir-driver/verify-all-zero.mlir index b298fa20015f..b4bc669c54a6 100644 --- a/mlir/test/rocmlir-driver/verify-all-zero.mlir +++ b/mlir/test/rocmlir-driver/verify-all-zero.mlir @@ -1,3 +1,4 @@ +// REQUIRES: rocm-runner // RUN: rocmlir-driver --host-pipeline=runner %s \ // RUN: | mlir-runner -O2 --shared-libs=%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void \ // RUN: | FileCheck %s diff --git a/mlir/test/rocmlir-gen/attention-kernel-causal.mlir b/mlir/test/rocmlir-gen/attention-kernel-causal.mlir index e5c6fdf4b180..ba2dab5430cb 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-causal.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-causal.mlir @@ -53,10 +53,12 @@ // CHECK: %[[qkTensor:.*]] = tosa.select %[[mask3]], %[[negInf]], %[[sqkTensorCast]] : (tensor<1x512x1024xi1>, tensor<1x512x1024xf32>, tensor<1x512x1024xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : (tensor<1x512x1024xf32>) -> tensor<1x512x1xf32> -// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : (tensor<1x512x1024xf32>, tensor<1x512x1xf32>) -> tensor<1x512x1024xf32> +// CHECK-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : (tensor<1x512x1xf32>, tensor<1x512x1xf32>) -> tensor<1x512x1xf32> +// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %[[safeSqkMaxs]] : (tensor<1x512x1024xf32>, tensor<1x512x1xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : (tensor<1x512x1024xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : (tensor<1x512x1024xf32>) -> tensor<1x512x1xf32> -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : (tensor<1x512x1xf32>) -> tensor<1x512x1xf32> +// CHECK-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : (tensor<1x512x1xf32>, tensor<1x512x1xf32>) -> tensor<1x512x1xf32> +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : (tensor<1x512x1xf32>) -> tensor<1x512x1xf32> // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : (tensor<1x512x1024xf32>, tensor<1x512x1xf32>, tensor<1xi8>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : (tensor<1x512x1024xf32>) -> tensor<1x512x1024xf32> // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : (tensor<1x512x1024xf32>, tensor<1x1024x32xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x512x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-f16.mlir b/mlir/test/rocmlir-gen/attention-kernel-f16.mlir index 1ff4a51a22a7..a306afaddd14 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-f16.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-f16.mlir @@ -38,10 +38,12 @@ // CHECK-DAG: %[[sqkTensor:.*]] = tosa.mul %[[qkTensor]], %[[scaleTensor:.*]], %{{.*}} : ([[squareShape]], [[squareShape]], tensor<1xi8>) -> [[squareShape]] // CHECK-DAG: %[[sqkTensorCast:.*]] = tosa.cast %[[sqkTensor]] : ([[squareShape]]) -> [[squareShapeF32:tensor<.*>]] // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[sqkTensorCast]] {{.*}} : ([[squareShapeF32]]) -> [[reducedShape:tensor<.*>]] -// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %{{.*}} : ([[squareShapeF32]], [[reducedShape]]) -> [[squareShapeF32]] +// CHECK-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %[[safeSqkMaxs]] : ([[squareShapeF32]], [[reducedShape]]) -> [[squareShapeF32]] // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShapeF32]]) -> [[squareShapeF32]] // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShapeF32]]) -> [[reducedShape]] -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShapeF32]], [[reducedShape]], tensor<1xi8>) -> [[squareShapeF32]] // CHECK-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShapeF32]]) -> [[squareShape]] // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} {acc_type = f32} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf16>, tensor<1xf16>) -> [[squareShape:tensor<.*>]] diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir index f6015ca01685..7326a2c26ac8 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir @@ -54,7 +54,8 @@ // CHECK: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1x1024xf32>, !tosa.shape<3>) -> tensor<4x1x1024xf32> // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] @@ -63,7 +64,8 @@ // CHECK-DAG: %[[logL:.*]] = tosa.log %[[expsSumsTensorCasted]] : (tensor<4x1x1xf32>) -> tensor<4x1x1xf32> // CHECK-DAG: %[[resultLse:.*]] = tosa.add %[[logL]], %[[sqkMaxsCasted]] : (tensor<4x1x1xf32>, tensor<4x1x1xf32>) -> tensor<4x1x1xf32> -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK-DAG: %[[softmaxTensorCasted:.*]] = tosa.cast %[[softmaxTensor]] : (tensor<4x1x1024xf32>) -> tensor<4x1x1024xf32> // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCasted]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} {acc_type = f32} : (tensor<4x1x1024xf32>, tensor<4x1024x32xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x1x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir index 8db886fb11d9..7a8cb30d0439 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir @@ -53,7 +53,8 @@ // CHECK-DAG: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1024x1024xf32>, !tosa.shape<3>) -> tensor<4x1024x1024xf32> // CHECK-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] @@ -62,7 +63,8 @@ // CHECK-DAG: %[[logL:.*]] = tosa.log %[[expsSumsTensorCast]] : (tensor<4x1024x1xf32>) -> tensor<4x1024x1xf32> // CHECK-DAG: %[[resultLse:.*]] = tosa.add %[[logL]], %[[sqkMaxsCast]] : (tensor<4x1024x1xf32>, tensor<4x1024x1xf32>) -> tensor<4x1024x1xf32> -// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> tensor<4x1024x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir index eea4debe77b5..8898fed0f176 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir @@ -68,10 +68,12 @@ // CHECK_SCALE: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1024x1024xf32>, !tosa.shape<3>) -> tensor<4x1024x1024xf32> // CHECK_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_SCALE-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[qkTensor]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape]], tensor<1xf32>, tensor<1xf32>) -> tensor<4x1024x32xf32> @@ -131,10 +133,12 @@ // CHECK_NO_SCALE: %[[qkTensor:.*]] = tosa.reshape %[[qkTensorBeforeReshape]], %{{.*}} : (tensor<1x4x1024x1024xf32>, !tosa.shape<3>) -> tensor<4x1024x1024xf32> // CHECK_NO_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_NO_SCALE-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensor]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> tensor<4x1024x32xf32> diff --git a/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir b/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir index 55327150eab5..04851d320411 100644 --- a/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel-gqa.mlir @@ -33,10 +33,12 @@ // CHECK_SCALE-DAG: %[[sqkTensor:.*]] = tosa.mul %[[qkTensor]], %[[scaleTensor:.*]], %{{.*}} : ([[squareShape]], [[squareShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkTensorCast:.*]] = tosa.cast %[[sqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[sqkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_SCALE-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] @@ -75,10 +77,12 @@ // CHECK_NO_SCALE: %[[qkTensor:.*]] = tosa.matmul %[[queriesTensor:.*]], %[[keysTensor:.*]], %{{.*}}, %{{.*}} : ([[queriesShape:tensor<.*>]], [[keysShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[squareShape:tensor<.*>]] // CHECK_NO_SCALE: %[[qkTensorCast:.*]] = tosa.cast %[[qkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_NO_SCALE-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] diff --git a/mlir/test/rocmlir-gen/attention-kernel.mlir b/mlir/test/rocmlir-gen/attention-kernel.mlir index a809757077ae..976016e851ff 100644 --- a/mlir/test/rocmlir-gen/attention-kernel.mlir +++ b/mlir/test/rocmlir-gen/attention-kernel.mlir @@ -26,10 +26,12 @@ // CHECK_SCALE-DAG: %[[sqkTensor:.*]] = tosa.mul %[[qkTensor]], %[[scaleTensor:.*]], %{{.*}} : ([[squareShape]], [[squareShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkTensorCast:.*]] = tosa.cast %[[sqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[sqkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_SCALE-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedSqkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] @@ -61,10 +63,12 @@ // CHECK_NO_SCALE: %[[qkTensor:.*]] = tosa.matmul %[[queriesTensor:.*]], %[[keysTensor:.*]], %{{.*}}, %{{.*}} : ([[queriesShape:tensor<.*>]], [[keysShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[squareShape:tensor<.*>]] // CHECK_NO_SCALE: %[[qkTensorCast:.*]] = tosa.cast %[[qkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[qkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]] -// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] +// CHECK_NO_SCALE-DAG: %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[normilizedQkTensor:.*]] = tosa.sub %[[qkTensorCast]], %[[safeSqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsTensor:.*]] = tosa.exp %[[normilizedQkTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[expsSumsTensor:.*]] = tosa.reduce_sum %[[expsTensor]] {{.*}} : ([[squareShape]]) -> [[reducedShape]] -// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %{{.*}} : ([[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[safeExpsSums:.*]] = tosa.maximum %[[expsSumsTensor]], %{{.*}} : ([[reducedShape]], [[reducedShape]]) -> [[reducedShape]] +// CHECK_NO_SCALE-DAG: %[[invExpsSums:.*]] = tosa.reciprocal %[[safeExpsSums]] : ([[reducedShape]]) -> [[reducedShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensor:.*]] = tosa.mul %[[expsTensor]], %[[invExpsSums]], %{{.*}} : ([[squareShape]], [[reducedShape]], tensor<1xi8>) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[softmaxTensorCast:.*]] = tosa.cast %[[softmaxTensor]] : ([[squareShape]]) -> [[squareShape]] // CHECK_NO_SCALE-DAG: %[[resultTensor:.*]] = tosa.matmul %[[softmaxTensorCast]], %[[valuesTensor:.*]], %{{.*}}, %{{.*}} : ([[squareShape]], [[valuesShape:tensor<.*>]], tensor<1xf32>, tensor<1xf32>) -> [[valuesShape]] diff --git a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp index 0579c6bf6e9a..6a711cc44905 100644 --- a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp +++ b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp @@ -2982,7 +2982,8 @@ static Value addTensorArgToBlock(OpBuilder &builder, Location loc, static Value createFloatSplatTensor(OpBuilder builder, Location loc, RankedTensorType type, const APFloat &value) { - assert(isa(type.getElementType())); + assert(isa(type.getElementType()) && + "expected a float element type"); DenseElementsAttr valueAttr = DenseFPElementsAttr::get(type, value); return tosa::ConstOp::create(builder, loc, valueAttr.getType(), valueAttr); } @@ -2993,7 +2994,8 @@ static Value createFloatSplatTensor(OpBuilder builder, Location loc, rock::createAPFloat(type.getElementType(), value); APFloat fpVal = floatRes.first; auto status = floatRes.second; - assert(status == APFloat::opOK); + assert(status == APFloat::opOK && + "failed to create exact floating-point splat value"); return createFloatSplatTensor(builder, loc, type, fpVal); } From c0e1fbded9842605d5bdd6ecb16d2885d0fc033b Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Wed, 12 Aug 2026 19:41:58 +0000 Subject: [PATCH 3/5] [Attention] Add fully masked E2E coverage Exercise direct attention/LSE and split-KV recombination when a causal sliding-window row has no eligible keys. Co-authored-by: Cursor --- .../rock-attention-fully-masked.mlir | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir diff --git a/mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir b/mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir new file mode 100644 index 000000000000..de1a730c8efa --- /dev/null +++ b/mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir @@ -0,0 +1,19 @@ +// RUN: rocmlir-gen --arch %arch --operation attention -current_seq_len=2 -sliding_window_size=1 --causal -return_lse -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -rand 1 -rand_type float -pv \ +// RUN: | rocmlir-driver --host-pipeline=highlevel \ +// RUN: | rocmlir-driver -c \ +// RUN: | mlir-runner -O2 --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 --entry-point-result=void \ +// RUN: | FileCheck %s --check-prefix=DIRECT +// RUN: rocmlir-gen --arch %arch --operation attention -current_seq_len=2 -sliding_window_size=1 --causal -return_lse -split_kv=8 -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -rand 1 -rand_type float -pv \ +// RUN: | rocmlir-driver --host-pipeline=highlevel \ +// RUN: | rocmlir-driver -c \ +// RUN: | mlir-runner -O2 --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 --entry-point-result=void \ +// RUN: | FileCheck %s --check-prefix=SPLITKV + +// A causal sliding window can leave a query with no eligible keys. Both the +// GPU kernel and the CPU reference define that row's contribution as zero. +// The direct path also returns the fully masked row's -inf LSE; the split-KV +// path exercises host recombination of fully masked partial results. + +// DIRECT: [1 1 1] +// DIRECT-NEXT: [1 1 1] +// SPLITKV: [1 1 1] From 31ec9b3d6c925da01e3ac016511988aaaed1487f Mon Sep 17 00:00:00 2001 From: Umang Yadav <29876643+umangyadav@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:48:06 -0400 Subject: [PATCH 4/5] [Rock] Guard inactive blockwise reduction threads (#2447) * [Rock] Guard inactive blockwise reduction threads The NR-small LDS tree reduction derives per-thread LDS coordinates from tid without checking that the thread owns a reduction slice. When blockSize is not divisible by the product of the non-reduction dims, the surplus workitems compute coordinates that either alias a valid thread's slice or run past the logical workspace. The shortfall grows when rthreads is further clamped to divide the reduction dimension, which is the shape the attention configs hit. Because no barrier separates the reduction load loop from the result store, an idle workitem in one wave can double-count a partial that a valid workitem in another wave has already written. Guard both threadwise loops on tid < rthreads * nrDimProd, keeping the barriers workgroup-uniform so the whole workgroup still reaches them. The cross-lane DPP, permlane and ds_swizzle fast paths already require exact thread packing, so the guard is only emitted for the LDS fallback. Covered by two lit sections: one where blockSize leaves the remainder idle, and one where the rthreads clamp does. Both fail without the guard. Negative EXACT checks assert that exact-packing configs stay guard-free. Co-authored-by: Cursor * [Rock] Strengthen inactive reduction test coverage Clarify the clamped-thread hazard and protect another exact-packing fallback from unnecessary guards. Co-authored-by: Cursor * [Rock] Clarify inactive reduction test comment Document that the barrier-placement check remains valid independently of guard canonicalization. Co-authored-by: Cursor * [Rock] Harden inactive reduction thread guard Make the DPP factoring dependency explicit, avoid unused guard IR, and keep unrelated builder behavior unchanged. Co-authored-by: Cursor * [Rock] Generalize inactive reduction test comments Keep the regression documentation focused on reusable test behavior rather than the motivating attention configurations. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- .../Transforms/BlockwiseGemmToThreadwise.cpp | 102 ++++++++++++------ .../lowering_blockwise_broadcast_reduce.mlir | 74 +++++++++++++ 2 files changed, 145 insertions(+), 31 deletions(-) diff --git a/mlir/lib/Dialect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp b/mlir/lib/Dialect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp index c1f6c19ce34f..fecd528292cd 100644 --- a/mlir/lib/Dialect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp +++ b/mlir/lib/Dialect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp @@ -1096,13 +1096,10 @@ struct BlockwiseReduceRewritePattern // This should only be used if product non-reduction dims is // less than number threads in a block. // - // Given a input tensor : D0, ... , Dr , ... , DN to reduce, - // This function creates a view that maps the space of - // [D0, ... , Dr , ... , DN] --> [nrtid, rtid, rIter] where - // nrtid = tid / product(non-reduction dims) is a reduction subgroup leader. - // rtid = tid % product(non-reduction dims) is thread idx within a reduction - // subgroup. Size of the dimension 'rtid' is the number of threads - // that'd participate in the reduction + // Given an input tensor D0, ... , Dr , ... , DN to reduce, this function + // creates a view with upper coordinates [nrtid, rtid, rIter]. nrtid selects + // a non-reduction point, while rtid and rIter select a reduction participant + // and its elements. The caller chooses how to factor tid into nrtid and rtid. ArrayAttr createThreadViewforNRSmallerThanThreads( Location loc, ArrayRef toReduceShape, int64_t blockSize, size_t reduceAxis, PatternRewriter &rewriter) const { @@ -1949,6 +1946,45 @@ struct BlockwiseReduceRewritePattern canUsePermlaneSwap_NRSmall); assert(!canUseDsSwizzleBpermute_NRSmall_LdsSkip || canUseDsSwizzleBpermute_NRSmall); + + // The tree layout can leave threads at the end of the block idle + // when blockSize is not divisible by the number of non-reduction + // elements; the shortfall grows when rthreads is further reduced to + // divide the reduction dimension. Fold the valid (rtid, nrtid) + // rectangle to a linear thread bound. Without the guard, idle + // workitems can race valid LDS updates through aliased coordinates + // or access beyond the logical workspace. Barriers remain outside so + // the whole workgroup reaches them. + int64_t activeReductionThreadCount = + maxActiveReductionThreads * nonReductionDimSizeProduct; + assert(activeReductionThreadCount <= blockSize && + "active reduction threads must fit in the block"); + bool hasInactiveReductionThreads = + activeReductionThreadCount < blockSize; + assert(!(hasInactiveReductionThreads && canUseDPP) && + "linear active-thread bound assumes the tree tid factoring"); + Value isActiveReductionThread; + auto emitForActiveReductionThread = [&](auto &&emit) { + if (!hasInactiveReductionThreads) { + emit(rewriter); + return; + } + if (!isActiveReductionThread) { + Value activeReductionThreadCountVal = + arith::ConstantIndexOp::create(rewriter, loc, + activeReductionThreadCount); + isActiveReductionThread = arith::CmpIOp::create( + rewriter, loc, arith::CmpIPredicate::ult, tid, + activeReductionThreadCountVal); + } + scf::IfOp ifActive = + scf::IfOp::create(rewriter, loc, isActiveReductionThread, + /*withElseRegion=*/false); + OpBuilder thenBuilder = + ifActive.getThenBodyBuilder(rewriter.getListener()); + emit(thenBuilder); + }; + // Two different tid → (rtid, nrtid) factorings are used: // // DPP path: rtid = tid % clusterSize, nrtid = tid / clusterSize. @@ -2041,25 +2077,27 @@ struct BlockwiseReduceRewritePattern Value initVal = getReductionInitValue(op, rewriter); FillOp::create(rewriter, loc, accReg, initVal); - TransformingForOp reductionLoop = TransformingForOp::create( - rewriter, loc, ArrayRef(inits), - ArrayRef{threadToLDSViewTrs}, - ArrayRef(bounds), ArrayRef(strides), - /*forceUnroll=*/true, /*useIndexDiffs=*/true); - { - PatternRewriter::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(reductionLoop.getBody()); + auto emitThreadwiseReduction = [&](OpBuilder &builder) { + TransformingForOp reductionLoop = TransformingForOp::create( + builder, loc, ArrayRef(inits), + ArrayRef{threadToLDSViewTrs}, + ArrayRef(bounds), ArrayRef(strides), + /*forceUnroll=*/true, /*useIndexDiffs=*/true); + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(reductionLoop.getBody()); Block::BlockArgListType LDSLoadCoords = reductionLoop.getLowerCoords(/*domain=*/0); Value loadVal = - InBoundsLoadOp::create(rewriter, loc, loadTypeInputReg, + InBoundsLoadOp::create(builder, loc, loadTypeInputReg, workspaceLDSBuffer, LDSLoadCoords); - Value loadAcc = InBoundsLoadOp::create(rewriter, loc, elemType, + Value loadAcc = InBoundsLoadOp::create(builder, loc, elemType, accReg, zeroConstantOp); - Value reduced = createReducingOp(op, loadVal, loadAcc, rewriter); - InBoundsStoreOp::create(rewriter, loc, reduced, accReg, + Value reduced = createReducingOp(op, loadVal, loadAcc, builder); + InBoundsStoreOp::create(builder, loc, reduced, accReg, zeroConstantOp); - } + }; + + emitForActiveReductionThread(emitThreadwiseReduction); } if (canUsePermlaneSwap_NRSmall) { @@ -2304,21 +2342,23 @@ struct BlockwiseReduceRewritePattern SmallVector bounds{1, 1, 1}; SmallVector strides{1, 1, 1}; - TransformingForOp storeLoop = TransformingForOp::create( - rewriter, loc, ArrayRef(inits), - ArrayRef{threadToLDSViewTrs}, - ArrayRef(bounds), ArrayRef(strides), - /*forceUnroll=*/true, /*useIndexDiffs=*/true); - { - PatternRewriter::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(storeLoop.getBody()); + auto emitThreadwiseResultStore = [&](OpBuilder &builder) { + TransformingForOp storeLoop = TransformingForOp::create( + builder, loc, ArrayRef(inits), + ArrayRef{threadToLDSViewTrs}, + ArrayRef(bounds), ArrayRef(strides), + /*forceUnroll=*/true, /*useIndexDiffs=*/true); + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(storeLoop.getBody()); Block::BlockArgListType LDSStoreCoords = storeLoop.getLowerCoords(/*domain=*/0); - Value loadVal = InBoundsLoadOp::create(rewriter, loc, elemType, + Value loadVal = InBoundsLoadOp::create(builder, loc, elemType, accReg, zeroConstantOp); - InBoundsStoreOp::create(rewriter, loc, loadVal, + InBoundsStoreOp::create(builder, loc, loadVal, workspaceLDSBuffer, LDSStoreCoords); - } + }; + + emitForActiveReductionThread(emitThreadwiseResultStore); LDSBarrierOp::create(rewriter, loc); } diff --git a/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir b/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir index 6654cc7a6044..0709781d9cbe 100644 --- a/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir +++ b/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir @@ -1,4 +1,8 @@ // RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-opt -rock-blockwise-gemm-to-threadwise -canonicalize -split-input-file | FileCheck %s +// Keep exact-packing checks in a separate invocation because CHECK and EXACT +// both label the same functions; combining prefixes would require a second +// match for each function. +// RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-opt -rock-blockwise-gemm-to-threadwise -canonicalize -split-input-file | FileCheck %s --check-prefix=EXACT // CHECK-DAG: #[[MAP:.*]] = affine_map<(d0) -> (d0, 0)> // CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)> @@ -77,6 +81,9 @@ func.func @rock_blockwise_reducesum_nr_threads_gt_blocksize(%input_reg : memref< #inputView_tid = #rock.transform_map (0, d0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [8] -> [1, 8]> #inputView_iter = #rock.transform_map (d0, 0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [4] -> [4, 1]> // CHECK-LABEL: func @rock_blockwise_reducesum_rthreads_fix +// EXACT-LABEL: func @rock_blockwise_reducesum_rthreads_fix +// EXACT-NOT: arith.cmpi ult +// EXACT: return func.func @rock_blockwise_reducesum_rthreads_fix(%input_reg : memref<4xf32, #gpu.address_space>, %output_reg : memref<4xf32, #gpu.address_space>, %ws_lds : memref<32xf32, #gpu.address_space>) attributes{rock.arch = "##TOKEN_ARCH##", block_size = 8 : i32, grid_size = 2 : i32, rock.kernel} { // Compute rthread index and nr index from tid // blockSize=8, nrDimProd=4, rTid=2, cs=2 -> cs*nrDimProd=8==blockSize @@ -127,6 +134,9 @@ func.func @rock_blockwise_reducesum_rthreads_fix(%input_reg : memref<4xf32, #gpu // CHECK-DAG: #[[TMAP13:.*]] = #rock.transform_map<#[[MAP7]] // CHECK: func @rock_blockwise_reducesum_nr_threads_lt_blocksize +// EXACT-LABEL: func @rock_blockwise_reducesum_nr_threads_lt_blocksize +// EXACT-NOT: arith.cmpi ult +// EXACT: return // CHECK-DAG: %[[ZEROFP:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0 : index @@ -210,6 +220,9 @@ func.func @rock_blockwise_reducesum_nr_threads_lt_blocksize(%input_reg : memref< #inputView_iter = #rock.transform_map (d0, 0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [5] -> [5, 1]> // CHECK-LABEL: func @rock_blockwise_reducesum_nonpow2_nrdimprod +// EXACT-LABEL: func @rock_blockwise_reducesum_nonpow2_nrdimprod +// EXACT-NOT: arith.cmpi ult +// EXACT: return // CHECK-DAG: %[[TID:.*]] = rock.workitem_id : index // CHECK: %[[RTID:.*]] = arith.divsi %[[TID]], %c5 // CHECK: %[[NRTID:.*]] = arith.remsi %[[TID]], %c5 @@ -220,3 +233,64 @@ func.func @rock_blockwise_reducesum_nonpow2_nrdimprod(%input_reg : memref<5xf32, rock.blockwise_broadcast_reduce sum [#inputView][#inputView_tid][#inputView_iter]%input_reg into %output_reg using %ws_lds {axis = 1 : index, blockSize = 15 : i32, nrDimPerThread = 5 : index} : memref<5xf32, #gpu.address_space> using memref<75xf32, #gpu.address_space> into memref<5xf32, #gpu.address_space> return } + +// ----- + +// Verify floor division can leave workitems inactive. blockSize=24 and +// nrDimProd=5 produce four reduction threads per row and 20 active workitems. + +#inputView = #rock.transform_map (d1, d0)> by [ ["r"] at [1]>, ["nr_per_bid"] at [0]>] bounds = [24, 5] -> [5, 24]> +#inputView_tid = #rock.transform_map (0, d0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [24] -> [1, 24]> +#inputView_iter = #rock.transform_map (d0, 0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [5] -> [5, 1]> + +// CHECK-LABEL: func @rock_blockwise_reducesum_inactive_threads +// CHECK-DAG: %[[TID:.*]] = rock.workitem_id : index +// CHECK-DAG: %[[ACTIVE_COUNT:.*]] = arith.constant 20 : index +// CHECK: %[[ACTIVE:.*]] = arith.cmpi ult, %[[TID]], %[[ACTIVE_COUNT]] : index +// CHECK: scf.if %[[ACTIVE]] { +// CHECK: rock.in_bounds_load {{.*}} : memref<120xf32, #gpu.address_space>, index -> vector<2xf32> +// CHECK: rock.in_bounds_store {{.*}} : f32 -> memref<120xf32, #gpu.address_space>, index +// The two closing braces end the result-store loop and active-thread guard. +// The barrier must follow the guard so it remains workgroup-uniform. +// CHECK: rock.yield +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: rock.lds_barrier + +func.func @rock_blockwise_reducesum_inactive_threads(%input_reg : memref<5xf32, #gpu.address_space>, %output_reg : memref<5xf32, #gpu.address_space>, %ws_lds : memref<120xf32, #gpu.address_space>) attributes{rock.arch = "##TOKEN_ARCH##", block_size = 24 : i32, grid_size = 8 : i32, rock.kernel} { + rock.blockwise_broadcast_reduce sum [#inputView][#inputView_tid][#inputView_iter]%input_reg into %output_reg using %ws_lds {axis = 1 : index, blockSize = 24 : i32, nrDimPerThread = 5 : index} : memref<5xf32, #gpu.address_space> using memref<120xf32, #gpu.address_space> into memref<5xf32, #gpu.address_space> + return +} + +// ----- + +// Verify the divisibility clamp can independently leave workitems inactive. +// blockSize=20 and nrDimProd=3 give 6 reduction threads per row, but 6 does not +// divide the 20-element reduction dimension, so rthreads drops to 5 and only +// 15 workitems remain active. + +#inputView = #rock.transform_map (d1, d0)> by [ ["r"] at [1]>, ["nr_per_bid"] at [0]>] bounds = [20, 3] -> [3, 20]> +#inputView_tid = #rock.transform_map (0, d0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [20] -> [1, 20]> +#inputView_iter = #rock.transform_map (d0, 0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [3] -> [3, 1]> + +// Check that the view uses five reduction threads with four elements each and +// does not pad the reduction dimension. +// CHECK-DAG: ["rDim"] at [1]> +// CHECK-DAG: ["rDim"] at [1]> + +// CHECK-LABEL: func @rock_blockwise_reducesum_rthreads_clamped_inactive +// CHECK-DAG: %[[TID:.*]] = rock.workitem_id : index +// CHECK-DAG: %[[ACTIVE_COUNT:.*]] = arith.constant 15 : index +// CHECK: %[[ACTIVE:.*]] = arith.cmpi ult, %[[TID]], %[[ACTIVE_COUNT]] : index +// CHECK: scf.if %[[ACTIVE]] { +// CHECK: rock.in_bounds_load {{.*}} : memref<60xf32, #gpu.address_space>, index -> vector<4xf32> +// CHECK: rock.in_bounds_store {{.*}} : f32 -> memref<60xf32, #gpu.address_space>, index +// CHECK: rock.yield +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: rock.lds_barrier + +func.func @rock_blockwise_reducesum_rthreads_clamped_inactive(%input_reg : memref<3xf32, #gpu.address_space>, %output_reg : memref<3xf32, #gpu.address_space>, %ws_lds : memref<60xf32, #gpu.address_space>) attributes{rock.arch = "##TOKEN_ARCH##", block_size = 20 : i32, grid_size = 8 : i32, rock.kernel} { + rock.blockwise_broadcast_reduce sum [#inputView][#inputView_tid][#inputView_iter]%input_reg into %output_reg using %ws_lds {axis = 1 : index, blockSize = 20 : i32, nrDimPerThread = 3 : index} : memref<3xf32, #gpu.address_space> using memref<60xf32, #gpu.address_space> into memref<3xf32, #gpu.address_space> + return +} From 893f0f5b68f8a320481e914eee9fdaaf22c33888 Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Tue, 18 Aug 2026 19:38:51 +0000 Subject: [PATCH 5/5] add nrlargerthanthread test --- .../lowering_blockwise_broadcast_reduce.mlir | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir b/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir index 0709781d9cbe..28f8b7828f5e 100644 --- a/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir +++ b/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir @@ -294,3 +294,29 @@ func.func @rock_blockwise_reducesum_rthreads_clamped_inactive(%input_reg : memre rock.blockwise_broadcast_reduce sum [#inputView][#inputView_tid][#inputView_iter]%input_reg into %output_reg using %ws_lds {axis = 1 : index, blockSize = 20 : i32, nrDimPerThread = 3 : index} : memref<3xf32, #gpu.address_space> using memref<60xf32, #gpu.address_space> into memref<3xf32, #gpu.address_space> return } + +// ----- + +// Verify the NR-Large path pads the non-reduction dimension when blockSize +// does not evenly divide nrDimProd. blockSize=4 and nrDimProd=6 require two +// non-reduction iterations per thread and two padded positions. + +#inputView = #rock.transform_map (d1, d0)> by [ ["r"] at [1]>, ["nr_per_bid"] at [0]>] bounds = [4, 6] -> [6, 4]> +#inputView_tid = #rock.transform_map (0, d0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [4] -> [1, 4]> +#inputView_iter = #rock.transform_map (d0, 0)> by [ ["nr_per_bid", "r"] at [0, 1]>] bounds = [6] -> [6, 1]> + +// CHECK-DAG: ["nrDim"] at [0]> +// CHECK-DAG: ["nrDim"] at [0]> + +// CHECK-LABEL: func @rock_blockwise_reducesum_nrlarge_uneven +// CHECK: affine.for {{.*}} = 0 to 2 +// CHECK: rock.transforming_for +// CHECK-SAME: bounds [1, 1, 4] strides [1, 1, 4] +// CHECK: vector.reduction +// CHECK: rock.lds_barrier +// CHECK: rock.threadwise_read_into + +func.func @rock_blockwise_reducesum_nrlarge_uneven(%input_reg : memref<6xf32, #gpu.address_space>, %output_reg : memref<6xf32, #gpu.address_space>, %ws_lds : memref<24xf32, #gpu.address_space>) attributes{rock.arch = "##TOKEN_ARCH##", block_size = 4 : i32, grid_size = 8 : i32, rock.kernel} { + rock.blockwise_broadcast_reduce sum [#inputView][#inputView_tid][#inputView_iter]%input_reg into %output_reg using %ws_lds {axis = 1 : index, blockSize = 4 : i32, nrDimPerThread = 6 : index} : memref<6xf32, #gpu.address_space> using memref<24xf32, #gpu.address_space> into memref<6xf32, #gpu.address_space> + return +}