[Codegen][Tuner] verifier for the default tuning spec

Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
iree-org · Dec 19, 2024 · b35931b · b35931b
1 parent e553425
commit b35931b
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 2 deletions.
diff --git a/compiler/plugins/target/ROCM/builtins/tuning/iree_default_tuning_spec_gfx942.mlir b/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.default_tuning_spec } {
 
 transform.named_sequence @apply_op_config(%op: !transform.any_op {transform.readonly},
                                         %config: !transform.any_param {transform.readonly}) {

diff --git a/compiler/src/iree/compiler/Codegen/Common/test/verify_tuning_specs.mlir b/compiler/src/iree/compiler/Codegen/Common/test/verify_tuning_specs.mlir
@@ -51,3 +51,9 @@ 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 default tuning specification must include an operation with the symbol name '__kernel_config'}}
+module @iree_default_tuning_spec attributes { iree_codegen.default_tuning_spec } {
+}
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h b/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 kTuningDefaultSpecAttrName =
+    "iree_codegen.default_tuning_spec";
 constexpr StringLiteral kTuningSpecEntrypointAttrName =
     "iree_codegen.tuning_spec_entrypoint";
 constexpr StringLiteral kSerializedTuningSpecAttrName =

diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp
@@ -51,8 +51,11 @@ 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` :
+  //   - For the `ModuleOp` operation ( representing the default spec):
+  //     - Ensure the module contains one operation with the symbol
+  //       name `__kernel_config`. If not, emit an error.
   // - If the attribute's name matches `kTuningSpecEntrypointAttrName`
   // ("iree_codegen.tuning_spec_entrypoint"):
   //   1. The attribute value must be a UnitAttr.
@@ -63,6 +66,19 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op,
   //         b. It must have exactly one argument type, and the argument must be
   //         of type `transform::AnyOpType`.
 
+  if (symbol == kTuningDefaultSpecAttrName) {
+    if (auto moduleOp = dyn_cast<ModuleOp>(op)) {
+      if (!llvm::any_of(moduleOp.getOps(), [](auto &op) {
+            return SymbolTable::getSymbolName(&op).getValue() ==
+                   kKernelConfigSpecName;
+          })) {
+        return moduleOp.emitError()
+               << "The default tuning specification must include an "
+                  "operation with the symbol name '__kernel_config'.";
+      }
+    }
+  }
+
   if (symbol != kTuningSpecEntrypointAttrName)
     return success();