[GPU][Codegen] Allowing mfma for narrow problem config sizes (#19615)
The motivation of this PR is convolution performance for resnet50
configs. With this PR (and a few pending ones), conv performance with
igemm pipeline get decent speedup in situation where a standalone
dimension size is smaller than intrinsic size. (Take dispatch 69 as
example, the select tile m:7, n:512, k:4608 will be rejected from mfma
because m tile is smaller than intrinsic size of 16). This happens
because previously we are too defensive about when to use intrinsic: in
situation when alignment is not required, we still enforce mfma to be
picked up only when m/n/k tiles are all larger than intrinsic size.
With @nirvedhmeshram's https://github.com/iree-org/iree/pull/19271 and
https://github.com/iree-org/iree/pull/19484, padding is allowed in tile
and fuse matmul and igemm tile and fuse pipelines, it is no longer
necessary to be as conservative as before. I am therefore getting rid of
the conditional check that blocks mfma from being picked up.
This will impact a few pipelines that use `canTargetIntrinsic()`:
- `LLVMGPUPadAndVectorDistribute` will allow narrow m/n/k dimension
sizes for batch matmul
- In `iree-codegen-rocdl-configuration-pipeline`, will allow narrow
m/n/k dimension sizes for matmul (instead of warp reduction)
---------
Signed-off-by: jerryyin <zhuoryin@amd.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp
index 8cab5c6..f8e30f3 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp
@@ -246,14 +246,20 @@
return failure(); // Cannot use this intrinsic for misaligned cases.
}
- // Cannot use the intrinsic when the tile size is greater than problem size.
- // Because tiling is a no-op, and we can't infer tiling sizes from IR.
- if (!mustBeAligned && (problem.mSizes.back() < intrinsic.mSizes[0] ||
- problem.nSizes.back() < intrinsic.nSizes[0] ||
- problem.kSizes.back() < intrinsic.kSizes[0])) {
+ // TODO: Figure out what the precise cutoff is, this may be machine dependent.
+ // In situation when alignment isn't required, we disallow intrinsics to be
+ // picked if the tile size is too small. For example, this will force a matmul
+ // with a tiny dimension to not use MFMA instructions because of the
+ // additional overhead that comes with it. However, 4 is only an approximation
+ // to boundary between matvec and matmul. The actual heuristic can be
+ // established after we sweep the different tile sizes for a problem config.
+ // Once a precise threshold is established, replace 4 with the threshold and
+ // remove this todo.
+ if (!mustBeAligned &&
+ (problem.mSizes.back() < 4 || problem.nSizes.back() < 4 ||
+ problem.kSizes.back() < 4)) {
return failure();
}
-
return success();
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_gfx942.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_gfx942.mlir
index d71e7ed..373b67b 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_gfx942.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_gfx942.mlir
@@ -278,8 +278,8 @@
flow.dispatch.tensor.store %7, %2, offsets = [0, 0, 0], sizes = [64, 968, 4], strides = [1, 1, 1] : tensor<64x968x4xf16> -> !flow.dispatch.tensor<writeonly:tensor<64x968x4xf16>>
return
}
-// Check that we don't support LLVMGPUPadAndVectorDistribute for narrow N/M atm.
-// CHECK-NOT: #iree_codegen.translation_info<pipeline = LLVMGPUPadAndVectorDistribute
+// Check that we support LLVMGPUPadAndVectorDistribute for narrow N/M atm.
+// CHECK: #iree_codegen.translation_info<pipeline = LLVMGPUPadAndVectorDistribute
// CHECK-LABEL: func.func @narrow_n_batch_matmul_64x968x4x320_f16()
// -----
