[SPIRV] Do not multibuffer GEMMs with fused leading ops (#14712)
Fixes #14711
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
index 9af72af..33f49ce 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -97,6 +97,28 @@
return fusedWithOp;
}
+// Check if there is a fused linalg op present in the backward slice of either
+// input.
+static bool hasFusedLeadingOp(linalg::LinalgOp matmul) {
+ assert(matmul.getNumDpsInputs() == 2 && "matmul expected to have two inputs");
+
+ BackwardSliceOptions options;
+ options.inclusive = true;
+
+ // Get the backward slice of the lhs.
+ SetVector<Operation *> backwardSlice;
+ getBackwardSlice(matmul->getOperand(0), &backwardSlice, options);
+
+ // Get the backward slice of the rhs and take the union.
+ SetVector<Operation *> rhsBackwardSlice;
+ getBackwardSlice(matmul->getOperand(1), &rhsBackwardSlice, options);
+ backwardSlice.set_union(rhsBackwardSlice);
+
+ return llvm::any_of(backwardSlice, [](Operation *op) {
+ return llvm::isa<linalg::LinalgOp>(op);
+ });
+}
+
//===----------------------------------------------------------------------===//
// Convolution Default Configuration
//===----------------------------------------------------------------------===//
@@ -730,6 +752,12 @@
storeStage = 1;
}
+ // TODO: Enable multibuffering with leading elementwise.
+ if (hasFusedLeadingOp(op)) {
+ pipelineDepth = 0;
+ storeStage = 1;
+ }
+
// Try to adjust tiling sizes to fit in shared memory.
auto usePromotionPipeline =
enablePromotion &&
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_amd_matmul.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_amd_matmul.mlir
index e894ae2..39f1a68 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_amd_matmul.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_amd_matmul.mlir
@@ -185,3 +185,78 @@
// CHECK: func.func @batch_matmul_f16_1x4096x4096x512()
// CHECK: linalg.batch_matmul
// CHECK-SAME: lowering_config = #[[CONFIG]]
+
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
+ #hal.descriptor_set.layout<0, bindings = [
+ #hal.descriptor_set.binding<0, storage_buffer>,
+ #hal.descriptor_set.binding<1, storage_buffer>,
+ #hal.descriptor_set.binding<2, storage_buffer>
+ ]>
+]>
+hal.executable @matmul_multi_reduce_i4xf32xf32 {
+ hal.executable.variant @vulkan_spirv_fb, target = <"vulkan", "vulkan-spirv-fb", {
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.6, [Shader], []>, AMD:DiscreteGPU, #spirv.resource_limits<
+ max_compute_shared_memory_size = 65536,
+ max_compute_workgroup_invocations = 1024,
+ max_compute_workgroup_size = [1024, 1024, 1024],
+ subgroup_size = 64>>
+ }> {
+ hal.executable.export @matmul_multi_reduce_i4xf32xf32 layout(#pipeline_layout)
+ builtin.module {
+ func.func @matmul_multi_reduce_i4xf32xf32() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = hal.interface.constant.load[0] : i32
+ %1 = hal.interface.constant.load[1] : i32
+ %2 = hal.interface.constant.load[2] : i32
+ %3 = hal.interface.constant.load[3] : i32
+ %4 = hal.interface.constant.load[4] : i32
+ %5 = arith.index_castui %0 : i32 to index
+ %6 = arith.index_castui %1 : i32 to index
+ %7 = arith.index_castui %2 : i32 to index
+ %8 = arith.index_castui %3 : i32 to index
+ %9 = arith.index_castui %4 : i32 to index
+ %10 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%5) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<11008x32x128xi4>>
+ %11 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%6) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<11008x32xf32>>
+ %12 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%7) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<11008x32xf32>>
+ %13 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%8) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<512x32x128xf32>>
+ %14 = hal.interface.binding.subspan set(0) binding(3) type(storage_buffer) alignment(64) offset(%9) : !flow.dispatch.tensor<writeonly:tensor<512x11008xf32>>
+ %15 = flow.dispatch.tensor.load %10, offsets = [0, 0, 0], sizes = [11008, 32, 128], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<11008x32x128xi4>> -> tensor<11008x32x128xi4>
+ %16 = flow.dispatch.tensor.load %11, offsets = [0, 0], sizes = [11008, 32], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<11008x32xf32>> -> tensor<11008x32xf32>
+ %17 = flow.dispatch.tensor.load %12, offsets = [0, 0], sizes = [11008, 32], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<11008x32xf32>> -> tensor<11008x32xf32>
+ %18 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0], sizes = [512, 32, 128], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<512x32x128xf32>> -> tensor<512x32x128xf32>
+ %19 = tensor.empty() : tensor<512x11008xf32>
+ %20 = tensor.empty() : tensor<11008x32x128xf32>
+ %21 = linalg.fill ins(%cst : f32) outs(%19 : tensor<512x11008xf32>) -> tensor<512x11008xf32>
+ %22 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%15, %16, %17 : tensor<11008x32x128xi4>, tensor<11008x32xf32>, tensor<11008x32xf32>) outs(%20 : tensor<11008x32x128xf32>) {
+ ^bb0(%in: i4, %in_0: f32, %in_1: f32, %out: f32):
+ %24 = arith.extui %in : i4 to i32
+ %25 = arith.uitofp %24 : i32 to f32
+ %26 = arith.subf %25, %in_1 : f32
+ %27 = arith.mulf %26, %in_0 : f32
+ linalg.yield %27 : f32
+ } -> tensor<11008x32x128xf32>
+ %23 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%18, %22 : tensor<512x32x128xf32>, tensor<11008x32x128xf32>) outs(%21 : tensor<512x11008xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %24 = arith.mulf %in, %in_0 : f32
+ %25 = arith.addf %24, %out : f32
+ linalg.yield %25 : f32
+ } -> tensor<512x11008xf32>
+ flow.dispatch.tensor.store %23, %14, offsets = [0, 0], sizes = [512, 11008], strides = [1, 1] : tensor<512x11008xf32> -> !flow.dispatch.tensor<writeonly:tensor<512x11008xf32>>
+ return
+ }
+ }
+ }
+}
+
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 128, 1, 16]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<SPIRVMatmulPromoteVectorize>
+// CHECK: hal.executable.export public @matmul_multi_reduce_i4xf32xf32
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK-SAME: workgroup_size = [32 : index, 8 : index, 1 : index]
+// CHECK: func.func @matmul_multi_reduce_i4xf32xf32()
+// CHECK: linalg.generic
+// CHECK: linalg.generic
+// CHECK-SAME: lowering_config = #[[CONFIG]]
