[Codegen] Change CompilationInfoAttr to take a lowering config interf…

…ace (iree-org#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.

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp

@@ -338,15 +338,17 @@ LoweringConfigAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 
 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 @@ CompilationInfoAttr::verify(function_ref<InFlightDiagnostic()> emitError,
           translationInfo.getCodegenSpec(), translationInfo.getWorkgroupSize(),
           translationInfo.getSubgroupSize(),
           translationInfo.getConfiguration()))) {
-    return failure();
+    return emitError() << "invalid translation info: " << translationInfo;
   }
   return success();
 }

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td

@@ -286,21 +286,19 @@ def IREECodegen_CompilationInfoAttr :
     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
   );
 

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/lowering_config_attr.mlir

@@ -33,7 +33,7 @@ module {
 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 {
 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
   }

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 @@ hal.executable @matmul_f32_128x256x64 {
 // 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 = [

compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_cooperative_ops.mlir

@@ -865,7 +865,7 @@ hal.executable public @generic_batch_matmul_32x128x512x64 {
 ]>
 
 #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 {

compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_matmul_promotion.mlir

@@ -173,7 +173,7 @@ hal.executable @matmul_f16_128x256x64 {
 ]>
 
 #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 {

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}

tests/e2e/matmul/generate_e2e_matmul_tests.py

@@ -569,7 +569,7 @@ def generate_function(
         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}, "

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 @@ func.func @lowering_config_test() {
 // 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>