[Codegen][Tuner] improve verifier for the default attribute (#20173)

This PR enhances the verification of the default attribute
`iree_codegen.tuning_spec_with_default_entrypoint`.

- The `__kernel_config `named sequence operation must include the
`iree_codegen.tuning_spec_entrypoint` attribute, on which the
verification will be performed.

- The `__kernel_config` named sequence operation must contain only a
single `foreach_match` operation.

---------

Signed-off-by: Bangtian Liu <liubangtian@gmail.com>

Author: bangtianliu

compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir

@@ -35,8 +35,10 @@
 // Check that both the user tuning spec and the default spec get linked and
 // materialized. The user spec should have precedence over the default one.
 
+// TODO: Re-add the check for iree_codegen.tuning_spec_with_default_entrypoint
+// once new linking is added and the output IR can pass verification for the default attribute.
+
 // BOTH-LABEL: module @iree_linked_tuning_spec
-// BOTH-SAME:    iree_codegen.tuning_spec_with_default_entrypoint
 // BOTH-SAME:    transform.with_named_sequence
 // BOTH-LABEL:   module @mmt_tile_and_fuse_spec_0 attributes {transform.with_named_sequence}
 // BOTH-LABEL:     transform.named_sequence @main

compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp

@@ -83,7 +83,9 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef<NamedSequenceOp> specsToLink) {
       0, hasConsumedSequences ? kArgConsumedAttrName : kArgReadOnlyAttrName,
       builder.getUnitAttr());
   newSpec->setAttr(kTuningSpecEntrypointAttrName, builder.getUnitAttr());
-  module->setAttr(kTuningSpecDefaultEntrypointAttrName, builder.getUnitAttr());
+  // TODO: Re-enable default attribute as below once new linking lands.
+  // module->setAttr(kTuningSpecDefaultEntrypointAttrName,
+  // builder.getUnitAttr());
 
   Region &region = newSpec.getRegion();
   Block *body = builder.createBlock(&region, region.begin(),

compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs.mlir

@@ -10,8 +10,10 @@
 
 // Check that the final tuning spec is as expected when the user tuning spec is provided.
 
+// TODO: Add the check for default attribute (`iree_codegen.tuning_spec_with_default_entrypoint`) here
+//       once the merging logic supports cases beyond a single `foreach_match` operation.
+
 // CHECK-LABEL: module @iree_linked_tuning_spec
-// CHECK-SAME:    iree_codegen.tuning_spec_with_default_entrypoint
 // CHECK-SAME:    transform.with_named_sequence
 // CHECK-LABEL:   module @user_spec_0 attributes {transform.with_named_sequence}
 // CHECK-LABEL:     transform.named_sequence @hello
@@ -31,6 +33,7 @@
 // SKIPLINK-LABEL: module  @user_spec
 // SKIPLINK-SAME:    iree_codegen.tuning_spec_with_default_entrypoint
 // SKIPLINK-SAME:    transform.with_named_sequence
+// SKIPLINK:         transform.print  {name = "Hello Tuning Spec"}
 // SKIPLINK-NOT:    module @{{.+}}
 // SKIPLINK:        module attributes
 // SKIPLINK-SAME:     iree_codegen.tuning_spec_mlirbc = dense<{{.+}}> : vector<{{[0-9]+}}xi8>

compiler/src/iree/compiler/Codegen/Common/test/tuning_spec_default.mlir

@@ -1,9 +1,20 @@
 // RUN: iree-opt %s
 
 module @user_spec attributes { transform.with_named_sequence, iree_codegen.tuning_spec_with_default_entrypoint } {
-  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op
-    attributes { iree_codegen.tuning_spec_entrypoint } {
-    transform.print {name = "Hello Tuning Spec", skip_regions}
-    transform.yield %arg0 : !transform.any_op
+  transform.named_sequence @match(%arg: !transform.any_op {transform.readonly}) -> (!transform.any_op) {
+    transform.print {name = "Hello Tuning Spec"}
+    transform.yield %arg : !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly}) {
+    transform.yield
+  }
+
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+      %res = transform.foreach_match in %arg0
+        @match -> @apply_op_config
+        : (!transform.any_op) -> (!transform.any_op)
+      transform.yield %res : !transform.any_op
   }
 }

