Add new tile size for fp16 and fallback if vectorization is not possible (#3882)
Add new tile size for fp16 and fallback to none vector version for size not aligned on tile size.
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
index 0600b24..7f0198a 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
@@ -61,6 +61,14 @@
return std::make_tuple(nprocs_x, nprocs / nprocs_x);
}
+namespace {
+struct LaunchConfigInfo {
+ std::array<int64_t, 3> workgroupSize = {1, 1, 1};
+ std::array<int64_t, 3> numSubgroups = {1, 1, 1};
+ bool vectorize = false;
+};
+} // namespace
+
/// For a given operation `op`, compute the following configurations according
/// to SPIR-V `targetEnv` and `options`:
/// 1) number of tiling levels and tile sizes to use (updates `tileSizes`),
@@ -71,8 +79,7 @@
static LogicalResult getOpLaunchConfig(T op, const spirv::TargetEnv &targetEnv,
const SPIRVCodegenOptions &options,
TileSizesListType &tileSizes,
- std::array<int64_t, 3> &workgroupSize,
- std::array<int64_t, 3> &numSubgroups) {
+ LaunchConfigInfo &config) {
return op.emitError("undefined launch config for tiled operation");
}
@@ -82,14 +89,13 @@
const spirv::TargetEnv &targetEnv,
const SPIRVCodegenOptions &options,
TileSizesListType &tileSizes,
- std::array<int64_t, 3> &workgroupSize,
- std::array<int64_t, 3> &numSubgroups) {
+ LaunchConfigInfo &config) {
unsigned maxWorkgroupSize = targetEnv.getResourceLimits()
.max_compute_workgroup_invocations()
.getInt();
- std::tie(workgroupSize[0], workgroupSize[1]) =
+ std::tie(config.workgroupSize[0], config.workgroupSize[1]) =
distributeProcs2D(maxWorkgroupSize);
- workgroupSize[2] = 1;
+ config.workgroupSize[2] = 1;
// This is just being hard-wired for now to be minimal viable, but this can be
// decided better when we have better estimates of device charecteristics.
const int64_t nRowsPerWorkitem = 1;
@@ -102,9 +108,9 @@
tileSizeK = 32;
}
assert(tileSizes.empty());
- SmallVector<int64_t, 4> ts = {nBatchesPerWorkitem,
- nRowsPerWorkitem * workgroupSize[1],
- nColsPerWorkitem * workgroupSize[0], tileSizeK};
+ SmallVector<int64_t, 4> ts = {
+ nBatchesPerWorkitem, nRowsPerWorkitem * config.workgroupSize[1],
+ nColsPerWorkitem * config.workgroupSize[0], tileSizeK};
tileSizes.emplace_back(std::move(ts));
return success();
}
@@ -209,16 +215,36 @@
std::array<int64_t, 3> &numSubgroups) {
if (targetEnv.getVendorID() != spirv::Vendor::ARM) return failure();
+ auto lhsType = op.inputs()[0].getType().cast<MemRefType>();
+ auto rhsType = op.inputs()[1].getType().cast<MemRefType>();
+ assert(lhsType.getElementType() == rhsType.getElementType());
+ // Pick ideal tile size based on the type.
+ SmallVector<int64_t, 4> workgroupLevelTs;
+ if (lhsType.getElementType().isF16()) {
+ workgroupLevelTs.append({16, 64, 8});
+ } else {
+ workgroupLevelTs.append({8, 64, 4});
+ }
+
+ // Fall back to the none vectorize path for cases we don't handle.
+ if (!lhsType.hasStaticShape() || !rhsType.hasStaticShape() ||
+ lhsType.getDimSize(0) % workgroupLevelTs[0] != 0 ||
+ rhsType.getDimSize(0) % workgroupLevelTs[1] != 0 ||
+ lhsType.getDimSize(1) % workgroupLevelTs[2] != 0) {
+ return failure();
+ }
+
workgroupSize[0] = targetEnv.getResourceLimits().subgroup_size().getInt();
workgroupSize[1] = 1;
workgroupSize[2] = 1;
- SmallVector<int64_t, 4> ts = {8, 64, 4};
- tileSizes.emplace_back(ts);
+ tileSizes.emplace_back(workgroupLevelTs);
// No tiling at the subgroup level since this target doesn't use subgroup op
// or shared memory.
tileSizes.emplace_back();
- SmallVector<int64_t, 4> threadTs = {ts[0], ts[1] / workgroupSize[0], ts[2]};
- tileSizes.emplace_back(threadTs);
+ SmallVector<int64_t, 4> invocationLevelTs = {
+ workgroupLevelTs[0], workgroupLevelTs[1] / workgroupSize[0],
+ workgroupLevelTs[2]};
+ tileSizes.emplace_back(invocationLevelTs);
return success();
}
@@ -227,25 +253,27 @@
const spirv::TargetEnv &targetEnv,
const SPIRVCodegenOptions &options,
TileSizesListType &tileSizes,
- std::array<int64_t, 3> &workgroupSize,
- std::array<int64_t, 3> &numSubgroups) {
+ LaunchConfigInfo &config) {
if (options.enableVectorization &&
succeeded(getConfigForCooperativeMatmul(op, targetEnv, options, tileSizes,
- workgroupSize, numSubgroups))) {
+ config.workgroupSize,
+ config.numSubgroups))) {
+ config.vectorize = true;
return success();
} else if (options.enableVectorization &&
succeeded(getTargetSpecificConfig(op, targetEnv, options,
- tileSizes, workgroupSize,
- numSubgroups))) {
+ tileSizes, config.workgroupSize,
+ config.numSubgroups))) {
+ config.vectorize = true;
return success();
}
unsigned maxWorkgroupSize = targetEnv.getResourceLimits()
.max_compute_workgroup_invocations()
.getInt();
- std::tie(workgroupSize[0], workgroupSize[1]) =
+ std::tie(config.workgroupSize[0], config.workgroupSize[1]) =
distributeProcs2D(maxWorkgroupSize);
- workgroupSize[2] = 1;
+ config.workgroupSize[2] = 1;
const int nRowsPerWorkitem = 1;
const int nColsPerWorkitem = 1;
int64_t tileSizeK = 0;
@@ -255,8 +283,9 @@
tileSizeK = 32;
}
assert(tileSizes.empty());
- SmallVector<int64_t, 4> ts = {nRowsPerWorkitem * workgroupSize[1],
- nColsPerWorkitem * workgroupSize[0], tileSizeK};
+ SmallVector<int64_t, 4> ts = {nRowsPerWorkitem * config.workgroupSize[1],
+ nColsPerWorkitem * config.workgroupSize[0],
+ tileSizeK};
tileSizes.emplace_back(std::move(ts));
return success();
}
@@ -266,8 +295,7 @@
const spirv::TargetEnv &targetEnv,
const SPIRVCodegenOptions &options,
TileSizesListType &tileSizes,
- std::array<int64_t, 3> &workgroupSize,
- std::array<int64_t, 3> &numSubgroups) {
+ LaunchConfigInfo &config) {
unsigned maxWorkgroupSize = targetEnv.getResourceLimits()
.max_compute_workgroup_invocations()
.getInt();
@@ -275,7 +303,7 @@
int64_t tileSizeY = maxWorkgroupSize / tileSizeX;
SmallVector<int64_t, 4> ts = {1, tileSizeY, tileSizeX};
tileSizes.emplace_back(std::move(ts));
- workgroupSize = {tileSizeX, tileSizeY, 1};
+ config.workgroupSize = {tileSizeX, tileSizeY, 1};
return success();
}
@@ -283,8 +311,7 @@
static LogicalResult getPoolingOpLaunchConfig(
PoolingOpTy op, const spirv::TargetEnv &targetEnv,
const SPIRVCodegenOptions &options, TileSizesListType &tileSizes,
- std::array<int64_t, 3> &workgroupSize,
- std::array<int64_t, 3> &numSubgroups) {
+ LaunchConfigInfo &config) {
unsigned maxWorkgroupSize = targetEnv.getResourceLimits()
.max_compute_workgroup_invocations()
.getInt();
@@ -299,7 +326,7 @@
ts[ts.size() - 2] = tileSizeY;
ts[ts.size() - 1] = tileSizeX;
tileSizes.emplace_back(std::move(ts));
- workgroupSize = {tileSizeX, tileSizeY, 1};
+ config.workgroupSize = {tileSizeX, tileSizeY, 1};
return success();
}
@@ -308,10 +335,9 @@
LogicalResult getOpLaunchConfig( \
opName op, const spirv::TargetEnv &targetEnv, \
const SPIRVCodegenOptions &options, TileSizesListType &tileSizes, \
- std::array<int64_t, 3> &workgroupSize, \
- std::array<int64_t, 3> &numSubgroups) { \
+ LaunchConfigInfo &config) { \
return getPoolingOpLaunchConfig(op, targetEnv, options, tileSizes, \
- workgroupSize, numSubgroups); \
+ config); \
}
DEFINE_POOLING_OP_CONFIG(linalg::PoolingMaxOp)
@@ -354,7 +380,7 @@
spirv::TargetEnv targetEnv(spirv::lookupTargetEnv(*linalgOps.begin()));
Optional<linalg::LinalgOp> rootOperation = {};
-
+ LaunchConfigInfo config;
for (linalg::LinalgOp linalgOp : linalgOps) {
#define DISPATCH(opName) \
if (auto op = dyn_cast<opName>(linalgOp.getOperation())) { \
@@ -365,7 +391,7 @@
rootOperation = linalgOp; \
TileSizesListType &tileSizesInfo = tileSizes[setKey(*rootOperation)]; \
if (failed(getOpLaunchConfig(op, targetEnv, options, tileSizesInfo, \
- workgroupSize, numSubgroups))) { \
+ config))) { \
return failure(); \
} \
continue; \
@@ -380,7 +406,9 @@
#undef DISPATCH
}
-
+ workgroupSize = config.workgroupSize;
+ numSubgroups = config.numSubgroups;
+ vectorize = config.vectorize;
if (!rootOperation) {
// No root operations found. Dont need to do anything.
return success();
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h
index 17f3bc5..fd3c0a2 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h
+++ b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h
@@ -110,6 +110,9 @@
return !getTileSizes(op, level).empty();
}
+ /// Use vectorize transformations.
+ bool useVectorize() const { return vectorize; }
+
protected:
/// Current tile size configuration per operation. They key used here to
/// retrieve the tile size information per operation is the value of a StrAttr
@@ -125,6 +128,9 @@
/// Number of subgroups that are logically distributed along x, y & z.
std::array<int64_t, 3> numSubgroups;
+ /// Use vectorization.
+ bool vectorize = false;
+
private:
/// Retrieves the key to use to get the `tileSizes` for a given
/// `operation`. Returns llvm::None on failure.
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
index 39b6bbf..34c6632 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
@@ -596,7 +596,7 @@
});
}
- if (options.enableVectorization) {
+ if (launchConfig.useVectorize()) {
{
OwningRewritePatternList secondLevelTilingPatterns;
populateTilingToSubgroupPatterns(context, launchConfig,
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_fused_vectorization.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_fused_vectorization.mlir
index 2bcaf93..33e1131 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_fused_vectorization.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_fused_vectorization.mlir
@@ -50,3 +50,56 @@
// CHECK: scf.yield
// CHECK-COUNT-8: vector.transfer_write %[[FOR_RES]]
// CHECK: return
+
+// -----
+
+module attributes {
+ spv.target_env =
+ #spv.target_env<#spv.vce<v1.3,
+ [Shader, Float64, Float16, Int64, Int16, Int8, StorageBuffer16BitAccess,
+ StorageUniform16, StoragePushConstant16, StorageBuffer8BitAccess,
+ UniformAndStorageBuffer8BitAccess, StoragePushConstant8, GroupNonUniform,
+ GroupNonUniformVote, GroupNonUniformArithmetic, GroupNonUniformBallot,
+ GroupNonUniformShuffle, GroupNonUniformShuffleRelative, VariablePointers,
+ VariablePointersStorageBuffer],
+ [SPV_KHR_16bit_storage, SPV_KHR_8bit_storage,
+ SPV_KHR_storage_buffer_storage_class, SPV_KHR_variable_pointers]>,
+ ARM: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}>} {
+ func @matmul_static_shape_f16()
+ attributes {vkspv.num_workgroups_fn = @matmul_static_shape__num_workgroups__} {
+ %arg0 = iree.placeholder for "interface buffer"
+ {binding = @legacy_io::@arg0, operand_result_num = 0 : i32} : memref<4096x4096xf16>
+ %arg1 = iree.placeholder for "interface buffer"
+ {binding = @legacy_io::@arg1, operand_result_num = 1 : i32} : memref<4096x4096xf16>
+ %ret0 = iree.placeholder for "interface buffer"
+ {binding = @legacy_io::@ret0, operand_result_num = 2 : i32} : memref<4096x4096xf16>
+ %cst = constant 0.000000e+00 : f16
+ linalg.fill(%ret0, %cst) : memref<4096x4096xf16>, f16
+ linalg.matmul ins(%arg0, %arg1 : memref<4096x4096xf16>, memref<4096x4096xf16>)
+ outs(%ret0 : memref<4096x4096xf16>)
+ return
+ }
+ func @matmul_static_shape__num_workgroups__
+ (!shapex.ranked_shape<[4096, 4096]>, !shapex.ranked_shape<[4096, 4096]>,
+ !shapex.ranked_shape<[4096, 4096]>) -> (index, index, index)
+ attributes {sym_visibility = "private"}
+ hal.interface @legacy_io attributes {sym_visibility = "private"} {
+ hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+ hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+ hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+ }
+}
+
+// CHECK-LABEL: func @matmul_static_shape_f16
+// CHECK-COUNT-16: vector.transfer_write
+// CHECK-COUNT-16: vector.transfer_read
+// CHECK: %[[FOR_RES:.+]]:16 = scf.for
+// CHECK-COUNT-40: vector.transfer_read
+// CHECK-COUNT-64: vector.contract
+// CHECK: scf.yield
+// CHECK-COUNT-16: vector.transfer_write %[[FOR_RES]]
+// CHECK: return
