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 index 53f56887d8df..1ce1401b16c0 100644 --- 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.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 index f049f79aede3..fd6f915ef3a7 100644 --- 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,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} diff --git a/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp index 1cdbfb1a821d..6b1502c23aac 100644 --- a/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp @@ -72,6 +72,8 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef specsToLink) { Type anyOpType = builder.getType(); 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( loc, kKernelConfigSpecName, TypeAttr::get(specType), /*sym_visibility=*/StringAttr{}, @@ -81,6 +83,7 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef specsToLink) { 0, hasConsumedSequences ? kArgConsumedAttrName : kArgReadOnlyAttrName, builder.getUnitAttr()); newSpec->setAttr(kTuningSpecEntrypointAttrName, builder.getUnitAttr()); + module->setAttr(kTuningSpecDefaultEntrypointAttrName, builder.getUnitAttr()); Region ®ion = newSpec.getRegion(); Block *body = builder.createBlock(®ion, region.begin(), diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp index db36cbf3b36b..c1fa52e94db0 100644 --- a/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp +++ b/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()); diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs.mlir index 8e2a081d614b..bd1e7edb353b 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs.mlir +++ b/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} 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 index aede375adb5b..6a3c6b292ecd 100644 --- 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,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 + } +} diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h index 59bf10a71a0a..cae3f0829599 100644 --- 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 kTuningSpecDefaultEntrypointAttrName = + "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 index 4a2281eef60e..8b02d9dbf6fd 100644 --- 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,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(op)) { + if (!llvm::any_of(moduleOp.getOps(), + [](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(); diff --git a/docs/website/docs/reference/tuning.md b/docs/website/docs/reference/tuning.md index 5dea9b484282..4a598b5e9c57 100644 --- a/docs/website/docs/reference/tuning.md +++ b/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