[spirv] NFC: clean up setDefaultOpConfig a bit (#8981)
The check for untiled dimensions is not needed now.
Also remove a duplicated test.
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
index 4ce42b2..f4c40c4 100644
--- a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
+++ b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -23,6 +23,7 @@
#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#define DEBUG_TYPE "iree-spirv-kernel-config"
@@ -431,9 +432,9 @@
// configuration for the corresponding GPU workgroup dimension.
int64_t wgDim = 0;
for (auto shapeDim : llvm::reverse(partitionedLoops)) {
- // Skip untiled or dynamic dimensions.
- // TODO: Skip size-1 dimensions in Flow level tiling and distribution.
- if (loopBounds.getValue()[shapeDim] <= 0) continue;
+ int64_t loopBound = loopBounds.getValue()[shapeDim];
+ // Skip dynamic dimensions.
+ if (ShapedType::isDynamic(loopBound)) continue;
// Try to find some power of two that can devide the current shape dim
// size. This vector keeps the candidate tile sizes.
@@ -449,24 +450,23 @@
candidates.push_back(i);
}
LLVM_DEBUG({
- llvm::dbgs() << "Candidates tile sizes: [";
+ llvm::dbgs() << "Candidate tile sizes: [";
llvm::interleaveComma(candidates, llvm::dbgs());
llvm::dbgs() << "]\n";
});
for (int64_t candidate : candidates) {
- if (loopBounds.getValue()[shapeDim] % candidate != 0) {
+ if (loopBound % candidate != 0) {
if (!lossFactor) continue;
// Skip this candidate if it causes many threads to be idle.
- int64_t idleThreads =
- candidate - (loopBounds.getValue()[shapeDim] % candidate);
+ int64_t idleThreads = candidate - (loopBound % candidate);
if (idleThreads > candidate / *lossFactor) continue;
}
- LLVM_DEBUG(llvm::dbgs() << "Chosen Candiate " << candidate << "\n");
// Found a suitable candidate. Try to let each thread handle 4
// elements if this is the workgroup x dimension.
workgroupTileSizes[shapeDim] = candidate;
+ LLVM_DEBUG(llvm::dbgs() << "Chosen tile size: " << candidate << "\n");
if (vectorizable && wgDim == 0 && !lossFactor && candidate % 4 == 0) {
threadTileSizes[shapeDim] = 4;
workgroupSize[wgDim] = candidate / 4;
diff --git a/iree/compiler/Codegen/SPIRV/test/config_default_linalg_ext_ops.mlir b/iree/compiler/Codegen/SPIRV/test/config_default_linalg_ext_ops.mlir
index 7dc3add..e8727a5 100644
--- a/iree/compiler/Codegen/SPIRV/test/config_default_linalg_ext_ops.mlir
+++ b/iree/compiler/Codegen/SPIRV/test/config_default_linalg_ext_ops.mlir
@@ -184,47 +184,3 @@
// CHECK: func @static_3d_fft_stage3()
// CHECK: iree_linalg_ext.fft
// CHECK-SAME: lowering_config = #[[CONFIG]]
-
-// -----
-
-#executable_layout = #hal.executable.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.executable private @copy_op {
- hal.executable.variant @vulkan_spirv_fb, target = <"vulkan", "vulkan-spirvfb", {
- spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Unknown:IntegratedGPU, {
- max_compute_shared_memory_size = 32768 : i32,
- max_compute_workgroup_invocations = 512 : i32,
- max_compute_workgroup_size = dense<512> : vector<3xi32>,
- subgroup_size = 16 : i32}>
- }> {
- hal.executable.entry_point @copy_op layout(#executable_layout)
- builtin.module {
- func.func @copy_op() {
- %offset_y = hal.interface.constant.load[0] : index
- %offset_x = hal.interface.constant.load[1] : index
- %source_size_y = hal.interface.constant.load[2] : index
- %source_size_x = hal.interface.constant.load[3] : index
- %dest_size_y = hal.interface.constant.load[4] : index
- %dest_size_x = hal.interface.constant.load[5] : index
- %source = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : memref<?x?xf32>{%source_size_y, %source_size_x}
- %dest = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) : memref<?x?xf32>{%dest_size_y, %dest_size_x}
- linalg.generic {
- indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
- iterator_types = ["parallel", "parallel"]}
- ins(%source : memref<?x?xf32>) outs(%dest : memref<?x?xf32>) {
- ^bb0(%b0 : f32, %b1 : f32):
- linalg.yield %b0 : f32
- }
- return
- }
- }
- }
-}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 16], [1, 1]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<SPIRVDistribute>
-// CHECK: linalg.generic
-// CHECK-SAME: lowering_config = #[[CONFIG]]