address reviewer comments, fix ci errors

Signed-off-by: Bangtian Liu <[email protected]>
iree-org · Dec 19, 2024 · e5350b7 · e5350b7
1 parent b35931b
commit e5350b7
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 12 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, iree_codegen.default_tuning_spec } {
+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}) {

diff --git a/compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir b/compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir
@@ -19,7 +19,7 @@
 
 // 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 attributes {iree_codegen.tuning_spec_with_default_entrypoint, transform.with_named_sequence}
 // DEFAULT-LABEL:   transform.named_sequence @__kernel_config
 // DEFAULT-SAME:      attributes {iree_codegen.tuning_spec_entrypoint}
 
@@ -37,7 +37,7 @@
 // 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}
-// BOTH-LABEL:   module @iree_default_tuning_spec_gfx942_1 attributes {transform.with_named_sequence}
+// BOTH-LABEL:   module @iree_default_tuning_spec_gfx942_1 attributes {iree_codegen.tuning_spec_with_default_entrypoint, transform.with_named_sequence}
 // BOTH:           transform.named_sequence @__kernel_config
 // BOTH-SAME:        attributes {iree_codegen.tuning_spec_entrypoint}
 // BOTH:         transform.named_sequence @__kernel_config

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
@@ -55,5 +55,5 @@ module @foo_module attributes { transform.with_named_sequence } {
 // -----
 
 // 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 } {
+module @iree_default_tuning_spec attributes { iree_codegen.tuning_spec_with_default_entrypoint } {
 }
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h
@@ -46,7 +46,7 @@ namespace mlir::iree_compiler {
 //===----------------------------------------------------------------------===//
 constexpr StringLiteral kConfigAttrName = "lowering_config";
 constexpr StringLiteral kTuningDefaultSpecAttrName =
-    "iree_codegen.default_tuning_spec";
+    "iree_codegen.tuning_spec_with_default_entrypoint";
 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
@@ -52,10 +52,8 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op,
   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 `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.
@@ -69,12 +67,17 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op,
   if (symbol == kTuningDefaultSpecAttrName) {
     if (auto moduleOp = dyn_cast<ModuleOp>(op)) {
       if (!llvm::any_of(moduleOp.getOps(), [](auto &op) {
-            return SymbolTable::getSymbolName(&op).getValue() ==
-                   kKernelConfigSpecName;
+            if (auto namedSeqOp = dyn_cast<transform::NamedSequenceOp>(&op)) {
+              return SymbolTable::getSymbolName(namedSeqOp)
+                  .getValue()
+                  .contains(kKernelConfigSpecName);
+            }
+            return false;
           })) {
         return moduleOp.emitError()
                << "The default tuning specification must include an "
-                  "operation with the symbol name '__kernel_config'.";
+                  "operation with the symbol name '"
+               << kKernelConfigSpecName << "'.";
       }
     }
   }