[Codegen][GPU] Enable 3-stage pipelining with hipblaslt compute->write->read ordering (#22788)
3-stage uses different cluster ordering to maximize distance between
read and use. Stage definitions remain the same. The changes are mostly
trivial and the new pipelined code work out of box.
On MI300, the tracked gemm configs improve performance by 4.5% on
average for NN and NT layouts.
Right now there's no easy way to experiment this through a compiler
flag, and the only way to experiment this is to change the tablegen
default number of stage to 3. I plan to expose it properly in the next
PR.
---------
Signed-off-by: jerryyin <zhuoryin@amd.com>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/ROCDLPrefetchSharedMemoryCopy.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/ROCDLPrefetchSharedMemoryCopy.cpp
index d5cf5c5..863ab7a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/ROCDLPrefetchSharedMemoryCopy.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/ROCDLPrefetchSharedMemoryCopy.cpp
@@ -371,7 +371,7 @@
};
if (numStages == 2) {
- // Two-stage pipelining: readStage in stage 0, compute+write in stage 1.
+ // Two-stage pipelining: read+write in stage 0, compute in stage 1.
for (Operation *op : stages.readStage)
assignOp(op, /*stage=*/0);
for (Operation *op : stages.writeStage)
@@ -379,6 +379,8 @@
for (Operation *op : stages.computeStage)
assignOp(op, /*stage=*/1);
} else {
+ // Three-stage pipelining: read in stage 0, write in stage 1, compute in
+ // stage 2.
for (Operation *op : stages.readStage)
assignOp(op, /*stage=*/0);
for (Operation *op : stages.writeStage)
@@ -389,32 +391,54 @@
}
// Populates cluster IDs for each operation based on stage groupings.
-// Operations are assigned cluster IDs in execution order: read -> compute ->
-// write. Each stage group gets its own cluster ID.
+// Cluster ordering determines execution order within each iteration:
+// - 2-stage pipeline: read -> compute -> write
+// (read i+1, compute i, write i+1)
+// - 3-stage pipeline: compute -> write -> read
+// (compute i, write i+1, read i+2)
static void
-populateOpToClusterMap(const StageClassification &stages,
+populateOpToClusterMap(const StageClassification &stages, unsigned numStages,
llvm::DenseMap<Operation *, unsigned> &opToCluster) {
-
unsigned clusterID = 0;
- // Assign cluster ID 0 to all read stage operations
- for (Operation *op : stages.readStage) {
- opToCluster[op] = clusterID;
- }
- ++clusterID;
+ if (numStages == 2) {
+ // 2-stage pipeline: read first, then compute, then write
+ // This allows reading for next iteration while computing current
+ for (Operation *op : stages.readStage) {
+ opToCluster[op] = clusterID;
+ }
+ ++clusterID;
- // Assign cluster ID 1 to all compute stage operations
- for (Operation *op : stages.computeStage) {
- opToCluster[op] = clusterID;
- }
- ++clusterID;
+ for (Operation *op : stages.computeStage) {
+ opToCluster[op] = clusterID;
+ }
+ ++clusterID;
- // Assign cluster ID 2 to all write stage operations
- for (Operation *op : stages.writeStage) {
- opToCluster[op] = clusterID;
+ for (Operation *op : stages.writeStage) {
+ opToCluster[op] = clusterID;
+ }
+ ++clusterID;
+ } else {
+ // 3-stage pipeline: compute first, then write, then read
+ // This maximizes distance between read and use
+ for (Operation *op : stages.computeStage) {
+ opToCluster[op] = clusterID;
+ }
+ ++clusterID;
+
+ for (Operation *op : stages.writeStage) {
+ opToCluster[op] = clusterID;
+ }
+ ++clusterID;
+
+ for (Operation *op : stages.readStage) {
+ opToCluster[op] = clusterID;
+ }
+ ++clusterID;
}
- LDBG() << "Built opToCluster map with " << clusterID + 1 << " clusters";
+ LDBG() << "Built opToCluster map with " << clusterID << " clusters "
+ << "(numStages=" << numStages << ")";
}
// Builds the final schedule using opToCluster and opToStage mappings.
@@ -622,11 +646,12 @@
return forOp;
}
- // Multi-stage pipelining (numStages > 2) is not yet implemented.
- if (numStages > 2) {
- LDBG()
- << "Multi-stage pipelining with numStages > 2 is not yet implemented";
- return failure();
+ // For global->shared->register data flow, we have 3 operation groups (read,
+ // write, compute), so 3 stages is the maximum meaningful pipeline depth.
+ if (numStages > 3) {
+ LDBG() << "numStages=" << numStages
+ << " requested but capping to 3 (maximum for read, write, compute)";
+ numStages = 3;
}
// Compute stage classification using the new refactored approach
@@ -645,7 +670,7 @@
// Step 1: Populate standalone opToCluster map
llvm::DenseMap<Operation *, unsigned> opToCluster;
- populateOpToClusterMap(stages, opToCluster);
+ populateOpToClusterMap(stages, numStages, opToCluster);
// Step 2: Use opToCluster + opToStage to build the final schedule
std::vector<std::pair<Operation *, unsigned>> finalSchedule;
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/prefetch_shared_memory.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/prefetch_shared_memory.mlir
index 4f7fdc9..3f5aff2 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/prefetch_shared_memory.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/prefetch_shared_memory.mlir
@@ -1,5 +1,6 @@
// RUN: iree-opt -pass-pipeline="builtin.module(func.func(iree-llvmgpu-prefetch-shared-memory),cse,canonicalize)" %s --split-input-file | FileCheck %s --check-prefixes=CHECK-ALL,CHECK
// RUN: iree-opt -pass-pipeline="builtin.module(func.func(iree-llvmgpu-prefetch-shared-memory{num-stages=1}))" %s --split-input-file | FileCheck %s --check-prefixes=CHECK-ALL,CHECK-1STAGE
+// RUN: iree-opt -pass-pipeline="builtin.module(func.func(iree-llvmgpu-prefetch-shared-memory{num-stages=3}))" %s --split-input-file | FileCheck %s --check-prefixes=CHECK-ALL,CHECK-3STAGE
// CHECK-ALL-LABEL: @prefetch_add
// CHECK-SAME: (%[[GLOBAL:.*]]: memref<128xf32>)
@@ -20,7 +21,16 @@
// CHECK: vector.transfer_write %[[PRO_READ]], %[[SHARED]]
// CHECK: %[[OUT:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C127]] step %[[C1]] iter_args(%[[ARG:.*]] = %[[CST]])
// CHECK-1STAGE: scf.for %{{.*}} = %c0 to %c128 step %c1
+ // 3-stage prologue: 2 reads
+ // CHECK-3STAGE: vector.transfer_read %arg0
+ // CHECK-3STAGE: gpu.barrier
+ // CHECK-3STAGE: vector.transfer_write
+ // CHECK-3STAGE: vector.transfer_read %arg0
+ // CHECK-3STAGE: arith.constant 2 : index
+ // CHECK-3STAGE: arith.subi %c128
+ // CHECK-3STAGE: scf.for %{{.*}} = %c0 to %{{.*}} step %c1 iter_args
%0 = scf.for %arg1 = %c0 to %c128 step %c1 iter_args(%arg2 = %cst) -> (vector<1xf32>) {
+ // 2-stage ordering: read -> compute -> write
// CHECK-DAG: %[[IVPLUS1:.*]] = arith.addi %[[IV]], %[[C1]] : index
// CHECK: %[[KER_READ:.*]] = vector.transfer_read %[[GLOBAL]][%[[IVPLUS1]]]
%1 = vector.transfer_read %arg0[%arg1], %cst_0 : memref<128xf32>, vector<1xf32>
@@ -34,12 +44,34 @@
// CHECK: amdgpu.sched_barrier allow = <none>
// CHECK: vector.transfer_write %[[KER_READ]], %[[SHARED]]
// CHECK: scf.yield %[[COMPUTE]]
+
+ // 3-stage ordering: compute -> write -> read
+ // CHECK-3STAGE: gpu.barrier
+ // CHECK-3STAGE: vector.transfer_read %alloc
+ // CHECK-3STAGE: arith.addf
+ // CHECK-3STAGE: gpu.barrier
+ // CHECK-3STAGE: amdgpu.sched_barrier allow = <none>
+ // CHECK-3STAGE: vector.transfer_write
+ // CHECK-3STAGE: arith.constant 2
+ // CHECK-3STAGE: arith.addi
+ // CHECK-3STAGE: vector.transfer_read %arg0
+ // CHECK-3STAGE: scf.yield
scf.yield %3 : vector<1xf32>
}
+ // 2-stage epilogue: 1 iteration
// CHECK: gpu.barrier
// CHECK: %[[EPI_READ:.*]] = vector.transfer_read %[[SHARED]][%[[C0]]]
// CHECK: %[[EPI_COMPUTE:.*]] = arith.addf %[[EPI_READ]], %[[OUT]]
// CHECK: vector.transfer_write %[[EPI_COMPUTE]], %[[GLOBAL]][%[[C0]]]
+
+ // 3-stage epilogue: 2 iterations
+ // CHECK-3STAGE: gpu.barrier
+ // CHECK-3STAGE: vector.transfer_read %alloc
+ // CHECK-3STAGE: arith.addf
+ // CHECK-3STAGE: vector.transfer_write
+ // CHECK-3STAGE: vector.transfer_read %alloc
+ // CHECK-3STAGE: arith.addf
+ // CHECK-3STAGE: vector.transfer_write
vector.transfer_write %0, %arg0[%c0] {in_bounds = [true]} : vector<1xf32>, memref<128xf32>
return
}