compiler/src/iree/compiler/Codegen/Common/test/verify_tuning_specs.mlir

@@ -18,7 +18,7 @@ module @foo_module attributes { transform.with_named_sequence } {
 // -----
 
 module @foo_module attributes { transform.with_named_sequence } {
-  // expected-error @+1{{Tuning spec entry point expected to have a single any_op argument}}
+  // expected-error @+1{{Must take one 'any_op' (required by 'iree_codegen.tuning_spec_entrypoint')}}
   transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}, %arg1: !transform.any_op {transform.readonly}) -> !transform.any_op
     attributes { iree_codegen.tuning_spec_entrypoint } {
     transform.yield %arg0 : !transform.any_op
@@ -28,15 +28,15 @@ module @foo_module attributes { transform.with_named_sequence } {
 // -----
 
 module @foo_module attributes { transform.with_named_sequence } {
-  // expected-error @+1{{Tuning spec entry point expected to have a single any_op argument}}
+  // expected-error @+1{{Must take one 'any_op' (required by 'iree_codegen.tuning_spec_entrypoint')}}
   transform.named_sequence @foo(%arg0: i32) -> !transform.any_op
     attributes { iree_codegen.tuning_spec_entrypoint } {}
 }
 
 // -----
 
 module @foo_module attributes { transform.with_named_sequence } {
-  // expected-error @+1{{Tuning spec entry point expected to return any_op}}
+  // expected-error @+1{{Must return one 'any_op' (required by 'iree_codegen.tuning_spec_entrypoint')}}
   transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}) -> i32
     attributes { iree_codegen.tuning_spec_entrypoint } {
     %0 = arith.constant 0 : i32
@@ -47,22 +47,256 @@ module @foo_module attributes { transform.with_named_sequence } {
 // -----
 
 module @foo_module attributes { transform.with_named_sequence } {
-  // expected-error @+1{{Tuning spec entry point expected to return any_op}}
+  // expected-error @+1{{Must return one 'any_op' (required by 'iree_codegen.tuning_spec_entrypoint')}}
   transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly})
     attributes { iree_codegen.tuning_spec_entrypoint } {}
 }
 
 // -----
 
-// expected-error @+1{{The tuning specification must include a named sequence with the symbol name '__kernel_config'}}
+// expected-error @+1{{Missing named sequence '__kernel_config' (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
 module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
 }
 
 // -----
 
