[Codegen] Allow multiple reduction dimensions in VectorDistribution (#18868)
This PR adds support for multiple k dimensions in VectorDistribution
contract codegen.
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
index 45b3f4b..ff6493b 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
@@ -1280,9 +1280,6 @@
llvm::errs() << "Getting mma layouts for:\n" << contractOp << "\n";
llvm::errs() << "For schedule: " << *this << "\n";
});
- if (opInfo.getKDims().size() != 1) {
- return contractOp->emitError("Unimplemented: > 1 k dims");
- }
int64_t rank = contractOp.getIteratorTypesArray().size();
auto mmaAttr = llvm::cast<MMAAttr>(getIntrinsic());
@@ -1450,6 +1447,10 @@
aSubgroupSizes[dim] = subgroupMBasis[i];
aSubgroupStrides[dim] = subgroupMStrides[i];
}
+ for (auto [kDim, lhsKDim] :
+ llvm::zip_equal(opInfo.getKDims(), opInfo.lhsKDim)) {
+ aBatchSizes[lhsKDim] = bounds[kDim];
+ }
aBatchSizes[afk] = bounds[opInfo.getKDims().back()] / intrinsicK;
auto aLayout = createNestedLayout(context, aRank, afm, afk,
@@ -1470,6 +1471,10 @@
bSubgroupSizes[dim] = subgroupNBasis[i];
bSubgroupStrides[dim] = subgroupNStrides[i];
}
+ for (auto [kDim, rhsKDim] :
+ llvm::zip_equal(opInfo.getKDims(), opInfo.rhsKDim)) {
+ bBatchSizes[rhsKDim] = bounds[kDim];
+ }
bBatchSizes[bfk] = bounds[opInfo.getKDims().back()] / intrinsicK;
auto bLayout = createNestedLayout(context, bRank, bfk, bfn,
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_gfx940.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_gfx940.mlir
index 1c460f7..610e114 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_gfx940.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_gfx940.mlir
@@ -166,6 +166,55 @@
// -----
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+
+hal.executable @matmul_multiple_k {
+ hal.executable.variant public @rocm target(<"rocm", "rocm-hsaco-fb">) {
+ hal.executable.export public @matmul_multiple_k 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 @matmul_multiple_k() attributes {translation_info = #iree_codegen.translation_info<LLVMGPUVectorDistribute workgroup_size = [256, 1, 1] subgroup_size = 64, {mma_schedule = #iree_gpu.mma_schedule<intrinsic = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>, subgroup_m_count = 1, subgroup_n_count = 4>}>} {
+ %cst = arith.constant 0.000000e+00 : f16
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(<bindings = [#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">, #hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">, #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>) binding(0) alignment(64) offset(%c0) flags("ReadOnly|Indirect") : !flow.dispatch.tensor<readonly:tensor<2x128x64x2048xf16>>
+ %1 = hal.interface.binding.subspan layout(<bindings = [#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">, #hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">, #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>) binding(1) alignment(64) offset(%c0) flags("ReadOnly|Indirect") : !flow.dispatch.tensor<readonly:tensor<10x128x64x2048xf16>>
+ %2 = hal.interface.binding.subspan layout(<bindings = [#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">, #hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">, #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>) binding(2) alignment(64) offset(%c0) flags(Indirect) : !flow.dispatch.tensor<writeonly:tensor<2x10x64x64xf16>>
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2, 128, 64, 2048], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2x128x64x2048xf16>> -> tensor<2x128x64x2048xf16>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0, 0], sizes = [10, 128, 64, 2048], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<10x128x64x2048xf16>> -> tensor<10x128x64x2048xf16>
+ %5 = tensor.empty() : tensor<2x10x64x64xf16>
+ %6 = linalg.fill ins(%cst : f16) outs(%5 : tensor<2x10x64x64xf16>) -> tensor<2x10x64x64xf16>
+ %7 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d4, d2, d5)>, affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d4, d3, d5)>, affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%3, %4 : tensor<2x128x64x2048xf16>, tensor<10x128x64x2048xf16>) outs(%6 : tensor<2x10x64x64xf16>) attrs = {lowering_config = #iree_gpu.lowering_config<{reduction = [0, 0, 0, 0, 1, 128], workgroup = [1, 1, 64, 64, 0, 0]}>} {
+ ^bb0(%in: f16, %in_0: f16, %out: f16):
+ %8 = arith.mulf %in, %in_0 : f16
+ %9 = arith.addf %8, %out : f16
+ linalg.yield %9 : f16
+ } -> tensor<2x10x64x64xf16>
+ flow.dispatch.tensor.store %7, %2, offsets = [0, 0, 0, 0], sizes = [2, 10, 64, 64], strides = [1, 1, 1, 1] : tensor<2x10x64x64xf16> -> !flow.dispatch.tensor<writeonly:tensor<2x10x64x64xf16>>
+ return
+ }
+ }
+ }
+}
+
+// Check if we can handle multiple reduction dimensions and that they generate
+// one coalesced loop.
+
+// CHECK-LABEL: func.func @matmul_multiple_k
+// CHECK: scf.for %[[IV:.+]] = %c0 to %c2048 step %c1
+// CHECK: affine.delinearize_index %[[IV]] into (%c128, %c16)
+// CHECK-COUNT-32: amdgpu.mfma
+// CHECK: scf.yield
+// CHECK-COUNT-4: vector.transfer_write {{.+}} {in_bounds = [true, true]} : vector<4x1xf16>, memref<2x10x64x64xf16, #hal.descriptor_type<storage_buffer>>
+
+// -----
+
// Basic f8, f8 -> f32 matmul.
#config = #iree_gpu.lowering_config<{workgroup = [64, 64, 0], reduction = [0, 0, 256]}>
diff --git a/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.cpp
index 83ffaf8..94cc749 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.cpp
@@ -18,7 +18,7 @@
// Returns the (LHS K, RHS K) dimension index pair.
std::pair<int, int> VectorContractOpInfo::getOperandKIndex() const {
- return std::make_pair(lhsKDim, rhsKDim);
+ return std::make_pair(lhsKDim.back(), rhsKDim.back());
}
// Returns the result (M, N) dimension index pair.
@@ -55,9 +55,12 @@
opInfo.outNDims.push_back(
*maps[2].getResultPosition(getAffineDimExpr(n, ctx)));
}
- int64_t k = contractionDims.k.back();
- opInfo.lhsKDim = *maps[0].getResultPosition(getAffineDimExpr(k, ctx));
- opInfo.rhsKDim = *maps[1].getResultPosition(getAffineDimExpr(k, ctx));
+ for (auto k : contractionDims.k) {
+ opInfo.lhsKDim.push_back(
+ *maps[0].getResultPosition(getAffineDimExpr(k, ctx)));
+ opInfo.rhsKDim.push_back(
+ *maps[1].getResultPosition(getAffineDimExpr(k, ctx)));
+ }
opInfo.lhsUnitDims = maps[0].getBroadcastDims();
opInfo.rhsUnitDims = maps[1].getBroadcastDims();
diff --git a/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.h b/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.h
index 101bf27..b8bde25 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/VectorOpUtils.h
@@ -49,9 +49,9 @@
int64_t getBatchCount() const { return contractionDims.batch.size(); }
SmallVector<int64_t> lhsMDims;
- int64_t lhsKDim;
+ SmallVector<int64_t> lhsKDim;
SmallVector<int64_t> rhsNDims;
- int64_t rhsKDim;
+ SmallVector<int64_t> rhsKDim;
SmallVector<int64_t> outMDims;
SmallVector<int64_t> outNDims;
