[CodeGen][SPIRV] Improved Matmul Tiling Heuristic for VideoCoreVII (#24559)
The generic SPIR-V matmul configuration is not tuned for the VideoCore
VII. Empirically, the heuristic added here yields on-average better
results on ResNet-style workloads by choosing workgroup/thread/reduction
tilings that seem to favor this GPU.
Only matmul ops with 8/16/32-bit element types and 2D/3D iteration
(matmul / batch matmul) are handled; anything else returns `failure()`
and falls through to the default SPIRV configuration.
One that seems to get on average better results with resnet.
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/VideoCoreConfig.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/VideoCoreConfig.cpp
index 229d85a..4ff8cb5 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/VideoCoreConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/VideoCoreConfig.cpp
@@ -26,6 +26,180 @@
namespace mlir::iree_compiler::detail {
+/// Repeatedly halves `startValue` until it evenly divides `dim`, returning the
+/// result (at least 1). When `startValue` is a power of two this is the largest
+/// power-of-two <= `startValue` that divides `dim`. `startValue` is not
+/// required to be a power of two.
+static int64_t largestDivisorByHalving(int64_t dim, int64_t startValue) {
+ if (dim < 0) {
+ return 1;
+ }
+ while (dim % startValue != 0) {
+ startValue >>= 1;
+ }
+ return std::max(startValue, int64_t(1));
+}
+
+/// Sets a SPIR-V "BaseVectorize" lowering config for a (batch) matmul on the
+/// Broadcom VideoCore VII GPU.
+///
+/// The op is treated as the standard contraction (Bx)MxK * (Bx)KxN = (Bx)MxN,
+/// with the B/M/N/K loop indices. Only static shapes of f32 are handled
+/// everything else will be rejected.
+///
+/// Four levels of tiling are selected:
+/// 1. Workgroup size - aim for about ~256 invocations
+/// 2. Workgroup tile - the M tile is the largest power-of-two dividing M up
+/// to
+/// ~4x the workgroup X count. The N tile takes the rest of a 1024-element
+/// budget. This biases sequential loads toward M.
+/// 3. Thread tile - each invocation handles 1/16th of the workgroup tile in
+/// M and N (at least 1).
+/// 4. Reduction tile - the K loop is vectorized by the largest power-of-two
+/// dividing K, capped at 4.
+LogicalResult
+setMatmulOpVideoCoreConfig(IREE::GPU::TargetAttr target, linalg::LinalgOp op,
+ std::array<int64_t, 2> bestWorkgroupSizeXY,
+ std::array<int64_t, 3> bestThreadTileSizeMNK) {
+ LLVM_DEBUG(llvm::dbgs() << "trying to deduce config as matmul...\n");
+ OpOperand *lhs = op.getDpsInputOperand(0);
+ OpOperand *rhs = op.getDpsInputOperand(1);
+
+ // The following tiling heuristic will ignore any operations that do not
+ // have statically known shapes.
+ auto lhsType = llvm::cast<ShapedType>(lhs->get().getType());
+ auto rhsType = llvm::cast<ShapedType>(rhs->get().getType());
+ // This routine can in principle handle 8/16/32-bit element types, but the
+ // caller (setVideoCoreMatmulConfig) currently restricts inputs to f32.
+ auto elementBits =
+ static_cast<int>(IREE::Util::getTypeBitWidth(lhsType.getElementType()));
+ if (!llvm::is_contained({8, 16, 32}, elementBits)) {
+ return failure();
+ }
+
+ ArrayRef<int64_t> lhsShape = lhsType.getShape();
+ ArrayRef<int64_t> rhsShape = rhsType.getShape();
+ if (llvm::any_of(lhsShape, ShapedType::isDynamic)) {
+ return failure();
+ }
+ if (llvm::any_of(rhsShape, ShapedType::isDynamic)) {
+ return failure();
+ }
+
+ // Ensure that we only focus on batch matmuls or single matmuls, i.e. 2D or 3D
+ assert(llvm::is_contained({2u, 3u}, op.getNumParallelLoops()));
+
+ // Find the loop indices for the B, M, N, K dimensions. We tile the standard
+ // matmul (Bx)MxK * (Bx)KxN = (Bx)MxN, where M is the LHS parallel dim, N the
+ // RHS parallel dim and K the reduction dim (see getMatmulBMNKIndex).
+ int lastParallelDim = -1;
+ const auto [bIndex, mIndex, nIndex, kIndex] =
+ getMatmulBMNKIndex(op, &lastParallelDim);
+ if (mIndex < 0 || nIndex < 0 || kIndex < 0) {
+ return failure();
+ }
+ const bool isBM = bIndex >= 0;
+
+ // Get all the dimension sizes that we need for tiling.
+ SmallVector<int64_t> loopRanges = op.getStaticLoopRanges();
+ const unsigned numLoops = loopRanges.size();
+ const int64_t dimM = loopRanges[mIndex];
+ const int64_t dimN = loopRanges[nIndex];
+ const int64_t dimK = loopRanges[kIndex];
+
+ int64_t bestX = bestWorkgroupSizeXY[0], bestY = bestWorkgroupSizeXY[1];
+ LLVM_DEBUG({
+ llvm::dbgs() << "best thread tile size (M, N, K) = ("
+ << bestThreadTileSizeMNK[0] << ", " << bestThreadTileSizeMNK[1]
+ << ", " << bestThreadTileSizeMNK[2] << ")\n";
+ llvm::dbgs() << "best workgroup size (X, Y) = (" << bestX << ", " << bestY
+ << ")\n";
+ });
+
+ // Aim for ~256 invocations per workgroup, balanced to the output shape: the Y
+ // (N) dimension uses half of dimN capped at 16, and X (M) takes the rest of
+ // the ~256 budget (capped at dimM). These larger workgroups are load-bearing
+ // for the VideoCore VII -- sizing the workgroup down to match the tile
+ // measured several times slower on the V3D.
+ SmallVector<int64_t, 3> workgroupSize(3, 1); // (X, Y, Z)
+ workgroupSize[1] = std::min(dimN >> 1, int64_t(16));
+ workgroupSize[0] =
+ std::min(std::max((bestX * bestY) / workgroupSize[1], int64_t(1)), dimM);
+
+ SmallVector<int64_t> workgroupTileSizes(numLoops, 0);
+ // Batch is simply tiled to 1 for now. Could be improved if batch is
+ // large and the spatial dimensions are small compared to the amount
+ // of available threads.
+ if (isBM) {
+ workgroupTileSizes[bIndex] = 1;
+ }
+ // To maximize the number of elements loaded sequentially we give the M
+ // dimension the larger workgroup tile and hand the rest of the budget (1024
+ // elements total) to the N dimension.
+ workgroupTileSizes[mIndex] =
+ largestDivisorByHalving(dimM, workgroupSize[0] * 4);
+ workgroupTileSizes[nIndex] =
+ largestDivisorByHalving(dimN, 1024 / workgroupTileSizes[mIndex]);
+
+ // Thread Tiling
+ SmallVector<int64_t> threadTileSizes(numLoops, 0);
+ // Batch is simply tiled to 1 for now. Remember workgroup is 1 and tile size
+ // is 1
+ if (isBM) {
+ threadTileSizes[bIndex] = workgroupTileSizes[bIndex] / workgroupSize[2];
+ }
+ // Each thread is given 1/16th of what the workgroup was given. Which should
+ // be skewed in favour of the M dimension to allow for coalesced memory
+ // accesses.
+ threadTileSizes[mIndex] =
+ std::max(workgroupTileSizes[mIndex] / 16, int64_t(1));
+ threadTileSizes[nIndex] =
+ std::max(workgroupTileSizes[nIndex] / 16, int64_t(1));
+
+ // The reduction tiling determines the width of the vector multiply-add in the
+ // inner loop. The largest power-of-two that divides K (so the K loop tiles
+ // evenly), searched downward from maxVectorization and capped at 4.
+ SmallVector<int64_t> reductionTileSizes(numLoops, 0);
+ // 32 is an empirical upper bound from the original tuning. It has no
+ // known correspondence to a specific hardware property.
+ int64_t maxVectorization = 32;
+ reductionTileSizes[kIndex] =
+ std::min(largestDivisorByHalving(dimK, maxVectorization), int64_t(4));
+
+ workgroupTileSizes.resize(lastParallelDim + 1);
+ threadTileSizes.resize(lastParallelDim + 1);
+
+ TileSizesListType tileSizes;
+ llvm::append_values(tileSizes, workgroupTileSizes, threadTileSizes,
+ reductionTileSizes);
+
+ LLVM_DEBUG({
+ llvm::dbgs() << "workgroup size (X, Y, X) = (" << workgroupSize[0] << ", "
+ << workgroupSize[1] << ", " << workgroupSize[2] << ")\n";
+ llvm::dbgs() << "workgroup tiling (M, N) = (" << workgroupTileSizes[mIndex]
+ << ", " << workgroupTileSizes[nIndex] << ")\n";
+ llvm::dbgs() << "thread tiling (M, N) = (" << threadTileSizes[mIndex]
+ << ", " << threadTileSizes[nIndex] << ")\n";
+ llvm::dbgs() << "reduction tiling (M, N, k) = ("
+ << reductionTileSizes[mIndex] << ", "
+ << reductionTileSizes[nIndex] << ','
+ << reductionTileSizes[kIndex] << ")\n";
+ });
+
+ // Sets the workgroup size on the dispatch function and adds the tiling for
+ // the MatMul to the corresponding linalg operation.
+ MLIRContext *ctx = op->getContext();
+ auto config = IREE::Codegen::LoweringConfigAttr::get(ctx, tileSizes);
+ auto pipelineAttr = IREE::GPU::SPIRVPipelineAttr::get(
+ ctx, IREE::GPU::SPIRVLoweringPipeline::BaseVectorize);
+ auto translationInfo = IREE::Codegen::TranslationInfoAttr::get(
+ ctx, pipelineAttr, SymbolRefAttr(), workgroupSize, /*subgroupSize=*/0,
+ DictionaryAttr());
+ auto funcOp = op->getParentOfType<mlir::FunctionOpInterface>();
+ return setOpConfigAndEntryPointFnTranslation(funcOp, op, config,
+ translationInfo);
+}
+
static LogicalResult setVideoCoreMatmulConfig(linalg::LinalgOp op,
IREE::GPU::TargetAttr target) {
auto inputType =
@@ -37,7 +211,7 @@
}
const std::array<int64_t, 2> workgroupXY = {16, 16};
const std::array<int64_t, 3> threadMNK = {4, 4, 4};
- return setMatmulOpConfig(target, op, workgroupXY, threadMNK);
+ return setMatmulOpVideoCoreConfig(target, op, workgroupXY, threadMNK);
}
static int64_t getWorkgroupTiling(int64_t &remainingThreads,
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_broadcom_matmul.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_broadcom_matmul.mlir
index fa28f5f..5d2568a 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_broadcom_matmul.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_broadcom_matmul.mlir
@@ -24,7 +24,7 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 64], [4, 4], [0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 16], [4, 1], [0, 0, 4]{{\]}}>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 16, 1]>
// CHECK: func.func @matmul_1024x2048x512()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -57,8 +57,8 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 8], [1, 4], [0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [2, 64, 1]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[2, 8], [1, 1], [0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [21, 12, 1]>
// CHECK: func.func @matmul_3136x24x96()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.matmul
@@ -91,7 +91,7 @@
}
// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[4, 64], [1, 4], [0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 4, 1]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 16, 1]>
// CHECK: func.func @matmul_196x64x192()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.matmul
@@ -119,8 +119,8 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[128, 32], [4, 4], [0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [8, 32, 1]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 16], [4, 1], [0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 16, 1]>
// CHECK: func.func @matmul_12544x96x16()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.matmul
@@ -152,8 +152,8 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 32], [1, 4], [0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [8, 1, 1]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 32], [1, 2], [0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 16, 1]>
// CHECK: func.func @matmul_49x160x576()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.matmul
@@ -191,8 +191,8 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[2, 128], [1, 4], [0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [32, 2, 1]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[2, 512], [1, 32], [0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [2, 16, 1]>
// CHECK: func.func @matmul_2x1024x576()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.matmul
@@ -222,7 +222,7 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 64, 64], [1, 4, 4], [0, 0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 64, 16], [1, 4, 1], [0, 0, 0, 4]{{\]}}>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 16, 1]>
// CHECK: func.func @batch_matmul_4x384x384()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -256,8 +256,8 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 2, 8], [1, 1, 4], [0, 0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [2, 2, 1]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 2, 8], [1, 1, 1], [0, 0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [2, 4, 1]>
// CHECK: func.func @batch_matmul_4x2x8()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.batch_matmul
@@ -298,8 +298,8 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 8, 64], [1, 1, 4], [0, 0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 8, 1]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 8, 128], [1, 1, 8], [0, 0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [8, 16, 1]>
// CHECK: func.func @generic_batch_matmul_32x2x512()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
@@ -351,8 +351,8 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 4, 64], [1, 1, 4], [0, 0, 0, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 4, 1]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 4, 256], [1, 1, 16], [0, 0, 0, 4]{{\]}}>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_gpu.spirv_pipeline<BaseVectorize> workgroup_size = [16, 16, 1]>
// CHECK: func.func @generic_batch_matmul_8x2500x512x4608()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
