[Codegen][Tuner] attr verifier for tuning specs

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

@@ -107,6 +107,16 @@ getUserTuningSpec(ModuleOp module, IREE::Codegen::IREECodegenDialect &dialect) {
            << clCodegenTuningSpecPath;
   }
 
+  // Iterate through all operations in the module to verify attributes
+  for (Operation &op : (*maybeTransformLibrary).getBody()->getOperations()) {
+    for (NamedAttribute attr : op.getAttrs()) {
+      if (failed(dialect.verifyOperationAttribute(&op, attr))) {
+        return op.emitError() << "Attribute verification failed for operation "
+                                 "in the user tuning spec";
+      }
+    }
+  }
+
   return *maybeTransformLibrary;
 }
 
@@ -138,8 +148,25 @@ getDefaultTuningSpec(ModuleOp module,
 
   // Load the library through the codegen dialect so that we cache the parsed
   // module.
-  return dialect.getOrParseTransformLibraryModule(defaultTuningSpecName,
-                                                  *defaultTuningSpecSource);
+  FailureOr<ModuleOp> defaultTransformLibrary =
+      dialect.getOrParseTransformLibraryModule(defaultTuningSpecName,
+                                               *defaultTuningSpecSource);
+  if (failed(defaultTransformLibrary)) {
+    return module->emitError()
+           << "Failed to parse default tuning spec transform library for '"
+           << arch << "'";
+  }
+  // Iterate through operations and validate their attributes
+  for (Operation &op : (*defaultTransformLibrary).getBody()->getOperations()) {
+    for (NamedAttribute attr : op.getAttrs()) {
+      if (failed(dialect.verifyOperationAttribute(&op, attr))) {
+        return op.emitError() << "Attribute verification failed for operation "
+                                 "in default tuning spec";
+      }
+    }
+  }
+
+  return *defaultTransformLibrary;
 }
 
 static FailureOr<DenseElementsAttr>
@@ -240,6 +267,17 @@ struct MaterializeTuningSpecsPass final
       return signalPassFailure();
     }
 
+    // Iterate through operations in the linked module and verify attributes
+    for (Operation &op : (linkedTuningSpec.get()).getBody()->getOperations()) {
+      for (NamedAttribute attr : op.getAttrs()) {
+        if (failed(dialect->verifyOperationAttribute(&op, attr))) {
+          op.emitError("Attribute verification failed for operation in linked "
+                       "tuning spec");
+          return signalPassFailure();
+        }
+      }
+    }
+
     if (failed(dumpFinalTuningSpecToDir(linkedTuningSpec.get()))) {
       return signalPassFailure();
     }

compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel

@@ -96,6 +96,7 @@ iree_lit_test_suite(
             "vector_layout_analysis.mlir",
             "vectorize_memref_copy.mlir",
             "vectorize_tensor_pad.mlir",
+            "verify_tuning_specs.mlir",
             "verify_workgroup_distribution.mlir",
             "vmvx_materialize_encoding.mlir",
         ],

compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt

@@ -92,6 +92,7 @@ iree_lit_test_suite(
     "vector_layout_analysis.mlir"
     "vectorize_memref_copy.mlir"
     "vectorize_tensor_pad.mlir"
+    "verify_tuning_specs.mlir"
     "verify_workgroup_distribution.mlir"
     "vmvx_materialize_encoding.mlir"
   TOOLS

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

@@ -0,0 +1,18 @@
+// RUN: iree-opt --no-implicit-module --verify-diagnostics -split-input-file --mlir-disable-threading %s
+
+module @td_module attributes { transform.with_named_sequence } {
+  module @foo_module attributes { transform.with_named_sequence } {
+    // expected-error @+1{{'iree_codegen.tuning_spec_entrypoint' attribute must be a UnitAttr}}
+    transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op
+      attributes { iree_codegen.tuning_spec_entrypoint = "foo" } {
+      transform.yield %arg0 : !transform.any_op
+    }
+    transform.named_sequence @bar(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op
+      attributes { iree_codegen.tuning_spec_entrypoint } {
+      transform.yield %arg0 : !transform.any_op
+    }
+    func.func @baz(%arg0: i32) -> () {
+      return
+    }
+  }
+}

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

@@ -10,6 +10,7 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp.inc"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/IR/DialectImplementation.h"
 
 namespace mlir::iree_compiler::IREE::Codegen {
@@ -45,4 +46,39 @@ void IREECodegenDialect::initialize() {
       >();
 }
 
+LogicalResult
+IREECodegenDialect::verifyOperationAttribute(Operation *op,
+                                             NamedAttribute attribute) {
+  StringRef symbol = attribute.getName().strref();
+  Attribute attr = attribute.getValue();
+
+  // Early return if the symbol is not "iree_codegen.tuning_spec_entrypoint"
+  if (symbol != kTuningSpecEntrypointAttrName)
+    return success();
+
+  // Verify that the attribute is a UnitAttr
+  if (!llvm::isa<UnitAttr>(attr)) {
+    return op->emitError("'") << symbol << "' attribute must be a UnitAttr";
+  }
+
+  if (auto namedSeqOp = dyn_cast<transform::NamedSequenceOp>(op)) {
+    ArrayRef<Type> resTypes = namedSeqOp.getFunctionType().getResults();
+    if (resTypes.size() != 1 || !isa<transform::AnyOpType>(resTypes[0])) {
+      return namedSeqOp.emitWarning()
+             << "Tuning spec entry point expected to return any_op";
+    }
+
+    ArrayRef<Type> argTypes = namedSeqOp.getArgumentTypes();
+    if (argTypes.size() != 1 || !isa<transform::AnyOpType>(argTypes[0])) {
+      return namedSeqOp.emitWarning()
+             << "Tuning spec entry point expected to have a "
+                "single any_op argument";
+    }
+
+    return success();
+  }
+
+  return success();
+}
+
 } // namespace mlir::iree_compiler::IREE::Codegen

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

@@ -68,6 +68,7 @@ def IREECodegen_Dialect : Dialect {
     std::mutex libraryMutex;
   }];
   let useDefaultAttributePrinterParser = 1;
+  let hasOperationAttrVerify = 1;
 }
 
 def AnyRankedTensorOrMemRefType : AnyTypeOf<[AnyRankedTensor, AnyMemRef]>;