Only add one LDS Barrier for one wave blockSizes

mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp

@@ -21,6 +21,7 @@
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/Rock/IR/GetRockInfo.h"
 #include "mlir/Dialect/Rock/IR/Rock.h"
 #include "mlir/Dialect/Rock/Passes.h"
 #include "mlir/Dialect/Rock/Transforms/RockMultibuffer.h"
@@ -139,8 +140,9 @@ struct RemoveBackToBackBarriersRewritePattern
 
   LogicalResult matchAndRewrite(rock::LDSBarrierOp op,
                                 PatternRewriter &rw) const override {
-    if (dyn_cast_or_null<rock::LDSBarrierOp>(op->getNextNode())) {
-      op->getNextNode()->erase();
+    if (auto nextBarrier =
+            dyn_cast_or_null<rock::LDSBarrierOp>(op->getNextNode())) {
+      rw.eraseOp(nextBarrier);
       return success();
     }
     return failure();
@@ -162,7 +164,12 @@ struct PushBarrierDownRewritePattern
       return failure();
 
     // Don't go over the terminator
-    if (!nextOp->getNextNode())
+    if (nextOp->hasTrait<OpTrait::IsTerminator>() ||
+        nextOp->hasTrait<OpTrait::ReturnLike>())
+      return failure();
+
+    // Don't push past another barrier - let RemoveBackToBackBarriers handle it
+    if (isa<rock::LDSBarrierOp>(nextOp))
       return failure();
 
     // We assume that operations that have a body may modify LDS
@@ -465,6 +472,71 @@ DagType pruneGraph(const DagType &dag) {
   return prunedGraph;
 }
 
+// Determine if the backward barrier can be skipped for single-wave kernels.
+//
+// For scheduleVersion 1 (Default) or 3 (DirectToLDSDefault), the loop
+// structure is:
+//   GlobalLoad -> DSWrite -> (fwd barrier) -> DSRead + MFMA
+//
+// The forward barrier ensures DSWrites complete before DSReads start.
+// For the loop-carried dependency (backward barrier), we need to ensure
+// DSReads from iteration i finish before DSWrites from iteration i+1.
+//
+// When blockSize <= waveSize (single wave), this is guaranteed because
+// GPU issues instructions in order within a wave - once DSReads have been
+// issued, they have read the data from the buffers, so DSWrites can proceed
+// without an explicit barrier.
+bool canSkipBackwardBarrierForOneWave(func::FuncOp func, scf::ForOp forOp) {
+  // Check if this is a single-wave kernel
+  auto maybeBlockSize = rock::getBlockSize(func);
+  if (failed(maybeBlockSize))
+    return false;
+
+  int64_t blockSize = maybeBlockSize->getInt();
+
+  // Check if arch attribute exists before calling getArchValue which
+  // triggers llvm_unreachable if arch is missing
+  if (!func->hasAttr("arch") && !func->hasAttr("mhal.arch"))
+    return false;
+
+  StringAttr arch = rock::getArchValue(func);
+  if (!arch)
+    return false;
+
+  int64_t waveSize = rock::lookupArchInfo(arch).waveSize;
+  bool isOneWave = (blockSize <= waveSize);
+  if (!isOneWave)
+    return false;
+
+  // for nested loops, it may require more analysis. For now, only support
+  // single loop.
+  int forOpCount = 0;
+  func.walk([&](scf::ForOp) { ++forOpCount; });
+  if (forOpCount != 1)
+    return false;
+
+  // Find the scheduleVersion from ThreadwiseGemmAccelOp within the loop.
+  // The scheduleVersion is stored in the params attribute of the op.
+  std::optional<int64_t> scheduleVersion;
+  forOp.walk([&](rock::ThreadwiseGemmAccelOp gemmOp) {
+    rock::RockAccelTuningParamAttrInterface params = gemmOp.getParams();
+    scheduleVersion = params.getScheduleVersion();
+  });
+
+  if (!scheduleVersion.has_value())
+    return false;
+
+  // Check if the schedule version supports skipping the backward barrier.
+  // Only scheduleVersion 1 (Default) and 3 (DirectToLDSDefault)
+  // have the loop structure that allows skipping the backward barrier.
+  bool canSkip = (*scheduleVersion == 1 || *scheduleVersion == 3);
+
+  LLVM_DEBUG(DBGS() << "canSkipBackwardBarrierForOneWave: isOneWave="
+                    << isOneWave << ", scheduleVersion=" << *scheduleVersion
+                    << ", canSkip=" << canSkip << "\n");
+  return canSkip;
+}
+
 // Utility function to place an empty stage before or after another `stage`. The
 // empty stage will contain an `lds_barrier` if `isBarrier` is set to true
 rock::StageOp placeEmptyStage(IRRewriter &rewriter, Location loc,
@@ -487,8 +559,8 @@ rock::StageOp placeEmptyStage(IRRewriter &rewriter, Location loc,
 // initiation interval twice as big and pipeline as usual. This function
 // takes also care to update the initiation interval, so that the caller
 // does not have to know how `placeBarrier` internally works.
-void placeBarriers(IRRewriter &rewriter, Location loc, scf::ForOp forOp,
-                   ArrayRef<rock::StageOp> stages,
+void placeBarriers(IRRewriter &rewriter, Location loc, func::FuncOp func,
+                   scf::ForOp forOp, ArrayRef<rock::StageOp> stages,
                    SetVector<rock::GpuAllocOp> &allocs,
                    SmallVector<rock::StageOp> &extendedStages,
                    int64_t &initiationInterval, int64_t numIterations) {
@@ -497,8 +569,9 @@ void placeBarriers(IRRewriter &rewriter, Location loc, scf::ForOp forOp,
   dag = pruneGraph(dag);
 
   // If there is a loop, we probably need a backward barrier, i.e.,
-  // an LDS barrier that takes the loop dependency into account
-  const bool addBackwardBarrier = numIterations > 1;
+  // an LDS barrier that takes the loop dependency into account.
+  bool canSkipBackwardBarrier = canSkipBackwardBarrierForOneWave(func, forOp);
+  const bool addBackwardBarrier = numIterations > 1 && !canSkipBackwardBarrier;
 
   DenseMap<rock::StageOp, int> timeSlotMap;
   int timeSlot = 0;
@@ -733,14 +806,14 @@ void RockPipeline::runOnOperation() {
 
       forOp.walk([&](rock::StageOp stageOp) { stages.push_back(stageOp); });
 
+      if (stages.empty())
+        continue;
+
       forOp.walk([](rock::LDSBarrierOp barrier) {
         if (!barrier->getParentOfType<rock::StageOp>())
           barrier->erase();
       });
 
-      if (stages.empty())
-        continue;
-
       LLVM_DEBUG(DBGS() << "Number of stages: " << stages.size() << "\n");
       LLVM_DEBUG(DBGS() << "Initiation Interval: " << ii << "\n");
       size_t numStages = stages.size();
@@ -762,9 +835,8 @@ void RockPipeline::runOnOperation() {
       SmallVector<rock::StageOp> extendedStages;
       // use "multiAllocs" to place LDS barriers, no need to explicitly place
       // barriers for registers or globals
-      placeBarriers(rewriter, loc, forOp, stages, multiAllocs, extendedStages,
-                    ii, numIterations);
-
+      placeBarriers(rewriter, loc, func, forOp, stages, multiAllocs,
+                    extendedStages, ii, numIterations);
       ScheduleType schedule;
       // use all "resources" to generate dependency graph and generate schedule
       createSchedule(extendedStages, resources, ii, schedule,
@@ -800,26 +872,13 @@ void RockPipeline::runOnOperation() {
     // Cleanup the stages
     {
       if (removeStages) {
-        RewritePatternSet patternsPushBarrier(&getContext());
-        // run PushBarrierDownRewritePattern before RemoveStagesRewritePattern,
-        // because the latter will remove the stages and their terminators
-        patternsPushBarrier.add<PushBarrierDownRewritePattern>(ctx);
-        if (failed(applyPatternsGreedily(func, std::move(patternsPushBarrier))))
-          return signalPassFailure();
-
-        // run RemoveStagesRewritePattern before
-        // RemoveBackToBackBarriersRewritePattern, because the latter expects to
-        // find no stages
-        RewritePatternSet patternsRemoveStages(&getContext());
-        patternsRemoveStages.add<RemoveStagesRewritePattern>(ctx);
+        RewritePatternSet patterns(&getContext());
+        patterns.add<RemoveStagesRewritePattern, PushBarrierDownRewritePattern,
+                     RemoveBackToBackBarriersRewritePattern>(&getContext());
         if (failed(
-                applyPatternsGreedily(func, std::move(patternsRemoveStages))))
-          return signalPassFailure();
-
-        RewritePatternSet patternsBackToBack(&getContext());
-        patternsBackToBack.add<RemoveBackToBackBarriersRewritePattern>(ctx);
-        if (failed(applyPatternsGreedily(func, std::move(patternsBackToBack))))
+                applyPatternsGreedily(getOperation(), std::move(patterns)))) {
           return signalPassFailure();
+        }
       }
     }
   }

mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir

@@ -4,7 +4,7 @@
 // COUNT-COUNT-1: rock.lds_barrier
 
 module {
-  func.func @pipeline_loop_in_scf_if(%arg0: memref<128xf16>, %arg1: memref<128xf16>, %arg2: memref<128xf16>, %arg3: i32) attributes {block_size = 64 : i32, grid_size = 1 : i32, kernel} {
+  func.func @pipeline_loop_in_scf_if(%arg0: memref<128xf16>, %arg1: memref<128xf16>, %arg2: memref<128xf16>, %arg3: i32) attributes {arch = "amdgcn-amd-amdhsa:gfx90a", block_size = 64 : i32, grid_size = 1 : i32, kernel} {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
     %c4 = arith.constant 4 : index
@@ -66,6 +66,10 @@ module {
           // CHECK: arith.addf
           // CHECK: memref.store {{.*}}[%[[INNER_IV]]]
           // CHECK: }
+          // CHECK: %[[ALLOC_LDS_A:.*]] = rock.alloc() : memref<16xf16, #gpu.address_space<private>>
+          // CHECK: %[[ALLOC_LDS_B:.*]] = rock.alloc() : memref<16xf16, #gpu.address_space<private>>
+          // CHECK: %[[WID_LDS:.*]] = rock.workitem_id : index
+          // CHECK: memref.load %[[ALLOC_LDS_A]][%c0]
           // CHECK-NEXT: rock.lds_barrier
           affine.for %arg5 = 0 to 16 {
             %4 = memref.load %1[%arg5] : memref<64xf16, #gpu.address_space<workgroup>>
@@ -82,13 +86,7 @@ module {
         rock.lds_barrier
       } {pipeline = #rock.pipeline<2>}
 
-      // CHECK: %[[ALLOC_G:.*]] = rock.alloc() : memref<16xf16, #gpu.address_space<private>>
-      // CHECK: %[[ALLOC_H:.*]] = rock.alloc() : memref<16xf16, #gpu.address_space<private>>
-      // CHECK: %[[WID4:.*]] = rock.workitem_id : index
-      // CHECK: memref.load %[[ALLOC_G]][%c0]
-      // CHECK: memref.store {{.*}}, {{.*}}[%[[WID4]]]
-      // CHECK: memref.load %[[ALLOC_H]][%c0]
-      // CHECK: memref.store {{.*}}, {{.*}}[%[[WID4]]]
+      // CHECK: memref.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<64xf16, #gpu.address_space<workgroup>>
       // CHECK: }
       // CHECK-NOT: {pipeline = #rock.pipeline<2>}