From fd5699d30800e55558c6a67d5ba0490f2d515e9b Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Tue, 24 Feb 2026 15:20:49 +0000 Subject: [PATCH 1/6] Fix barrier placement for scheduleVersion = 1. It should appear before LDSRead and not before GlobalLoad. Also fix a bug where, if there are no stages found then it should preserve barriers in original loop --- .../Dialect/Rock/Transforms/RockPipeline.cpp | 48 ++--- .../Rock/rock-pipeline-early-exit.mlir | 12 +- .../test/Dialect/Rock/test_rock_pipeline.mlir | 187 ++++++++++++++++ .../Rock/test_rock_pipeline_nested.mlir | 202 ++++++++++++++++-- 4 files changed, 406 insertions(+), 43 deletions(-) diff --git a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp index 17e5e2534056..84a40b231c6a 100644 --- a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp +++ b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp @@ -139,8 +139,9 @@ struct RemoveBackToBackBarriersRewritePattern LogicalResult matchAndRewrite(rock::LDSBarrierOp op, PatternRewriter &rw) const override { - if (dyn_cast_or_null(op->getNextNode())) { - op->getNextNode()->erase(); + if (auto nextBarrier = + dyn_cast_or_null(op->getNextNode())) { + rw.eraseOp(nextBarrier); return success(); } return failure(); @@ -162,7 +163,12 @@ struct PushBarrierDownRewritePattern return failure(); // Don't go over the terminator - if (!nextOp->getNextNode()) + if (nextOp->hasTrait() || + nextOp->hasTrait()) + return failure(); + + // Don't push past another barrier - let RemoveBackToBackBarriers handle it + if (isa(nextOp)) return failure(); // We assume that operations that have a body may modify LDS @@ -733,14 +739,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()) 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(); @@ -764,6 +770,7 @@ void RockPipeline::runOnOperation() { // barriers for registers or globals placeBarriers(rewriter, loc, forOp, stages, multiAllocs, extendedStages, ii, numIterations); + LLVM_DEBUG(DBGS() << "ForOp: " << forOp << "\n"); ScheduleType schedule; // use all "resources" to generate dependency graph and generate schedule @@ -796,30 +803,23 @@ void RockPipeline::runOnOperation() { } } } + LLVM_DEBUG({ + DBGS() << "After remulti-buffer\n"; + DBGS() << "===============\n"; + DBGS() << "Operation: " << getOperation() << "\n"; + DBGS() << "===============\n"; + }); // Cleanup the stages { if (removeStages) { - RewritePatternSet patternsPushBarrier(&getContext()); - // run PushBarrierDownRewritePattern before RemoveStagesRewritePattern, - // because the latter will remove the stages and their terminators - patternsPushBarrier.add(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(ctx); + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext()); if (failed( - applyPatternsGreedily(func, std::move(patternsRemoveStages)))) - return signalPassFailure(); - - RewritePatternSet patternsBackToBack(&getContext()); - patternsBackToBack.add(ctx); - if (failed(applyPatternsGreedily(func, std::move(patternsBackToBack)))) + applyPatternsGreedily(getOperation(), std::move(patterns)))) { return signalPassFailure(); + } } } } diff --git a/mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir b/mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir index 592557e3e24e..93d7bfce7506 100644 --- a/mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir +++ b/mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir @@ -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> + // CHECK: %[[ALLOC_LDS_B:.*]] = rock.alloc() : memref<16xf16, #gpu.address_space> + // 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> @@ -82,13 +86,7 @@ module { rock.lds_barrier } {pipeline = #rock.pipeline<2>} - // CHECK: %[[ALLOC_G:.*]] = rock.alloc() : memref<16xf16, #gpu.address_space> - // CHECK: %[[ALLOC_H:.*]] = rock.alloc() : memref<16xf16, #gpu.address_space> - // 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> // CHECK: } // CHECK-NOT: {pipeline = #rock.pipeline<2>} diff --git a/mlir/test/Dialect/Rock/test_rock_pipeline.mlir b/mlir/test/Dialect/Rock/test_rock_pipeline.mlir index 582bdc24dcf2..9a5de2df69a5 100644 --- a/mlir/test/Dialect/Rock/test_rock_pipeline.mlir +++ b/mlir/test/Dialect/Rock/test_rock_pipeline.mlir @@ -2,6 +2,17 @@ // RUN: rocmlir-opt %s --rock-pipeline="rock-pipeline-remove-stages=true" | FileCheck %s --check-prefix=REMOVE-STAGES // CHECK-LABEL: rock_pipeline_3_stages_ii_1 +// REMOVE-STAGES-LABEL: rock_pipeline_3_stages_ii_1 +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_BUF:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_BUF]]{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_3_stages_ii_1(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -60,6 +71,22 @@ func.func @rock_pipeline_3_stages_ii_1(%input : memref<16xi8, #gpu.address_space } // CHECK-LABEL: rock_pipeline_3_stages_ii_2 +// REMOVE-STAGES-LABEL: rock_pipeline_3_stages_ii_2 +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS read + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_RD:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.load %[[LDS_RD]] + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_3_stages_ii_2(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -119,6 +146,19 @@ func.func @rock_pipeline_3_stages_ii_2(%input : memref<16xi8, #gpu.address_space // this test shouldn't pipeline loop but it would add barriers and multibuffer by 1 // CHECK-LABEL: rock_pipeline_3_stages_ii_2_less_iterations +// REMOVE-STAGES-LABEL: rock_pipeline_3_stages_ii_2_less_iterations +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS read + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_RD:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.load %[[LDS_RD]] +// REMOVE-STAGES: } +// Epilogue: barrier before LDS read +// REMOVE-STAGES-NEXT: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_3_stages_ii_2_less_iterations(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -170,6 +210,23 @@ func.func @rock_pipeline_3_stages_ii_2_less_iterations(%input : memref<16xi8, #g } // CHECK-LABEL: rock_pipeline_3_stages_ii_3 +// REMOVE-STAGES-LABEL: rock_pipeline_3_stages_ii_3 +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xi8, #gpu.address_space> + // Barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_RD:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.load %[[LDS_RD]] +// REMOVE-STAGES: } +// No barrier after loop for this function - all barriers are inside the loop +// REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_3_stages_ii_3(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -223,6 +280,14 @@ func.func @rock_pipeline_3_stages_ii_3(%input : memref<16xi8, #gpu.address_space // This test shouldn't do any pipelining as it doesn't have any stages but it should still multibuffer by 1 // CHECK-LABEL: rock_pipeline_no_stages_ii_1 +// REMOVE-STAGES-LABEL: rock_pipeline_no_stages_ii_1 +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_no_stages_ii_1(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -266,6 +331,28 @@ func.func @rock_pipeline_no_stages_ii_1(%input : memref<16xi8, #gpu.address_spac } // CHECK-LABEL: rock_pipeline_4_stages_ii_2 +// REMOVE-STAGES-LABEL: rock_pipeline_4_stages_ii_2 +// REMOVE-STAGES-NOT: rock.stage +// Prologue: barrier before LDS read +// REMOVE-STAGES: memref.store %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES-NEXT: rock.lds_barrier +// REMOVE-STAGES-NEXT: %[[LDS_PRO:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space +// REMOVE-STAGES-NEXT: memref.load %[[LDS_PRO]] +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xi8, #gpu.address_space> + // Barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_RD:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.load %[[LDS_RD]] +// REMOVE-STAGES: } +// REMOVE-STAGES-NEXT: rock.lds_barrier +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_4_stages_ii_2(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -333,6 +420,27 @@ func.func @rock_pipeline_4_stages_ii_2(%input : memref<16xi8, #gpu.address_space } // CHECK-LABEL: rock_pipeline_4_stages_ii_1_i8 +// REMOVE-STAGES-LABEL: rock_pipeline_4_stages_ii_1_i8 +// REMOVE-STAGES-NOT: rock.stage +// Prologue: barrier before LDS write +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES-NEXT: %[[LDS_PRO:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space +// REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_PRO]]{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: } +// Epilogue: barrier before LDS write, barrier before LDS read +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.store %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_4_stages_ii_1_i8(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -413,6 +521,27 @@ func.func @rock_pipeline_4_stages_ii_1_i8(%input : memref<16xi8, #gpu.address_sp } // CHECK-LABEL: rock_pipeline_4_stages_ii_1_f16 +// REMOVE-STAGES-LABEL: rock_pipeline_4_stages_ii_1_f16 +// REMOVE-STAGES-NOT: rock.stage +// Prologue: barrier before LDS write +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES-NEXT: %[[LDS_PRO:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space +// REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_PRO]]{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: } +// Epilogue: barrier before LDS write, barrier before LDS read +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_4_stages_ii_1_f16(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -489,6 +618,23 @@ func.func @rock_pipeline_4_stages_ii_1_f16(%input : memref<16xf16, #gpu.address_ // This test should adjust II to 2 to enable loop pipelining // CHECK-LABEL: rock_pipeline_4_stages_ii_1_f16_less_iterations +// REMOVE-STAGES-LABEL: rock_pipeline_4_stages_ii_1_f16_less_iterations +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS read + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_RD:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.load %[[LDS_RD]] + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: } +// REMOVE-STAGES-NEXT: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_4_stages_ii_1_f16_less_iterations(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -557,6 +703,27 @@ func.func @rock_pipeline_4_stages_ii_1_f16_less_iterations(%input : memref<16xf1 // this test should do loop pipelining without adjust II but notice that it emits scf.for loop with zero iterations. // CHECK-LABEL: rock_pipeline_4_stages_ii_1_f16_less_iterations_2 +// REMOVE-STAGES-LABEL: rock_pipeline_4_stages_ii_1_f16_less_iterations_2 +// REMOVE-STAGES-NOT: rock.stage +// Prologue: barrier before LDS write +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES-NEXT: %[[LDS_PRO:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space +// REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_PRO]]{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: } +// Epilogue: barrier before LDS write, barrier before LDS read +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_4_stages_ii_1_f16_less_iterations_2(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -699,6 +866,26 @@ func.func @rock_nopipeline(%input : memref<16xi8, #gpu.address_space>, % // The three-way rotation S2,S3,S4 -> S4,S3,S2 avoids private multi-buffering // for regB and regC. // REMOVE-STAGES-LABEL: rock_pipeline_5_stages_three_way_swap +// REMOVE-STAGES-NOT: rock.stage +// Prologue: barrier before LDS write (twice for 2-deep prologue) +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // Barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES-NEXT: %[[LDS_WR:.*]] = rock.extract_multibuffer{{.*}}#gpu.address_space + // REMOVE-STAGES-NEXT: memref.store %{{.*}}, %[[LDS_WR]]{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: } +// Epilogue: barrier before LDS write, barrier before LDS read +// REMOVE-STAGES: memref.load %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.store %{{.*}} : memref<16xi8, #gpu.address_space> +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return // CHECK-LABEL: rock_pipeline_5_stages_three_way_swap func.func @rock_pipeline_5_stages_three_way_swap(%input : memref<16xi8, #gpu.address_space>, %output : memref<16xi8, #gpu.address_space>){ %c0 = arith.constant 0 : index diff --git a/mlir/test/Dialect/Rock/test_rock_pipeline_nested.mlir b/mlir/test/Dialect/Rock/test_rock_pipeline_nested.mlir index 7fdfd334f200..f3a463094aac 100644 --- a/mlir/test/Dialect/Rock/test_rock_pipeline_nested.mlir +++ b/mlir/test/Dialect/Rock/test_rock_pipeline_nested.mlir @@ -2,6 +2,23 @@ // RUN: rocmlir-opt %s --rock-pipeline="rock-pipeline-remove-stages=true" | FileCheck %s --check-prefix=REMOVE-STAGES // REMOVE-STAGES-LABEL: rock_nopipeline +// No pipeline attribute - no barriers expected +// REMOVE-STAGES: rock.alloc() : memref<32xi8, #gpu.address_space> +// REMOVE-STAGES: rock.alloc() : memref<32xi8, #gpu.address_space> +// REMOVE-STAGES-NOT: rock.alloc() : memref<32xi8, #gpu.address_space> +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // REMOVE-STAGES-NOT: rock.lds_barrier + // REMOVE-STAGES-NOT: rock.extract_multibuffer + // REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // REMOVE-STAGES-NOT: rock.lds_barrier + // REMOVE-STAGES: scf.for + // REMOVE-STAGES-NOT: rock.stage + // REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_nopipeline(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -18,18 +35,6 @@ func.func @rock_nopipeline(%input : memref<16xf16, #gpu.address_space>, %lds0 = memref.view %rawLds0[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> %lds1 = memref.view %rawLds1[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> - // REMOVE-STAGES: rock.alloc() : memref<32xi8, #gpu.address_space> - // REMOVE-STAGES: rock.alloc() : memref<32xi8, #gpu.address_space> - // REMOVE-STAGES-NOT: rock.alloc() : memref<32xi8, #gpu.address_space> - - // REMOVE-STAGES-NOT: rock.stage - // REMOVE-STAGES: scf.for - // REMOVE-STAGES-NOT: rock.stage - // REMOVE-STAGES-NOT: rock.extract_multibuffer - // REMOVE-STAGES: scf.for - // REMOVE-STAGES-NOT: rock.stage - // REMOVE-STAGES: scf.for - // REMOVE-STAGES-NOT: rock.stage scf.for %idx = %c0 to %c4 step %c1 { scf.for %arg3 = %c0 to %c3 step %c1 { rock.stage { @@ -88,6 +93,25 @@ func.func @rock_nopipeline(%input : memref<16xf16, #gpu.address_space>, // one loop inside another loop, inner loop has pipeline<1> attribute // CHECK-LABEL: rock_pipeline_oneloop +// REMOVE-STAGES-LABEL: rock_pipeline_oneloop +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // Prologue: barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS write + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES-NEXT: rock.lds_barrier + // REMOVE-STAGES: } + // Epilogue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // Epilogue: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: return func.func @rock_pipeline_oneloop(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>){ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -175,6 +199,40 @@ func.func @rock_pipeline_oneloop(%input : memref<16xf16, #gpu.address_space attribute // CHECK-LABEL: rock_pipeline_twoloops +// REMOVE-STAGES-LABEL: rock_pipeline_twoloops +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // First inner pipelined loop + // Prologue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: } + // Epilogue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // Epilogue: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // Second inner pipelined loop + // Prologue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: } + // Epilogue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // Epilogue: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: return func.func @rock_pipeline_twoloops(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -323,6 +381,40 @@ func.func @rock_pipeline_twoloops(%input : memref<16xf16, #gpu.address_space attribute // CHECK-LABEL: rock_pipeline_twoloops_ii2 +// REMOVE-STAGES-LABEL: rock_pipeline_twoloops_ii2 +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // First inner pipelined loop (ii=2) + // Prologue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // Loop body: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: } + // Epilogue: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // Second inner pipelined loop (ii=2) + // Prologue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> + // Loop body: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // REMOVE-STAGES: } + // Epilogue: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: return func.func @rock_pipeline_twoloops_ii2(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -444,6 +536,31 @@ func.func @rock_pipeline_twoloops_ii2(%input : memref<16xf16, #gpu.address_space // two loops inside an outer loop (which is inside another outer loop), inner loops have pipeline<1> attribute // CHECK-LABEL: rock_pipeline_twoloops_triplenested +// REMOVE-STAGES-LABEL: rock_pipeline_twoloops_triplenested +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // REMOVE-STAGES: scf.for + // First inner pipelined loop + // Prologue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: } + // Epilogue barriers + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier + // Second inner pipelined loop + // Prologue: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: } + // Epilogue barriers + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_twoloops_triplenested(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -594,6 +711,41 @@ func.func @rock_pipeline_twoloops_triplenested(%input : memref<16xf16, #gpu.addr // two outer loops that each contain two inner loops, inner loops have pipeline<1> attribute // CHECK-LABEL: rock_pipeline_twoloops_twoouterloops +// REMOVE-STAGES-LABEL: rock_pipeline_twoloops_twoouterloops +// REMOVE-STAGES-NOT: rock.stage +// First outer loop +// REMOVE-STAGES: scf.for + // First inner pipelined loop + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: scf.for + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: } + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier + // Second inner pipelined loop + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: scf.for + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: } + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier +// Second outer loop +// REMOVE-STAGES: scf.for + // First inner pipelined loop + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: scf.for + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: } + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier + // Second inner pipelined loop + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: scf.for + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: } + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_twoloops_twoouterloops(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -867,6 +1019,18 @@ func.func @rock_pipeline_twoloops_twoouterloops(%input : memref<16xf16, #gpu.add // one loop inside a loop, inner loops have pipeline<1> attribute and no rock.stages // CHECK-LABEL: rock_pipeline_nestednostages +// REMOVE-STAGES-LABEL: rock_pipeline_nestednostages +// No stages to pipeline - but barriers are preserved when stages are empty +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // REMOVE-STAGES: scf.for + // Barriers preserved: barrier before LDS write + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.store %{{.*}} : memref<16xf16, #gpu.address_space> + // Barriers preserved: barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: memref.load %{{.*}} : memref<16xf16, #gpu.address_space> +// REMOVE-STAGES: return func.func @rock_pipeline_nestednostages(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -916,6 +1080,20 @@ func.func @rock_pipeline_nestednostages(%input : memref<16xf16, #gpu.address_spa // two loops inside an outer loop, inner loops have pipeline<2> attribute and two rock.stages // CHECK-LABEL: rock_pipeline_twoloops_ii_equal_numstages +// REMOVE-STAGES-LABEL: rock_pipeline_twoloops_ii_equal_numstages +// REMOVE-STAGES-NOT: rock.stage +// REMOVE-STAGES: scf.for + // First inner pipelined loop (ii=2, 2 stages - no prologue/epilogue) + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS write, barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier + // Second inner pipelined loop (ii=2, 2 stages - no prologue/epilogue) + // REMOVE-STAGES: scf.for + // Loop body: barrier before LDS write, barrier before LDS read + // REMOVE-STAGES: rock.lds_barrier + // REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return func.func @rock_pipeline_twoloops_ii_equal_numstages(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index From 7faecdcb8fe48b79556b17a074e949ff6d887a0e Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Tue, 24 Feb 2026 15:21:22 +0000 Subject: [PATCH 2/6] Remove debug prints --- mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp | 8 -------- 1 file changed, 8 deletions(-) diff --git a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp index 84a40b231c6a..bc3abe13ef26 100644 --- a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp +++ b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp @@ -770,8 +770,6 @@ void RockPipeline::runOnOperation() { // barriers for registers or globals placeBarriers(rewriter, loc, forOp, stages, multiAllocs, extendedStages, ii, numIterations); - LLVM_DEBUG(DBGS() << "ForOp: " << forOp << "\n"); - ScheduleType schedule; // use all "resources" to generate dependency graph and generate schedule createSchedule(extendedStages, resources, ii, schedule, @@ -803,12 +801,6 @@ void RockPipeline::runOnOperation() { } } } - LLVM_DEBUG({ - DBGS() << "After remulti-buffer\n"; - DBGS() << "===============\n"; - DBGS() << "Operation: " << getOperation() << "\n"; - DBGS() << "===============\n"; - }); // Cleanup the stages { From 573295405025232130b11b9e0ec842795c903931 Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Tue, 24 Feb 2026 20:15:53 +0000 Subject: [PATCH 3/6] Add single barrier for single wave kernels --- .../Dialect/Rock/Transforms/RockPipeline.cpp | 79 ++++- .../Rock/rock-pipeline-early-exit.mlir | 2 +- .../Rock/test_rock_pipeline_wave_barrier.mlir | 271 ++++++++++++++++++ mlir/test/e2e/CMakeLists.txt | 5 + mlir/test/e2e/GemmOneWaveBarrier.cfg | 3 + mlir/test/e2e/GemmOneWaveBarrier.toml | 33 +++ .../e2e/GemmOneWaveBarrierDirectToLDS.cfg | 5 + .../e2e/GemmOneWaveBarrierDirectToLDS.toml | 33 +++ mlir/test/e2e/GemmOneWaveBarrierFp8.cfg | 5 + mlir/test/e2e/GemmOneWaveBarrierFp8.toml | 33 +++ mlir/test/e2e/PrGemmOneWaveBarrier.cfg | 3 + mlir/test/e2e/PrGemmOneWaveBarrier.toml | 22 ++ .../e2e/PrGemmOneWaveBarrierDirectToLDS.cfg | 3 + .../e2e/PrGemmOneWaveBarrierDirectToLDS.toml | 22 ++ 14 files changed, 512 insertions(+), 7 deletions(-) create mode 100644 mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir create mode 100644 mlir/test/e2e/GemmOneWaveBarrier.cfg create mode 100644 mlir/test/e2e/GemmOneWaveBarrier.toml create mode 100644 mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.cfg create mode 100644 mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.toml create mode 100644 mlir/test/e2e/GemmOneWaveBarrierFp8.cfg create mode 100644 mlir/test/e2e/GemmOneWaveBarrierFp8.toml create mode 100644 mlir/test/e2e/PrGemmOneWaveBarrier.cfg create mode 100644 mlir/test/e2e/PrGemmOneWaveBarrier.toml create mode 100644 mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.cfg create mode 100644 mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.toml diff --git a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp index bc3abe13ef26..bfc4f4e6c242 100644 --- a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp +++ b/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" @@ -471,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 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, @@ -493,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 stages, +void placeBarriers(IRRewriter &rewriter, Location loc, func::FuncOp func, + scf::ForOp forOp, ArrayRef stages, SetVector &allocs, SmallVector &extendedStages, int64_t &initiationInterval, int64_t numIterations) { @@ -503,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 timeSlotMap; int timeSlot = 0; @@ -768,8 +835,8 @@ void RockPipeline::runOnOperation() { SmallVector 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, diff --git a/mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir b/mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir index 93d7bfce7506..ad0edbaf0dc7 100644 --- a/mlir/test/Dialect/Rock/rock-pipeline-early-exit.mlir +++ b/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 diff --git a/mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir b/mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir new file mode 100644 index 000000000000..4a635f24abe4 --- /dev/null +++ b/mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir @@ -0,0 +1,271 @@ +// RUN: rocmlir-opt %s --rock-pipeline="rock-pipeline-remove-stages=false" | FileCheck %s +// RUN: rocmlir-opt %s --rock-pipeline="rock-pipeline-remove-stages=true" | FileCheck %s --check-prefix=REMOVE-STAGES + +// This test file verifies the optimization that skips backward LDS barriers +// for single-wave kernels with specific schedule versions. + +// Test for single-wave kernel with scheduleVersion=1 (Default) +// When blockSize <= waveSize and scheduleVersion is 1 or 3, backward barriers should be skipped +// For scheduleVersion=1, the loop has 3 stages: GlobalRead, LDSWrite, LDSRead +// CHECK-LABEL: rock_pipeline_one_wave_schedule_v1 +// REMOVE-STAGES-LABEL: rock_pipeline_one_wave_schedule_v1 +// For single-wave with scheduleVersion=1, we should NOT see a second barrier in the loop +// (backward barrier is skipped) +// Prologue stores to LDS: +// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: scf.for +// Inside loop - only ONE barrier (forward), no backward barrier for single-wave +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: rock.threadwise_gemm_accel +// REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: } +// Epilogue barrier and LDS read: +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return +func.func @rock_pipeline_one_wave_schedule_v1(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 64 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + + // 128 f16 elements = 256 bytes + %rawLds = rock.alloc() : memref<256xi8, #gpu.address_space> + %rawRegA = rock.alloc() : memref<32xi8, #gpu.address_space> + %rawRegB = rock.alloc() : memref<32xi8, #gpu.address_space> + %matrixA = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixB = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixC = memref.alloc() : memref<1x1xvector<4xf32>, #gpu.address_space> + + %lds = memref.view %rawLds[%c0][] : memref<256xi8, #gpu.address_space> to memref<128xf16, #gpu.address_space> + %regA = memref.view %rawRegA[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + %regB = memref.view %rawRegB[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + + scf.for %arg3 = %c0 to %c16 step %c1 { + rock.stage { + %a = memref.load %input[%arg3] : memref<16xf16, #gpu.address_space> + memref.store %a, %regA[%arg3] : memref<16xf16, #gpu.address_space> + rock.yield + }{name="GlobalRead"} + rock.stage { + %a = memref.load %regA[%arg3] : memref<16xf16, #gpu.address_space> + memref.store %a, %lds[%arg3] : memref<128xf16, #gpu.address_space> + rock.yield + }{name="LDSWrite"} + rock.stage { + %a = memref.load %lds[%arg3] : memref<128xf16, #gpu.address_space> + memref.store %a, %regB[%arg3] : memref<16xf16, #gpu.address_space> + %tid = rock.workitem_id : index + rock.threadwise_gemm_accel %matrixC += %matrixA * %matrixB at[%tid, %tid, %tid] { + params = #rock.accel_gemm_params< + kpackPerBlock = 4, mPerBlock = 16, nPerBlock = 16, kpack = 8, + mPerWave = 16, nPerWave = 16, mnPerXdl = 16, splitKFactor = 1, + scheduleVersion = 1, outputSwizzle = 2, wavesPerEU = 0, + gridGroupSize = 0, forceUnroll = true> + } : memref<1x1xvector<4xf32>, #gpu.address_space> += memref<1x2xvector<4xf16>, #gpu.address_space> * memref<1x2xvector<4xf16>, #gpu.address_space> + rock.yield + }{name="LDSRead"} + }{pipeline = #rock.pipeline<2>} + + %out = memref.load %regB[%c0] : memref<16xf16, #gpu.address_space> + memref.store %out, %output[%c0] : memref<16xf16, #gpu.address_space> + return +} + +// Test for single-wave kernel with scheduleVersion=3 (DirectToLDSDefault) +// When blockSize <= waveSize and scheduleVersion is 1 or 3, backward barriers should be skipped +// For scheduleVersion=3, the loop has only 2 stages: GlobalRead (writes directly to LDS) and LDSRead +// CHECK-LABEL: rock_pipeline_one_wave_schedule_v3 +// REMOVE-STAGES-LABEL: rock_pipeline_one_wave_schedule_v3 +// For single-wave with scheduleVersion=3, we should NOT see backward barrier +// The 2-stage loop doesn't fully pipeline but still gets barrier optimization +// REMOVE-STAGES: scf.for +// Inside loop - only ONE barrier for single-wave with scheduleVersion=3 +// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES-NEXT: rock.lds_barrier +// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: rock.threadwise_gemm_accel +// No second barrier before end of loop body +// REMOVE-STAGES: } +// No barriers after loop for this test since it doesn't fully pipeline +// REMOVE-STAGES-NOT: rock.lds_barrier +// REMOVE-STAGES: return +func.func @rock_pipeline_one_wave_schedule_v3(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 64 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + + // 128 f16 elements = 256 bytes + %rawLds = rock.alloc() : memref<256xi8, #gpu.address_space> + %rawRegA = rock.alloc() : memref<32xi8, #gpu.address_space> + %rawRegB = rock.alloc() : memref<32xi8, #gpu.address_space> + %matrixA = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixB = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixC = memref.alloc() : memref<1x1xvector<4xf32>, #gpu.address_space> + + %lds = memref.view %rawLds[%c0][] : memref<256xi8, #gpu.address_space> to memref<128xf16, #gpu.address_space> + %regA = memref.view %rawRegA[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + %regB = memref.view %rawRegB[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + + // For scheduleVersion=3 (DirectToLDS), there are only 2 stages: + // Stage 1: GlobalRead - loads from global and writes directly to LDS + // Stage 2: LDSRead - reads from LDS and performs MFMA + scf.for %arg3 = %c0 to %c16 step %c1 { + rock.stage { + // GlobalRead stage: load from global and write DIRECTLY to LDS (Direct-to-LDS) + %a = memref.load %input[%arg3] : memref<16xf16, #gpu.address_space> + memref.store %a, %lds[%arg3] : memref<128xf16, #gpu.address_space> + rock.yield + }{name="GlobalRead"} + rock.stage { + // LDSRead stage: read from LDS and perform MFMA + %a = memref.load %lds[%arg3] : memref<128xf16, #gpu.address_space> + memref.store %a, %regB[%arg3] : memref<16xf16, #gpu.address_space> + %tid = rock.workitem_id : index + rock.threadwise_gemm_accel %matrixC += %matrixA * %matrixB at[%tid, %tid, %tid] { + params = #rock.accel_gemm_params< + kpackPerBlock = 4, mPerBlock = 16, nPerBlock = 16, kpack = 8, + mPerWave = 16, nPerWave = 16, mnPerXdl = 16, splitKFactor = 1, + scheduleVersion = 3, outputSwizzle = 2, wavesPerEU = 0, + gridGroupSize = 0, forceUnroll = true> + } : memref<1x1xvector<4xf32>, #gpu.address_space> += memref<1x2xvector<4xf16>, #gpu.address_space> * memref<1x2xvector<4xf16>, #gpu.address_space> + rock.yield + }{name="LDSRead"} + }{pipeline = #rock.pipeline<2>} + + %out = memref.load %regB[%c0] : memref<16xf16, #gpu.address_space> + memref.store %out, %output[%c0] : memref<16xf16, #gpu.address_space> + return +} + +// Test for multi-wave kernel with scheduleVersion=1 - should still have backward barrier +// When blockSize > waveSize, backward barriers should NOT be skipped +// CHECK-LABEL: rock_pipeline_multi_wave_schedule_v1 +// REMOVE-STAGES-LABEL: rock_pipeline_multi_wave_schedule_v1 +// For multi-wave, we SHOULD see TWO barriers in the loop (forward and backward) +// Prologue stores to LDS: +// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: scf.for +// Inside loop - TWO barriers for multi-wave (forward + backward) +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: rock.threadwise_gemm_accel +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: } +// Epilogue barrier and LDS read: +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return +func.func @rock_pipeline_multi_wave_schedule_v1(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 128 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + + // 128 f16 elements = 256 bytes + %rawLds = rock.alloc() : memref<256xi8, #gpu.address_space> + %rawRegA = rock.alloc() : memref<32xi8, #gpu.address_space> + %rawRegB = rock.alloc() : memref<32xi8, #gpu.address_space> + %matrixA = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixB = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixC = memref.alloc() : memref<1x1xvector<4xf32>, #gpu.address_space> + + %lds = memref.view %rawLds[%c0][] : memref<256xi8, #gpu.address_space> to memref<128xf16, #gpu.address_space> + %regA = memref.view %rawRegA[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + %regB = memref.view %rawRegB[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + + scf.for %arg3 = %c0 to %c16 step %c1 { + rock.stage { + %a = memref.load %input[%arg3] : memref<16xf16, #gpu.address_space> + memref.store %a, %regA[%arg3] : memref<16xf16, #gpu.address_space> + rock.yield + }{name="GlobalRead"} + rock.stage { + %a = memref.load %regA[%arg3] : memref<16xf16, #gpu.address_space> + memref.store %a, %lds[%arg3] : memref<128xf16, #gpu.address_space> + rock.yield + }{name="LDSWrite"} + rock.stage { + %a = memref.load %lds[%arg3] : memref<128xf16, #gpu.address_space> + memref.store %a, %regB[%arg3] : memref<16xf16, #gpu.address_space> + %tid = rock.workitem_id : index + rock.threadwise_gemm_accel %matrixC += %matrixA * %matrixB at[%tid, %tid, %tid] { + params = #rock.accel_gemm_params< + kpackPerBlock = 4, mPerBlock = 16, nPerBlock = 16, kpack = 8, + mPerWave = 16, nPerWave = 16, mnPerXdl = 16, splitKFactor = 1, + scheduleVersion = 1, outputSwizzle = 2, wavesPerEU = 0, + gridGroupSize = 0, forceUnroll = true> + } : memref<1x1xvector<4xf32>, #gpu.address_space> += memref<1x2xvector<4xf16>, #gpu.address_space> * memref<1x2xvector<4xf16>, #gpu.address_space> + rock.yield + }{name="LDSRead"} + }{pipeline = #rock.pipeline<2>} + + %out = memref.load %regB[%c0] : memref<16xf16, #gpu.address_space> + memref.store %out, %output[%c0] : memref<16xf16, #gpu.address_space> + return +} + +// Test for single-wave kernel with scheduleVersion=2 (DoubleBuffer) - should still have backward barrier +// scheduleVersion=2 does NOT allow skipping backward barrier even for single-wave +// CHECK-LABEL: rock_pipeline_one_wave_schedule_v2 +// REMOVE-STAGES-LABEL: rock_pipeline_one_wave_schedule_v2 +// For scheduleVersion=2, we SHOULD see TWO barriers even with single-wave +// Prologue stores to LDS: +// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: scf.for +// Inside loop - TWO barriers even for single-wave with scheduleVersion=2 +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: rock.threadwise_gemm_accel +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// REMOVE-STAGES: } +// Epilogue barrier and LDS read: +// REMOVE-STAGES: rock.lds_barrier +// REMOVE-STAGES: return +func.func @rock_pipeline_one_wave_schedule_v2(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 64 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + + // 128 f16 elements = 256 bytes + %rawLds = rock.alloc() : memref<256xi8, #gpu.address_space> + %rawRegA = rock.alloc() : memref<32xi8, #gpu.address_space> + %rawRegB = rock.alloc() : memref<32xi8, #gpu.address_space> + %matrixA = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixB = memref.alloc() : memref<1x2xvector<4xf16>, #gpu.address_space> + %matrixC = memref.alloc() : memref<1x1xvector<4xf32>, #gpu.address_space> + + %lds = memref.view %rawLds[%c0][] : memref<256xi8, #gpu.address_space> to memref<128xf16, #gpu.address_space> + %regA = memref.view %rawRegA[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + %regB = memref.view %rawRegB[%c0][] : memref<32xi8, #gpu.address_space> to memref<16xf16, #gpu.address_space> + + scf.for %arg3 = %c0 to %c16 step %c1 { + rock.stage { + %a = memref.load %input[%arg3] : memref<16xf16, #gpu.address_space> + memref.store %a, %regA[%arg3] : memref<16xf16, #gpu.address_space> + rock.yield + }{name="GlobalRead"} + rock.stage { + %a = memref.load %regA[%arg3] : memref<16xf16, #gpu.address_space> + memref.store %a, %lds[%arg3] : memref<128xf16, #gpu.address_space> + rock.yield + }{name="LDSWrite"} + rock.stage { + %a = memref.load %lds[%arg3] : memref<128xf16, #gpu.address_space> + memref.store %a, %regB[%arg3] : memref<16xf16, #gpu.address_space> + %tid = rock.workitem_id : index + rock.threadwise_gemm_accel %matrixC += %matrixA * %matrixB at[%tid, %tid, %tid] { + params = #rock.accel_gemm_params< + kpackPerBlock = 4, mPerBlock = 16, nPerBlock = 16, kpack = 8, + mPerWave = 16, nPerWave = 16, mnPerXdl = 16, splitKFactor = 1, + scheduleVersion = 2, outputSwizzle = 2, wavesPerEU = 0, + gridGroupSize = 0, forceUnroll = true> + } : memref<1x1xvector<4xf32>, #gpu.address_space> += memref<1x2xvector<4xf16>, #gpu.address_space> * memref<1x2xvector<4xf16>, #gpu.address_space> + rock.yield + }{name="LDSRead"} + }{pipeline = #rock.pipeline<2>} + + %out = memref.load %regB[%c0] : memref<16xf16, #gpu.address_space> + memref.store %out, %output[%c0] : memref<16xf16, #gpu.address_space> + return +} diff --git a/mlir/test/e2e/CMakeLists.txt b/mlir/test/e2e/CMakeLists.txt index ca0d7c2473f3..269163ca26ad 100644 --- a/mlir/test/e2e/CMakeLists.txt +++ b/mlir/test/e2e/CMakeLists.txt @@ -53,6 +53,8 @@ if (ROCMLIR_DRIVER_PR_E2E_TEST_ENABLED) PrLdsTransposeLoadAttention PrConvDirectToLDS PrAttentionDirectToLDS + PrGemmOneWaveBarrier + PrGemmOneWaveBarrierDirectToLDS ) set(GEN_MODE "") endif() @@ -102,6 +104,9 @@ if (ROCK_E2E_TEST_ENABLED) AttentionNonPowerOfTwoTileSize LdsTransposeLoad LdsTransposeLoadAttention + GemmOneWaveBarrier + GemmOneWaveBarrierDirectToLDS + GemmOneWaveBarrierFp8 ) endif() # Create a list for dummy files diff --git a/mlir/test/e2e/GemmOneWaveBarrier.cfg b/mlir/test/e2e/GemmOneWaveBarrier.cfg new file mode 100644 index 000000000000..f03c401eaf22 --- /dev/null +++ b/mlir/test/e2e/GemmOneWaveBarrier.cfg @@ -0,0 +1,3 @@ +# Require MFMA or WMMA support +if not config.arch_support_mfma and not config.arch_support_wmma: + config.unsupported = True diff --git a/mlir/test/e2e/GemmOneWaveBarrier.toml b/mlir/test/e2e/GemmOneWaveBarrier.toml new file mode 100644 index 000000000000..3937e6b64a5e --- /dev/null +++ b/mlir/test/e2e/GemmOneWaveBarrier.toml @@ -0,0 +1,33 @@ +# One-wave GEMM barrier optimization test (scheduleVersion=1) + +directory = "GemmOneWaveBarrier" +prefix = "rocmlir-gen" +suffix = "--operation gemm --arch %arch %pv %random_data %rocmlir_gen_flags | rocmlir-driver -c | 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 | FileCheck %s --check-prefix=" + +[[axis]] +name = "transA" +values = ["true", "false"] +prefix = "--transA=" + +[[axis]] +name = "transB" +values = ["true", "false"] +prefix = "--transB=" + +[[axis]] +name = "data type" +values = ["f32", "f16", "bf16", "i8"] +prefix = "-t " + +[[axis]] +name = "perf_config" +# One-wave: mPerBlock=mPerWave=32, nPerBlock=nPerWave=32, kpackPerBlock=16, scheduleVersion=1 +values = ["v4:32,32,16,32,32,32,1,1,1,2,0,0,1,1"] +prefix = "-perf_config=" + +[[suite]] +name = "gemm_one_wave_barrier" + +# gridSize = (m/32) * (n/32) = (1024/32) * (1024/32) = 32 * 32 = 1024 > 4*numCU +[[suite.test]] +config = "-g 1 -m 1024 -k 64 -n 1024" diff --git a/mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.cfg b/mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.cfg new file mode 100644 index 000000000000..65cc870f6efe --- /dev/null +++ b/mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.cfg @@ -0,0 +1,5 @@ +# Require direct_to_lds feature for scheduleVersion 3, and fp8 support +if not 'direct_to_lds_32b' in config.features and not 'direct_to_lds_128b' in config.features: + config.unsupported = True +if not config.arch_support_accel_fp8: + config.unsupported = True diff --git a/mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.toml b/mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.toml new file mode 100644 index 000000000000..ec2b14a86a87 --- /dev/null +++ b/mlir/test/e2e/GemmOneWaveBarrierDirectToLDS.toml @@ -0,0 +1,33 @@ +# One-wave GEMM barrier optimization test (scheduleVersion=3, DirectToLDS) + +directory = "GemmOneWaveBarrierDirectToLDS" +prefix = "rocmlir-gen" +suffix = "--operation gemm --arch %arch %pv %random_data %rocmlir_gen_flags | rocmlir-driver -c | 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 | FileCheck %s --check-prefix=" + +[[axis]] +name = "transA" +values = ["true", "false"] +prefix = "--transA=" + +[[axis]] +name = "transB" +values = ["true", "false"] +prefix = "--transB=" + +[[axis]] +name = "data type" +values = ["f32", "f16", "bf16", "i8", "fp8_fp8"] +prefix = "-t " + +[[axis]] +name = "perf_config" +# One-wave: mPerBlock=mPerWave=32, nPerBlock=nPerWave=32, kpackPerBlock=16, scheduleVersion=3 +values = ["v4:32,32,16,32,32,32,1,1,3,2,0,0,1,1"] +prefix = "-perf_config=" + +[[suite]] +name = "gemm_one_wave_barrier_direct_to_lds" + +# gridSize = (m/32) * (n/32) = (1024/32) * (1024/32) = 32 * 32 = 1024 > 4*numCU +[[suite.test]] +config = "-g 1 -m 1024 -k 64 -n 1024" diff --git a/mlir/test/e2e/GemmOneWaveBarrierFp8.cfg b/mlir/test/e2e/GemmOneWaveBarrierFp8.cfg new file mode 100644 index 000000000000..617d8f44098f --- /dev/null +++ b/mlir/test/e2e/GemmOneWaveBarrierFp8.cfg @@ -0,0 +1,5 @@ +# Require MFMA or WMMA support and fp8 acceleration +if not config.arch_support_mfma and not config.arch_support_wmma: + config.unsupported = True +if not config.arch_support_accel_fp8: + config.unsupported = True diff --git a/mlir/test/e2e/GemmOneWaveBarrierFp8.toml b/mlir/test/e2e/GemmOneWaveBarrierFp8.toml new file mode 100644 index 000000000000..7bda6b27e790 --- /dev/null +++ b/mlir/test/e2e/GemmOneWaveBarrierFp8.toml @@ -0,0 +1,33 @@ +# One-wave GEMM barrier optimization test for fp8 (scheduleVersion=1) + +directory = "GemmOneWaveBarrierFp8" +prefix = "rocmlir-gen" +suffix = "--operation gemm --arch %arch %pv %random_data %rocmlir_gen_flags | rocmlir-driver -c | 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 | FileCheck %s --check-prefix=" + +[[axis]] +name = "transA" +values = ["true", "false"] +prefix = "--transA=" + +[[axis]] +name = "transB" +values = ["true", "false"] +prefix = "--transB=" + +[[axis]] +name = "data type" +values = ["fp8_fp8"] +prefix = "-t " + +[[axis]] +name = "perf_config" +# One-wave: mPerBlock=mPerWave=32, nPerBlock=nPerWave=32, kpackPerBlock=16, scheduleVersion=1 +values = ["v4:32,32,16,32,32,32,1,1,1,2,0,0,1,1"] +prefix = "-perf_config=" + +[[suite]] +name = "gemm_one_wave_barrier_fp8" + +# gridSize = (m/32) * (n/32) = (1024/32) * (1024/32) = 32 * 32 = 1024 > 4*numCU +[[suite.test]] +config = "-g 1 -m 1024 -k 64 -n 1024" diff --git a/mlir/test/e2e/PrGemmOneWaveBarrier.cfg b/mlir/test/e2e/PrGemmOneWaveBarrier.cfg new file mode 100644 index 000000000000..4428522278dc --- /dev/null +++ b/mlir/test/e2e/PrGemmOneWaveBarrier.cfg @@ -0,0 +1,3 @@ +# Require MFMA or WMMA support for accelerated GEMM +if not (config.arch_support_mfma or config.arch_support_wmma): + config.unsupported = True diff --git a/mlir/test/e2e/PrGemmOneWaveBarrier.toml b/mlir/test/e2e/PrGemmOneWaveBarrier.toml new file mode 100644 index 000000000000..73ec9be13121 --- /dev/null +++ b/mlir/test/e2e/PrGemmOneWaveBarrier.toml @@ -0,0 +1,22 @@ +# One-wave GEMM barrier optimization test (scheduleVersion=1) + +directory = "PrGemmOneWaveBarrier" +prefix = "rocmlir-gen" +suffix = "--operation gemm --arch %arch %pv %random_data %rocmlir_gen_flags | rocmlir-driver -c | 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 | FileCheck %s --check-prefix=" + +[[axis]] +name = "data type" +values = ["f16"] +prefix = "-t " + +[[axis]] +name = "perf_config" +# One-wave: mPerBlock=mPerWave=32, nPerBlock=nPerWave=32, scheduleVersion=1 +values = ["v4:32,32,8,32,32,32,4,1,1,2,0,0,1,1"] +prefix = "-perf_config=" + +[[suite]] +name = "pr_gemm_one_wave_barrier" + +[[suite.test]] +config = "-g 1 -m 32 -k 64 -n 32" diff --git a/mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.cfg b/mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.cfg new file mode 100644 index 000000000000..5f04a2150299 --- /dev/null +++ b/mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.cfg @@ -0,0 +1,3 @@ +# Require direct_to_lds feature for scheduleVersion 3 +if not 'direct_to_lds_32b' in config.features and not 'direct_to_lds_128b' in config.features: + config.unsupported = True diff --git a/mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.toml b/mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.toml new file mode 100644 index 000000000000..b05be858a0e5 --- /dev/null +++ b/mlir/test/e2e/PrGemmOneWaveBarrierDirectToLDS.toml @@ -0,0 +1,22 @@ +# One-wave GEMM barrier optimization test (scheduleVersion=3, DirectToLDS) + +directory = "PrGemmOneWaveBarrierDirectToLDS" +prefix = "rocmlir-gen" +suffix = "--operation gemm --arch %arch %pv %random_data %rocmlir_gen_flags | rocmlir-driver -c | 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 | FileCheck %s --check-prefix=" + +[[axis]] +name = "data type" +values = ["f16"] +prefix = "-t " + +[[axis]] +name = "perf_config" +# One-wave: mPerBlock=mPerWave=32, nPerBlock=nPerWave=32, scheduleVersion=3 +values = ["v4:32,32,8,32,32,32,4,1,3,2,0,0,1,1"] +prefix = "-perf_config=" + +[[suite]] +name = "pr_gemm_one_wave_barrier_direct_to_lds" + +[[suite.test]] +config = "-g 1 -m 32 -k 64 -n 32" From e08393294bc78849ec20ed3825be418f97dfca71 Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Tue, 24 Feb 2026 20:35:35 +0000 Subject: [PATCH 4/6] just use CHECK --- .../Rock/test_rock_pipeline_wave_barrier.mlir | 109 ++++++++---------- 1 file changed, 51 insertions(+), 58 deletions(-) diff --git a/mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir b/mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir index 4a635f24abe4..a8bca60ed7bf 100644 --- a/mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir +++ b/mlir/test/Dialect/Rock/test_rock_pipeline_wave_barrier.mlir @@ -1,5 +1,4 @@ -// RUN: rocmlir-opt %s --rock-pipeline="rock-pipeline-remove-stages=false" | FileCheck %s -// RUN: rocmlir-opt %s --rock-pipeline="rock-pipeline-remove-stages=true" | FileCheck %s --check-prefix=REMOVE-STAGES +// RUN: rocmlir-opt %s --rock-pipeline="rock-pipeline-remove-stages=true" | FileCheck %s // This test file verifies the optimization that skips backward LDS barriers // for single-wave kernels with specific schedule versions. @@ -7,23 +6,21 @@ // Test for single-wave kernel with scheduleVersion=1 (Default) // When blockSize <= waveSize and scheduleVersion is 1 or 3, backward barriers should be skipped // For scheduleVersion=1, the loop has 3 stages: GlobalRead, LDSWrite, LDSRead -// CHECK-LABEL: rock_pipeline_one_wave_schedule_v1 -// REMOVE-STAGES-LABEL: rock_pipeline_one_wave_schedule_v1 -// For single-wave with scheduleVersion=1, we should NOT see a second barrier in the loop -// (backward barrier is skipped) + +// CHECK-LABEL: func.func @rock_pipeline_one_wave_schedule_v1 // Prologue stores to LDS: -// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: scf.for -// Inside loop - only ONE barrier (forward), no backward barrier for single-wave -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: rock.threadwise_gemm_accel -// REMOVE-STAGES-NOT: rock.lds_barrier -// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: } +// CHECK: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: scf.for +// Inside loop - only ONE barrier (forward), no backward barrier for single-wave +// CHECK: rock.lds_barrier +// CHECK: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: rock.threadwise_gemm_accel +// CHECK-NOT: rock.lds_barrier +// CHECK: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: } // Epilogue barrier and LDS read: -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: return +// CHECK: rock.lds_barrier +// CHECK: return func.func @rock_pipeline_one_wave_schedule_v1(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 64 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -75,21 +72,19 @@ func.func @rock_pipeline_one_wave_schedule_v1(%input : memref<16xf16, #gpu.addre // Test for single-wave kernel with scheduleVersion=3 (DirectToLDSDefault) // When blockSize <= waveSize and scheduleVersion is 1 or 3, backward barriers should be skipped // For scheduleVersion=3, the loop has only 2 stages: GlobalRead (writes directly to LDS) and LDSRead -// CHECK-LABEL: rock_pipeline_one_wave_schedule_v3 -// REMOVE-STAGES-LABEL: rock_pipeline_one_wave_schedule_v3 -// For single-wave with scheduleVersion=3, we should NOT see backward barrier -// The 2-stage loop doesn't fully pipeline but still gets barrier optimization -// REMOVE-STAGES: scf.for -// Inside loop - only ONE barrier for single-wave with scheduleVersion=3 -// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES-NEXT: rock.lds_barrier -// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: rock.threadwise_gemm_accel + +// CHECK-LABEL: func.func @rock_pipeline_one_wave_schedule_v3 +// CHECK: scf.for +// Inside loop - only ONE barrier for single-wave with scheduleVersion=3 +// CHECK: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK-NEXT: rock.lds_barrier +// CHECK: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: rock.threadwise_gemm_accel // No second barrier before end of loop body -// REMOVE-STAGES: } +// CHECK: } // No barriers after loop for this test since it doesn't fully pipeline -// REMOVE-STAGES-NOT: rock.lds_barrier -// REMOVE-STAGES: return +// CHECK-NOT: rock.lds_barrier +// CHECK: return func.func @rock_pipeline_one_wave_schedule_v3(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 64 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -140,22 +135,21 @@ func.func @rock_pipeline_one_wave_schedule_v3(%input : memref<16xf16, #gpu.addre // Test for multi-wave kernel with scheduleVersion=1 - should still have backward barrier // When blockSize > waveSize, backward barriers should NOT be skipped -// CHECK-LABEL: rock_pipeline_multi_wave_schedule_v1 -// REMOVE-STAGES-LABEL: rock_pipeline_multi_wave_schedule_v1 -// For multi-wave, we SHOULD see TWO barriers in the loop (forward and backward) + +// CHECK-LABEL: func.func @rock_pipeline_multi_wave_schedule_v1 // Prologue stores to LDS: -// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: scf.for -// Inside loop - TWO barriers for multi-wave (forward + backward) -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: rock.threadwise_gemm_accel -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: } +// CHECK: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: scf.for +// Inside loop - TWO barriers for multi-wave (forward + backward) +// CHECK: rock.lds_barrier +// CHECK: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: rock.threadwise_gemm_accel +// CHECK: rock.lds_barrier +// CHECK: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: } // Epilogue barrier and LDS read: -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: return +// CHECK: rock.lds_barrier +// CHECK: return func.func @rock_pipeline_multi_wave_schedule_v1(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 128 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -206,22 +200,21 @@ func.func @rock_pipeline_multi_wave_schedule_v1(%input : memref<16xf16, #gpu.add // Test for single-wave kernel with scheduleVersion=2 (DoubleBuffer) - should still have backward barrier // scheduleVersion=2 does NOT allow skipping backward barrier even for single-wave -// CHECK-LABEL: rock_pipeline_one_wave_schedule_v2 -// REMOVE-STAGES-LABEL: rock_pipeline_one_wave_schedule_v2 -// For scheduleVersion=2, we SHOULD see TWO barriers even with single-wave + +// CHECK-LABEL: func.func @rock_pipeline_one_wave_schedule_v2 // Prologue stores to LDS: -// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: scf.for -// Inside loop - TWO barriers even for single-wave with scheduleVersion=2 -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: memref.load {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: rock.threadwise_gemm_accel -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: memref.store {{.*}} : memref<128xf16, #gpu.address_space> -// REMOVE-STAGES: } +// CHECK: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: scf.for +// Inside loop - TWO barriers even for single-wave with scheduleVersion=2 +// CHECK: rock.lds_barrier +// CHECK: memref.load {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: rock.threadwise_gemm_accel +// CHECK: rock.lds_barrier +// CHECK: memref.store {{.*}} : memref<128xf16, #gpu.address_space> +// CHECK: } // Epilogue barrier and LDS read: -// REMOVE-STAGES: rock.lds_barrier -// REMOVE-STAGES: return +// CHECK: rock.lds_barrier +// CHECK: return func.func @rock_pipeline_one_wave_schedule_v2(%input : memref<16xf16, #gpu.address_space>, %output : memref<16xf16, #gpu.address_space>) attributes {block_size = 64 : i32, arch = "amdgcn-amd-amdhsa:gfx90a"} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index From 25dd1ad312df6e176e4f6808e66394aec5421fc3 Mon Sep 17 00:00:00 2001 From: Umang Yadav <29876643+umangyadav@users.noreply.github.com> Date: Tue, 24 Feb 2026 15:38:56 -0500 Subject: [PATCH 5/6] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../Dialect/Rock/Transforms/RockPipeline.cpp | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp index bfc4f4e6c242..271052f67770 100644 --- a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp +++ b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp @@ -474,17 +474,21 @@ DagType pruneGraph(const DagType &dag) { // Determine if the backward barrier can be skipped for single-wave kernels. // -// For scheduleVersion 1 (Default) or 3 (DirectToLDSDefault), the loop -// structure is: +// For scheduleVersion 1 (Default), 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. +// For scheduleVersion 3 (DirectToLDSDefault), GlobalLoad writes directly to +// LDS, so the loop structure is logically: +// GlobalLoad (to LDS) -> (fwd barrier) -> DSRead + MFMA +// +// In both cases, the forward barrier ensures LDS writes (explicit DSWrite or +// DirectToLDS GlobalLoad) complete before DSReads start. For the +// loop-carried dependency (backward barrier), we need to ensure DSReads from +// iteration i finish before LDS writes 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 +// issued, they have read the data from the buffers, so LDS writes can proceed // without an explicit barrier. bool canSkipBackwardBarrierForOneWave(func::FuncOp func, scf::ForOp forOp) { // Check if this is a single-wave kernel @@ -500,15 +504,14 @@ bool canSkipBackwardBarrierForOneWave(func::FuncOp func, scf::ForOp forOp) { 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 + // For nested loops, it may require more analysis. For now, only support // single loop. int forOpCount = 0; func.walk([&](scf::ForOp) { ++forOpCount; }); From c82e4d54f96e7df9394aaeb9ddb90473f8f0fb53 Mon Sep 17 00:00:00 2001 From: Umang Yadav Date: Mon, 23 Mar 2026 15:13:32 +0000 Subject: [PATCH 6/6] Merge fixes --- mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp index 4b99ba75af26..8f53fe35763c 100644 --- a/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp +++ b/mlir/lib/Dialect/Rock/Transforms/RockPipeline.cpp @@ -508,7 +508,6 @@ bool canSkipBackwardBarrierForOneWave(func::FuncOp func, scf::ForOp forOp) { StringAttr arch = rock::getArchValue(func); - int64_t waveSize = rock::lookupArchInfo(arch).waveSize; bool isOneWave = (blockSize <= waveSize); if (!isOneWave) @@ -841,8 +840,8 @@ void RockPipeline::runOnOperation() { SmallVector 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, @@ -899,7 +898,6 @@ void RockPipeline::runOnOperation() { patterns.add(&getContext()); if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) return signalPassFailure(); - } } } }