[Codegen][Tuner] verifier for the default tuning spec (#19525)

This PR adds the unit attribute`
iree_codegen.tuning_spec_with_default_entrypoint` to indicate the
default tuning spec (typically or user-provided tuning spec but can work
in the same manner) must contain one named sequence operation marked
with `__kernel_config`, also add the corresponding verification in
`verifyOperationAttribute` function.

This PR is relevant to task in
#19214: add [a discardable attr
verifier](https://mlir.llvm.org/docs/DefiningDialects/#discardable-attribute-verification)
for entry points iree_codegen.tuning_spec_entrypoint

Context:
Jakub proposed two approaches for verifying the default tuning
specification:
1. Implement a dedicated pass for verification.
2. Add a new attribute and update the verifyOperationAttribute function
accordingly.

After careful consideration, we agreed on the second approach to avoid
introducing an additional pass, ensuring a simple implementation.

---------

Signed-off-by: Bangtian Liu <[email protected]>

Author: bangtianliu

compiler/plugins/target/ROCM/builtins/tuning/iree_default_tuning_spec_gfx942.mlir

@@ -5,7 +5,7 @@
 // TODO(https://github.com/iree-org/iree/issues/19214): Add missing
 // configurations to this spec.
 
-module @iree_default_tuning_spec_gfx942 attributes { transform.with_named_sequence } {
+module @iree_default_tuning_spec_gfx942 attributes { transform.with_named_sequence, iree_codegen.tuning_spec_with_default_entrypoint } {
 
 transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly},
                                         %config: !transform.any_param {transform.readonly}) {

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

@@ -19,7 +19,9 @@
 
 // Check that the default tuning spec gets materialized without linking.
 
-// DEFAULT-LABEL: module @iree_default_tuning_spec_gfx942 attributes {transform.with_named_sequence}
+// DEFAULT-LABEL: module @iree_default_tuning_spec_gfx942
+// DEFAULT-SAME:    iree_codegen.tuning_spec_with_default_entrypoint
+// DEFAULT-SAME:    transform.with_named_sequence
 // DEFAULT-LABEL:   transform.named_sequence @__kernel_config
 // DEFAULT-SAME:      attributes {iree_codegen.tuning_spec_entrypoint}
 
@@ -33,7 +35,9 @@
 // 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.
 
-// BOTH-LABEL: module @iree_linked_tuning_spec attributes {transform.with_named_sequence}
+// 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
 // BOTH-SAME:        attributes {iree_codegen.tuning_spec_entrypoint}

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

@@ -72,6 +72,8 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef<NamedSequenceOp> specsToLink) {
   Type anyOpType = builder.getType<transform::AnyOpType>();
   FunctionType specType =
       builder.getFunctionType(TypeRange{anyOpType}, TypeRange{anyOpType});
+  // This code creates a named sequence operation that conforms to the
+  // requirements for tuning specifications with a default entry point.
   auto newSpec = builder.create<NamedSequenceOp>(
       loc, kKernelConfigSpecName, TypeAttr::get(specType),
       /*sym_visibility=*/StringAttr{},
@@ -81,6 +83,7 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef<NamedSequenceOp> specsToLink) {
       0, hasConsumedSequences ? kArgConsumedAttrName : kArgReadOnlyAttrName,
       builder.getUnitAttr());
   newSpec->setAttr(kTuningSpecEntrypointAttrName, builder.getUnitAttr());
+  module->setAttr(kTuningSpecDefaultEntrypointAttrName, builder.getUnitAttr());
 
   Region &region = newSpec.getRegion();
   Block *body = builder.createBlock(&region, region.begin(),

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

@@ -234,6 +234,11 @@ struct MaterializeTuningSpecsPass final
         UnitAttr::get(ctx));
     for (auto [idx, spec] : llvm::enumerate(allSpecs)) {
       ModuleOp clonedSpec = spec.clone();
+      // Drop the module-level attribute due to renamed entrypoints during
+      // linking.
+      if (clonedSpec->hasAttr(kTuningSpecDefaultEntrypointAttrName)) {
+        clonedSpec->removeAttr(kTuningSpecDefaultEntrypointAttrName);
+      }
       // Make sure there are no symbol name collisions.
       clonedSpec.setSymName(
           llvm::formatv("{}_{}", clonedSpec.getSymName().value(), idx).str());

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

@@ -5,7 +5,9 @@
 
 // Check that the final tuning spec is as expected when the user tuning spec is provided.
 
-// CHECK-LABEL: module @iree_linked_tuning_spec attributes {transform.with_named_sequence}
+// 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
 // CHECK-SAME:        attributes {iree_codegen.tuning_spec_entrypoint}

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

@@ -51,3 +51,18 @@ module @foo_module attributes { transform.with_named_sequence } {
   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'}}
+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'}}
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
+  func.func @__kernel_config(%arg0: i32) -> () {
+    return
+  }
+}

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

@@ -45,6 +45,8 @@ namespace mlir::iree_compiler {
 // Constant names.
 //===----------------------------------------------------------------------===//
 constexpr StringLiteral kConfigAttrName = "lowering_config";
+constexpr StringLiteral kTuningSpecDefaultEntrypointAttrName =
+    "iree_codegen.tuning_spec_with_default_entrypoint";
 constexpr StringLiteral kTuningSpecEntrypointAttrName =
     "iree_codegen.tuning_spec_entrypoint";
 constexpr StringLiteral kSerializedTuningSpecAttrName =

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

@@ -51,8 +51,9 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op,
                                              NamedAttribute attribute) {
   StringRef symbol = attribute.getName().strref();
   Attribute attr = attribute.getValue();
-
   // This function verifies the validity of a specific operation attribute.
+  // - If the attribute's name matches `kTuningDefaultSpecAttrName`, make
+  //   sure it contains a single named sequence op with name `__kernel_config`.
   // - If the attribute's name matches `kTuningSpecEntrypointAttrName`
   // ("iree_codegen.tuning_spec_entrypoint"):
   //   1. The attribute value must be a UnitAttr.
@@ -63,6 +64,20 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op,
   //         b. It must have exactly one argument type, and the argument must be
   //         of type `transform::AnyOpType`.
 
+  if (symbol == kTuningSpecDefaultEntrypointAttrName) {
+    if (auto moduleOp = dyn_cast<ModuleOp>(op)) {
+      if (!llvm::any_of(moduleOp.getOps<transform::NamedSequenceOp>(),
+                        [](transform::NamedSequenceOp op) {
+                          return op.getName() == kKernelConfigSpecName;
+                        })) {
+        return moduleOp.emitError()
+               << "The tuning specification must include a named "
+                  "sequence with the symbol name '"
+               << kKernelConfigSpecName << "'.";
+      }
+    }
+  }
+
   if (symbol != kTuningSpecEntrypointAttrName)
     return success();
 

docs/website/docs/reference/tuning.md

@@ -50,7 +50,7 @@ attempting the default one.
 ### Example
 
 ```mlir
-module @my_spec attributes { transform.with_named_sequence } {
+module @my_spec attributes { transform.with_named_sequence, iree_codegen.tuning_spec_with_default_entrypoint } {
 transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly},
                                           %config: !transform.any_param {transform.readonly}) {
   transform.annotate %op "compilation_info" = %config : !transform.any_op, !transform.any_param
@@ -123,6 +123,8 @@ that conform to the following format:
   `!transform.any_op`.
 * All entry points in the final tuning specs must either read
   (`transform.readonly`) or consume (`transform.consumed`) the argument.
+* The `iree_codegen.tuning_spec_with_default_entrypoint` attribute ensures that
+  the tuning spec includes a named sequence op with name `__kernel_config`.
 
 The tuning spec above attempts to match `linalg.generic` ops that correspond to the
 matmul operation with the RHS operand transposed (a.k.a. mmt) of shape