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/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/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir b/mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir index 6654cc7a6044..28f8b7828f5e 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,90 @@ 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 +} + +// ----- + +// 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 +} 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] 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..b4bc669c54a6 --- /dev/null +++ b/mlir/test/rocmlir-driver/verify-all-zero.mlir @@ -0,0 +1,37 @@ +// 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 + +// 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..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]], %[[sqkMaxs]] : (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 %[[expsSumsTensor]] : (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 5d253d8149c1..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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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 602ba6c910d3..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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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 c06bffc852c9..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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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 208c00f7627a..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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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 d040060e1571..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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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 f447919f2255..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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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]], %[[sqkMaxs]] : ([[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 %[[expsSumsTensor]] : ([[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-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..6a711cc44905 100644 --- a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp +++ b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp @@ -2979,24 +2979,32 @@ 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()) && + "expected a float element type"); + 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); + assert(status == APFloat::opOK && + "failed to create exact floating-point splat value"); - 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 +3367,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 +4574,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 +4610,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);