-// expected-error @+1{{The tuning specification must include a named sequence with the symbol name '__kernel_config'}}
+// expected-error @+1{{Missing named sequence '__kernel_config' (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
 module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
   func.func @__kernel_config(%arg0: i32) -> () {
     return
   }
 }
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  // expected-error @+1{{Missing attribute 'iree_codegen.tuning_spec_entrypoint' in named sequence '__kernel_config' (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+      -> (!transform.any_op) {
+      transform.yield %arg0 : !transform.any_op
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg: !transform.any_op) -> (!transform.any_op) {
+    transform.yield %arg : !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op) {
+    transform.yield
+  }
+
+  // expected-error @+1{{'__kernel_config' must contain exactly one block (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op)
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg: !transform.any_op {transform.readonly}) -> (!transform.any_op) {
+      transform.yield %arg : !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly}) {
+      transform.yield
+  }
+
+  // expected-error @+1{{'__kernel_config' must contain exactly two operations (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+
+      %res_a = transform.foreach_match in %arg0
+        @match -> @apply_op_config
+        : (!transform.any_op) -> (!transform.any_op)
+      %res_b = transform.foreach_match in %res_a
+        @match -> @apply_op_config
+        : (!transform.any_op) -> (!transform.any_op)
+      transform.yield %res_b : !transform.any_op
+  }
+}
+
+// -----
+
+// expected-error @+1{{Expected one named sequence with 'iree_codegen.tuning_spec_entrypoint', but found 2 (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+    transform.yield %arg0 : !transform.any_op
+  }
+
+  transform.named_sequence @main(%arg0: !transform.any_op {transform.readonly})
+      -> !transform.any_op attributes { iree_codegen.tuning_spec_entrypoint } {
+    transform.yield %arg0 : !transform.any_op
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+    transform.named_sequence @match(%arg: !transform.any_op {transform.readonly}) -> (!transform.any_op) {
+      transform.yield %arg : !transform.any_op
+    }
+
+    transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly}) {
+        transform.yield
+    }
+
+    // expected-error @+1{{'__kernel_config' must start with 'ForeachMatchOp' (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+    transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+        -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+        transform.print {name = "Hello"}
+        transform.yield %arg0 : !transform.any_op
+    }
+}
+
+// -----
+
+// expected-error @+1{{Expected one named sequence with 'iree_codegen.tuning_spec_entrypoint', but found 2 (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+    transform.yield %arg0 : !transform.any_op
+  }
+
+  module @extra_module attributes { iree_codegen.tuning_spec_entrypoint } {
+    transform.yield
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { transform.with_named_sequence, iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg: !transform.any_op {transform.readonly})
+    -> (!transform.any_op, !transform.any_op) {
+    transform.yield %arg, %arg : !transform.any_op, !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op1: !transform.any_op {transform.readonly}, %op2: !transform.any_op {transform.readonly})
+    -> (!transform.any_op) {
+    transform.yield %op1 : !transform.any_op
+  }
+
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+    -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+    // expected-error @+1 {{'ForeachMatchOp' must return exactly one 'any_op' result (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+    %res1, %res2 = transform.foreach_match in %arg0
+      @match -> @apply_op_config
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+
+    transform.yield %res1 : !transform.any_op
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { transform.with_named_sequence, iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg: !transform.any_op) -> (!transform.any_op) {
+    transform.yield %arg : !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op) {
+    transform.yield
+  }
+
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op)
+      -> (f32) attributes { iree_codegen.tuning_spec_entrypoint } {
+     // expected-error @+1 {{'ForeachMatchOp' must return exactly one 'any_op' result (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+    %res = transform.foreach_match in %arg0
+      @match -> @apply_op_config
+      : (!transform.any_op) -> (f32)
+
+    transform.yield %res : f32
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg1: !transform.any_op {transform.readonly}, %arg2: !transform.any_op {transform.readonly})
+        -> (!transform.any_op) {
+        transform.yield %arg1 : !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly}) {
+        transform.yield
+  }
+
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.consumed})
+    -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+    // expected-error @+1 {{'ForeachMatchOp' must take exactly one 'any_op' argument (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+    %res = transform.foreach_match in %arg0, %arg0
+      @match -> @apply_op_config
+      : (!transform.any_op, !transform.any_op) -> (!transform.any_op)
+
+    transform.yield %res : !transform.any_op
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg: index) -> (index) {
+    transform.yield %arg : index
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op) {
+    transform.yield
+  }
+
+  transform.named_sequence @__kernel_config(%arg0: index)
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+    // expected-error @+1 {{'ForeachMatchOp' must take exactly one 'any_op' argument (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+    %res = transform.foreach_match in %arg0
+      @match -> @apply_op_config
+      : (index) -> (!transform.any_op)
+
+    transform.yield %res : !transform.any_op
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg: !transform.any_op) -> (!transform.any_op) {
+    transform.yield %arg : !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op) {
+    transform.yield
+  }
+
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op)
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+     // expected-error @+1{{'ForeachMatchOp' must not have 'restrict_root' attribute (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+    %res = transform.foreach_match restrict_root in %arg0
+      @match -> @apply_op_config
+      : (!transform.any_op) -> (!transform.any_op)
+
+    transform.yield %res : !transform.any_op
+  }
+}
+
+// -----
+
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @match(%arg: !transform.any_op) -> (!transform.any_op) {
+    transform.yield %arg : !transform.any_op
+  }
+
+  transform.named_sequence @apply_op_config(%op: !transform.any_op) {
+    transform.yield
+  }
+
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op)
+      -> (!transform.any_op) attributes { iree_codegen.tuning_spec_entrypoint } {
+     // expected-error @+1{{'ForeachMatchOp' must not have 'flatten_results' attribute (required by 'iree_codegen.tuning_spec_with_default_entrypoint')}}
+    %res = transform.foreach_match flatten_results in %arg0
+      @match -> @apply_op_config
+      : (!transform.any_op) -> (!transform.any_op)
+
+    transform.yield %res : !transform.any_op
+  }
+}