[LLVMGPU] Add tests for Attention subgroup reduction pipeline (#19401)
Also removes an untested and wrong visitor in VectorLayoutAnalysis and a
fix for 0-d transfer_write distribution.
diff --git a/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp b/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
index aeac222..5030020 100644
--- a/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
@@ -136,9 +136,6 @@
RegionBranchPoint branchPoint,
MutableArrayRef<OpOperand> operands);
- void visitRegionBranchTerminatorOpInterface(RegionBranchOpInterface branch,
- RegionBranchPoint branchPoint);
-
DistributionLayout *getLatticeElement(Value val);
MLIRContext *ctx;
@@ -1006,9 +1003,6 @@
if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
visitRegionSuccessors(branch, RegionBranchPoint::parent(),
branch->getOpOperands());
-
- // Handle the propagation from scf.for to yield op.
- visitRegionBranchTerminatorOpInterface(branch, RegionBranchPoint::parent());
return;
}
@@ -1101,43 +1095,6 @@
}
}
-void EnforceLayout::visitRegionBranchTerminatorOpInterface(
- RegionBranchOpInterface branch, RegionBranchPoint branchPoint) {
- SmallVector<RegionSuccessor> successors;
- branch.getSuccessorRegions(branchPoint, successors);
- if (!branch.hasLoop())
- return;
- SmallVector<DistributionLayout *> resultLattices;
- for (Value result : branch->getResults()) {
- DistributionLayout *resultLattice = getLatticeElement(result);
- if (resultLattice->isUninitialized())
- continue;
- resultLattices.push_back(resultLattice);
- }
-
- // We do not support multiple results yet.
- if (resultLattices.size() != 1)
- return;
-
- for (RegionSuccessor successor : successors) {
- if (Region *succ = successor.getSuccessor()) {
- Operation *terminator = succ->back().getTerminator();
- if (scf::YieldOp yieldOp = dyn_cast<scf::YieldOp>(terminator)) {
- for (Value operand : yieldOp.getOperands()) {
- if (!isa<VectorType>(operand.getType())) {
- continue;
- }
- DistributionLayout *forwardLattice = getLatticeElement(operand);
- ChangeResult changed = forwardLattice->resolve(resultLattices[0]);
- propagateIfChanged(forwardLattice, changed);
- }
- }
- }
- }
-
- return;
-}
-
DistributionLayout *EnforceLayout::getLatticeElement(Value val) {
// Add dependency of operation on the analysis state.
assert(isa<VectorType>(val.getType()) && "Lattice value should be a vector");
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_reduction_gfx942.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_reduction_gfx942.mlir
index 6b6bcd8..62e1738 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_reduction_gfx942.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_reduction_gfx942.mlir
@@ -178,3 +178,98 @@
// CHECK: scf.yield
// CHECK: vector.transfer_write
+
+// -----
+
+// Tile Sizes:
+// workgroup = [1, 1, 0, 0, 32]
+// reduction = [0, 0, 0, 128, 0]
+// subgroup = [0, 0, 0, 0, 8]
+// thread = [0, 0, 32, 4, 4]
+
+// Counts:
+// subgroup = [1, 1, 1, 1, 4]
+// threads = [1, 1, 2, 32, 1]
+
+#config = #iree_gpu.lowering_config<{workgroup = [1, 1, 0, 0, 32],
+ reduction = [0, 0, 0, 128, 0],
+ promote_operands = [1, 2]}>
+
+#qk_config = #iree_gpu.lowering_config<{subgroup_basis = [[1, 1, 1, 1, 4], [0, 1, 2, 3]],
+ thread_basis = [[1, 1, 2, 32, 1], [0, 1, 2, 3]],
+ thread = [0, 0, 32, 4],
+ promote_operands = [1]}>
+
+#pv_config = #iree_gpu.lowering_config<{subgroup_basis = [[1, 1, 1, 1, 4], [0, 1, 3, 4]],
+ thread_basis = [[1, 1, 2, 32, 1], [0, 1, 3, 4]],
+ thread = [0, 0, 4, 4],
+ promote_operands = [1]}>
+
+#translation = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute
+ workgroup_size = [256, 1, 1]
+ subgroup_size = 64>
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+hal.executable private @attention_20x1x64x4096x64 {
+ hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb">) {
+ hal.executable.export public @attention_20x1x64x4096x64 ordinal(0) layout(#pipeline_layout) {
+ ^bb0(%arg0: !hal.device):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_slice
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @attention_20x1x64x4096x64() attributes {translation_info = #translation} {
+ %cst = arith.constant 1.250000e-01 : f16
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<20x1x64xf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(2) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>>
+ %3 = hal.interface.binding.subspan layout(#pipeline_layout) binding(3) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<20x1x64xf16>>
+ %4 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0], sizes = [20, 1, 64], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<20x1x64xf16>> -> tensor<20x1x64xf16>
+ %5 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0], sizes = [20, 4096, 64], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>> -> tensor<20x4096x64xf16>
+ %6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0], sizes = [20, 4096, 64], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>> -> tensor<20x4096x64xf16>
+ %7 = tensor.empty() : tensor<20x1x64xf16>
+ %8 = iree_linalg_ext.attention {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2)>,
+ affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>,
+ affine_map<(d0, d1, d2, d3, d4) -> ()>,
+ affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d4)>],
+ lowering_config = #config,
+ decomposition_config = {
+ qk_attrs = {lowering_config = #qk_config},
+ pv_attrs = {lowering_config = #pv_config}
+ }}
+ ins(%4, %5, %6, %cst : tensor<20x1x64xf16>, tensor<20x4096x64xf16>, tensor<20x4096x64xf16>, f16) outs(%7 : tensor<20x1x64xf16>) {
+ ^bb0(%score: f32):
+ iree_linalg_ext.yield %score : f32
+ } -> tensor<20x1x64xf16>
+ flow.dispatch.tensor.store %8, %3, offsets = [0, 0, 0], sizes = [20, 1, 64], strides = [1, 1, 1] : tensor<20x1x64xf16> -> !flow.dispatch.tensor<writeonly:tensor<20x1x64xf16>>
+ return
+ }
+ }
+ }
+}
+
+// CHECK-LABEL: func.func @attention_20x1x64x4096x64
+// CHECK: scf.for %{{.*}} = %c0 to %c4096 step %c128
+// QK Matmul
+// CHECK: vector.contract
+// CHECK-SAME: vector<1x1x32xf16>, vector<1x1x1x1x4x32xf16> into vector<1x1x4xf32>
+// CHECK-COUNT-4: gpu.subgroup_reduce add
+
+// QK Max
+// CHECK-COUNT-1: gpu.subgroup_reduce maximumf
+
+// PV Sum
+// CHECK-COUNT-1: gpu.subgroup_reduce add
+
+// PV Matmul
+// CHECK: vector.contract
+// CHECK-COUNT-8: gpu.subgroup_reduce add
+
+// CHECK: scf.yield
