[Codegen] Change CompilationInfoAttr to take a lowering config interface (#17752)
This allows specifying other lowering configs from the input. This is a
breaking change because now the lowering_config on
iree_codegen.compilation_info must be fully specified.
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp
index 118a781..f803103 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp
@@ -338,15 +338,17 @@
LogicalResult
CompilationInfoAttr::verify(function_ref<InFlightDiagnostic()> emitError,
- LoweringConfigAttr loweringConfig,
+ LoweringConfigAttrInterface loweringConfig,
TranslationInfoAttr translationInfo) {
if (!loweringConfig) {
return emitError() << "missing lowering config";
}
- if (failed(LoweringConfigAttr::verify(
- emitError, loweringConfig.getTilingLevels(),
- loweringConfig.getNativeVectorSize()))) {
- return failure();
+ if (auto defaultConfig = llvm::dyn_cast<LoweringConfigAttr>(loweringConfig)) {
+ if (failed(LoweringConfigAttr::verify(
+ emitError, defaultConfig.getTilingLevels(),
+ defaultConfig.getNativeVectorSize()))) {
+ return emitError() << "invalid lowering config: " << defaultConfig;
+ }
}
if (!translationInfo) {
return emitError() << "missing translation info";
@@ -356,7 +358,7 @@
translationInfo.getCodegenSpec(), translationInfo.getWorkgroupSize(),
translationInfo.getSubgroupSize(),
translationInfo.getConfiguration()))) {
- return failure();
+ return emitError() << "invalid translation info: " << translationInfo;
}
return success();
}
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td
index dad9349..cef0735 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td
@@ -286,21 +286,19 @@
Specifies the information that allows controlling the compilation
of operations like `linalg.matmul`/`linalg.*conv` within
IREE. This information is used to override the defaults used by
- the IREE compiler. Currently it is only valid to set this on
- `linalg.matmul`/`linalg.*conv*` operations.
+ the IREE compiler. If set on the input to the compiler, there is no
+ guarantee that the config survives until codegen. Named operations like
+ `linalg.matmul`/`linalg.*conv*` are more likely to retain their lowering
+ configurations.
TODO: It is expected that the `TranslationInfoAttr` and the
`LoweringConfigAttr` are specified. Currently there is no
verification that the values of the `LoweringConfigAttr` fully
specifies the behaviour of the compilation path chosen with
- `TranslationInfoAttr`. This could be added in the future. Note:
- Typically the values used for the first-level tiling in
- `LoweringConfigAttr` value in the `TranslationInfoAttr` are the
- same since the first-level of tile + distribute is already done
- at the `Flow` level. This verification is also a TODO.
+ `TranslationInfoAttr`. This could be added in the future.
}];
let parameters = (ins
- AttrParameter<"LoweringConfigAttr", "">:$loweringConfig,
+ AttrParameter<"LoweringConfigAttrInterface", "">:$loweringConfig,
AttrParameter<"TranslationInfoAttr", "">:$translationInfo
);
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/lowering_config_attr.mlir b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/lowering_config_attr.mlir
index 8979a72..a19c953 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/lowering_config_attr.mlir
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/lowering_config_attr.mlir
@@ -33,7 +33,7 @@
module {
func.func @test() attributes {
compilation_info = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = []>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = []>,
translation_info = <CPUDefault>>} {
return
}
@@ -48,7 +48,7 @@
module {
func.func @test() attributes {
compilation_info = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = []>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = []>,
translation_info = <CPUDefault workgroup_size = [16, 4, 1] subgroup_size = 32>>} {
return
}
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/lowering_matmul_promotion.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/lowering_matmul_promotion.mlir
index 2bb5062..27831cf 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/lowering_matmul_promotion.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/lowering_matmul_promotion.mlir
@@ -7,7 +7,7 @@
// Verify pipelining + multi-buffering.
#compilation = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[64, 64, 16]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64, 16]]>,
translation_info = <SPIRVMatmulPromoteVectorize workgroup_size = [16, 8, 1], {pipeline_depth = 2, store_stage = 1}>>
#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
@@ -204,7 +204,7 @@
// Check that fused transposed consumer elementwise op does not cause extra workgroup memory allocations.
#compilation = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[64, 256, 32]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 256, 32]]>,
translation_info = <SPIRVMatmulPromoteVectorize workgroup_size = [32, 8, 1], {pipeline_depth = 1, store_stage = 1}>>
#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_cooperative_ops.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_cooperative_ops.mlir
index deaaef5..8102908 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_cooperative_ops.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_cooperative_ops.mlir
@@ -865,7 +865,7 @@
]>
#compilation = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[1, 64, 64], [1, 16, 64], [0, 0, 0, 16], [1, 16, 16, 16]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 64, 64], [1, 16, 64], [0, 0, 0, 16], [1, 16, 16, 16]]>,
translation_info = <SPIRVCooperativeMatrixVectorize workgroup_size = [32, 4, 1] subgroup_size = 32, {pipeline_depth = 1, store_stage = 1}>>
hal.executable public @batch_matmul_f16_16x4096x4096x64_truncf_mulf {
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_promotion.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_promotion.mlir
index bdbad94..854fc57 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_promotion.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_promotion.mlir
@@ -173,7 +173,7 @@
]>
#user_config = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[16, 128, 16]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[16, 128, 16]]>,
translation_info = <SPIRVMatmulPromoteVectorize workgroup_size = [16, 8, 1], {pipeline_depth = 0, store_stage = 1}>>
hal.executable @matmul_f16_32x1280x1280 {
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/split_reduction.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/split_reduction.mlir
index 6b91b5e..b9c2742 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/split_reduction.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/split_reduction.mlir
@@ -1,7 +1,7 @@
// RUN: iree-opt --pass-pipeline='builtin.module(util.func(iree-flow-split-reduction-ops))' --iree-flow-split-matmul-reduction=4 %s | FileCheck %s
#compilation = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[64, 64, 0]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64, 0]]>,
translation_info = <CPUDefault>>
util.func public @matmul(%arg0: tensor<100x200xf32>, %arg1: tensor<200x300xf32>, %arg2: tensor<100x300xf32>) -> tensor<100x300xf32> {
%0 = linalg.matmul {compilation_info = #compilation}
diff --git a/tests/e2e/matmul/generate_e2e_matmul_tests.py b/tests/e2e/matmul/generate_e2e_matmul_tests.py
index 596f610..15f0e33 100644
--- a/tests/e2e/matmul/generate_e2e_matmul_tests.py
+++ b/tests/e2e/matmul/generate_e2e_matmul_tests.py
@@ -569,7 +569,7 @@
compilation_info_string = (
f"#compilation{generate_function.compilation_index} = "
"#iree_codegen.compilation_info<\n"
- f" lowering_config = <tile_sizes = {compilation_info.tile_sizes}>,\n"
+ f" lowering_config = #iree_codegen.lowering_config<tile_sizes = {compilation_info.tile_sizes}>,\n"
f" translation_info = <{compiler_pipeline} {compilation_info.workgroup_size_str()}\n"
f" {subgroup_size_str},\n"
f" {{ pipeline_depth = {compilation_info.software_pipeline_depth}, "
diff --git a/tests/e2e/regression/lowering_config.mlir b/tests/e2e/regression/lowering_config.mlir
index c2bfa8a..bbf6301 100644
--- a/tests/e2e/regression/lowering_config.mlir
+++ b/tests/e2e/regression/lowering_config.mlir
@@ -1,11 +1,11 @@
#compilation0 = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[32, 32], [1, 8, 0], [0, 0, 8], [0, 0, 0]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[32, 32], [1, 8, 0], [0, 0, 8], [0, 0, 0]]>,
translation_info = <CPUDoubleTilingExpert>>
#compilation1 = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[64, 64], [1, 4, 0], [0, 0, 4], [0, 0, 0]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64], [1, 4, 0], [0, 0, 4], [0, 0, 0]]>,
translation_info = <CPUDoubleTilingExpert>>
#compilation2 = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [{sizes=[32, 64], interchange=[1, 0]}, [1, 1, 0], [0, 0, 8], [0, 0, 0]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [{sizes=[32, 64], interchange=[1, 0]}, [1, 1, 0], [0, 0, 8], [0, 0, 0]]>,
translation_info = <CPUDoubleTilingExpert>>
func.func @lowering_config_test() {
@@ -24,11 +24,11 @@
// Conv dims: N, OH, OW, OC, KH, KW, (IC)
// Remove H
#conv_compilation0 = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[0, 7, 7, 64, 0, 0, 0], [1, 1, 7, 4, 0, 0, 0], [0, 0, 0, 0, 1, 3, 4], [0, 0, 0, 0, 0, 0, 0]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[0, 7, 7, 64, 0, 0, 0], [1, 1, 7, 4, 0, 0, 0], [0, 0, 0, 0, 1, 3, 4], [0, 0, 0, 0, 0, 0, 0]]>,
translation_info = <CPUConvTileAndDecomposeExpert>>
// Remove W
#conv_compilation1 = #iree_codegen.compilation_info<
- lowering_config = <tile_sizes = [[0, 7, 7, 64, 0, 0, 0], [1, 7, 1, 4, 0, 0, 0], [0, 0, 0, 0, 3, 1, 4], [0, 0, 0, 0, 0, 0, 0]]>,
+ lowering_config = #iree_codegen.lowering_config<tile_sizes = [[0, 7, 7, 64, 0, 0, 0], [1, 7, 1, 4, 0, 0, 0], [0, 0, 0, 0, 3, 1, 4], [0, 0, 0, 0, 0, 0, 0]]>,
translation_info = <CPUConvTileAndDecomposeExpert>>
func.func @conv() {
%input = util.unfoldable_constant dense<1.0> : tensor<36x7x7x512xf32>
