[Attention] Handle fully masked reference rows

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<int64_t> 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<ValueRange>(inits),
-                ArrayRef<Attribute>{threadToLDSViewTrs},
-                ArrayRef<int64_t>(bounds), ArrayRef<int64_t>(strides),
-                /*forceUnroll=*/true, /*useIndexDiffs=*/true);
-            {
-              PatternRewriter::InsertionGuard guard(rewriter);
-              rewriter.setInsertionPointToStart(reductionLoop.getBody());
+            auto emitThreadwiseReduction = [&](OpBuilder &builder) {
+              TransformingForOp reductionLoop = TransformingForOp::create(
+                  builder, loc, ArrayRef<ValueRange>(inits),
+                  ArrayRef<Attribute>{threadToLDSViewTrs},
+                  ArrayRef<int64_t>(bounds), ArrayRef<int64_t>(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<int64_t> bounds{1, 1, 1};
               SmallVector<int64_t> strides{1, 1, 1};
 
-              TransformingForOp storeLoop = TransformingForOp::create(
-                  rewriter, loc, ArrayRef<ValueRange>(inits),
-                  ArrayRef<Attribute>{threadToLDSViewTrs},
-                  ArrayRef<int64_t>(bounds), ArrayRef<int64_t>(strides),
-                  /*forceUnroll=*/true, /*useIndexDiffs=*/true);
-              {
-                PatternRewriter::InsertionGuard guard(rewriter);
-                rewriter.setInsertionPointToStart(storeLoop.getBody());
+              auto emitThreadwiseResultStore = [&](OpBuilder &builder) {
+                TransformingForOp storeLoop = TransformingForOp::create(
+                    builder, loc, ArrayRef<ValueRange>(inits),
+                    ArrayRef<Attribute>{threadToLDSViewTrs},
+                    ArrayRef<int64_t>(bounds), ArrayRef<int64_t>(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);
             }
 

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<double>(dataSize);
   double aveRelDiff = sumRelDiff / static_cast<double>(dataSize);
-  double err_RMS = sqrt(sumDiffSq) / (static_cast<double>(maxMag) *
-                                      sqrt(static_cast<double>(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<double>(maxMag) *
+                                           sqrt(static_cast<double>(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;

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<affine_map<(d0) -> (0, d0)> by [<Merge{1, 8} ["tid"] at [0] -> ["nr_per_bid", "r"] at [0, 1]>] bounds = [8] -> [1, 8]>
 #inputView_iter = #rock.transform_map<affine_map<(d0) -> (d0, 0)> by [<Merge{4, 1} ["iter"] at [0] -> ["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<private>>, %output_reg : memref<4xf32, #gpu.address_space<private>>, %ws_lds : memref<32xf32, #gpu.address_space<workgroup>>) 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<affine_map<(d0) -> (d0, 0)> by [<Merge{5, 1} ["iter"] at [0] -> ["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<private>> using memref<75xf32, #gpu.address_space<workgroup>> into memref<5xf32, #gpu.address_space<private>>
   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<affine_map<(d0, d1) -> (d1, d0)> by [<PassThrough ["tid"] at [0] -> ["r"] at [1]>, <PassThrough ["iter"] at [1] -> ["nr_per_bid"] at [0]>] bounds = [24, 5] -> [5, 24]>
+#inputView_tid = #rock.transform_map<affine_map<(d0) -> (0, d0)> by [<Merge{1, 24} ["tid"] at [0] -> ["nr_per_bid", "r"] at [0, 1]>] bounds = [24] -> [1, 24]>
+#inputView_iter = #rock.transform_map<affine_map<(d0) -> (d0, 0)> by [<Merge{5, 1} ["iter"] at [0] -> ["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<workgroup>>, index -> vector<2xf32>
+// CHECK:   rock.in_bounds_store {{.*}} : f32 -> memref<120xf32, #gpu.address_space<workgroup>>, 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<private>>, %output_reg : memref<5xf32, #gpu.address_space<private>>, %ws_lds : memref<120xf32, #gpu.address_space<workgroup>>) 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<private>> using memref<120xf32, #gpu.address_space<workgroup>> into memref<5xf32, #gpu.address_space<private>>
+  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<affine_map<(d0, d1) -> (d1, d0)> by [<PassThrough ["tid"] at [0] -> ["r"] at [1]>, <PassThrough ["iter"] at [1] -> ["nr_per_bid"] at [0]>] bounds = [20, 3] -> [3, 20]>
+#inputView_tid = #rock.transform_map<affine_map<(d0) -> (0, d0)> by [<Merge{1, 20} ["tid"] at [0] -> ["nr_per_bid", "r"] at [0, 1]>] bounds = [20] -> [1, 20]>
+#inputView_iter = #rock.transform_map<affine_map<(d0) -> (d0, 0)> by [<Merge{3, 1} ["iter"] at [0] -> ["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: <Unmerge{5, 4} ["rtid", "rIter"] at [1, 2] -> ["rDim"] at [1]>
+// CHECK-DAG: <Pad{0, 0} ["rDim"] at [1] -> ["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<workgroup>>, index -> vector<4xf32>
+// CHECK:   rock.in_bounds_store {{.*}} : f32 -> memref<60xf32, #gpu.address_space<workgroup>>, 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<private>>, %output_reg : memref<3xf32, #gpu.address_space<private>>, %ws_lds : memref<60xf32, #gpu.address_space<workgroup>>) 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<private>> using memref<60xf32, #gpu.address_space<workgroup>> into memref<3xf32, #gpu.address_space<private>>
+  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<affine_map<(d0, d1) -> (d1, d0)> by [<PassThrough ["tid"] at [0] -> ["r"] at [1]>, <PassThrough ["iter"] at [1] -> ["nr_per_bid"] at [0]>] bounds = [4, 6] -> [6, 4]>
+#inputView_tid = #rock.transform_map<affine_map<(d0) -> (0, d0)> by [<Merge{1, 4} ["tid"] at [0] -> ["nr_per_bid", "r"] at [0, 1]>] bounds = [4] -> [1, 4]>
+#inputView_iter = #rock.transform_map<affine_map<(d0) -> (d0, 0)> by [<Merge{6, 1} ["iter"] at [0] -> ["nr_per_bid", "r"] at [0, 1]>] bounds = [6] -> [6, 1]>
+
+// CHECK-DAG: <Unmerge{4, 2} ["tid", "nrIter"] at [0, 1] -> ["nrDim"] at [0]>
+// CHECK-DAG: <Pad{0, 2} ["nrDim"] at [0] -> ["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 <add>
+// CHECK: rock.lds_barrier
+// CHECK: rock.threadwise_read_into
+
+func.func @rock_blockwise_reducesum_nrlarge_uneven(%input_reg : memref<6xf32, #gpu.address_space<private>>, %output_reg : memref<6xf32, #gpu.address_space<private>>, %ws_lds : memref<24xf32, #gpu.address_space<workgroup>>) 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<private>> using memref<24xf32, #gpu.address_space<workgroup>> into memref<6xf32, #gpu.address_space<private>>
+  return
+}

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]