From b35931b12e18b43688679454eebe599c9a95e0e2 Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Wed, 18 Dec 2024 19:37:02 -0600 Subject: [PATCH 1/6] [Codegen][Tuner] verifier for the default tuning spec Signed-off-by: Bangtian Liu --- .../iree_default_tuning_spec_gfx942.mlir | 2 +- .../Common/test/verify_tuning_specs.mlir | 6 ++++++ .../Dialect/Codegen/IR/IREECodegenAttrs.h | 2 ++ .../Dialect/Codegen/IR/IREECodegenDialect.cpp | 18 +++++++++++++++++- 4 files changed, 26 insertions(+), 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 index 53f56887d8df..e48885cab04e 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.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 index aede375adb5b..19e5d72a2886 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,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 index 59bf10a71a0a..4b146ea9029c 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 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 index 4a2281eef60e..bd4c89bc5932 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,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(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(); From e5350b7ba28a0e940c147a72bcedca6ab28c279b Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Thu, 19 Dec 2024 00:16:54 -0600 Subject: [PATCH 2/6] address reviewer comments, fix ci errors Signed-off-by: Bangtian Liu --- .../tuning/iree_default_tuning_spec_gfx942.mlir | 2 +- .../ROCM/test/default_tuning_specs_amdgpu.mlir | 4 ++-- .../Common/test/verify_tuning_specs.mlir | 2 +- .../Dialect/Codegen/IR/IREECodegenAttrs.h | 2 +- .../Dialect/Codegen/IR/IREECodegenDialect.cpp | 17 ++++++++++------- 5 files changed, 15 insertions(+), 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 index e48885cab04e..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, 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 index f049f79aede3..0f48bf27ea3c 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,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 index 19e5d72a2886..691a522c550d 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 @@ -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 index 4b146ea9029c..d80d1a961341 100644 --- 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 index bd4c89bc5932..68e2fe723380 100644 --- 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(op)) { if (!llvm::any_of(moduleOp.getOps(), [](auto &op) { - return SymbolTable::getSymbolName(&op).getValue() == - kKernelConfigSpecName; + if (auto namedSeqOp = dyn_cast(&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 << "'."; } } } From b2d4cd2d5308462428e5ce008cf51629395b5e86 Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Thu, 19 Dec 2024 00:03:38 -0800 Subject: [PATCH 3/6] update tuning.md and test file Signed-off-by: Bangtian Liu --- .../Codegen/Common/test/verify_tuning_specs.mlir | 11 ++++++++++- .../Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h | 2 +- .../Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp | 6 +++--- docs/website/docs/reference/tuning.md | 4 +++- 4 files changed, 17 insertions(+), 6 deletions(-) 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 691a522c550d..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 @@ -54,6 +54,15 @@ 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'}} +// 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 d80d1a961341..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,7 +45,7 @@ namespace mlir::iree_compiler { // Constant names. //===----------------------------------------------------------------------===// constexpr StringLiteral kConfigAttrName = "lowering_config"; -constexpr StringLiteral kTuningDefaultSpecAttrName = +constexpr StringLiteral kTuningSpecDefaultEntrypointAttrName = "iree_codegen.tuning_spec_with_default_entrypoint"; constexpr StringLiteral kTuningSpecEntrypointAttrName = "iree_codegen.tuning_spec_entrypoint"; 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 68e2fe723380..afd1d6232b04 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp @@ -64,7 +64,7 @@ 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 (symbol == kTuningSpecDefaultEntrypointAttrName) { if (auto moduleOp = dyn_cast(op)) { if (!llvm::any_of(moduleOp.getOps(), [](auto &op) { if (auto namedSeqOp = dyn_cast(&op)) { @@ -75,8 +75,8 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op, return false; })) { return moduleOp.emitError() - << "The default tuning specification must include an " - "operation with the symbol name '" + << "The tuning specification must include a named " + "sequence with the symbol name '" << kKernelConfigSpecName << "'."; } } diff --git a/docs/website/docs/reference/tuning.md b/docs/website/docs/reference/tuning.md index 5dea9b484282..2679afbc18e5 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 marked with `__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 From 2da59a01b32048f3bb153cb231ce8b7b44c51b1a Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Thu, 2 Jan 2025 15:06:42 -0600 Subject: [PATCH 4/6] handle the unit attribute during linking Signed-off-by: Bangtian Liu --- .../target/ROCM/test/default_tuning_specs_amdgpu.mlir | 4 ++-- .../src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp | 4 ++++ .../compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp | 5 +++++ .../Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp | 5 ++--- 4 files changed, 13 insertions(+), 5 deletions(-) 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 0f48bf27ea3c..a23d32319bd3 100644 --- a/compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir +++ b/compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir @@ -33,11 +33,11 @@ // 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 attributes {iree_codegen.tuning_spec_with_default_entrypoint, 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} -// BOTH-LABEL: module @iree_default_tuning_spec_gfx942_1 attributes {iree_codegen.tuning_spec_with_default_entrypoint, transform.with_named_sequence} +// BOTH-LABEL: module @iree_default_tuning_spec_gfx942_1 attributes {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/LinkTuningSpecsPass.cpp b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp index 1cdbfb1a821d..dab6ebd29294 100644 --- a/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp @@ -81,6 +81,10 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef specsToLink) { 0, hasConsumedSequences ? kArgConsumedAttrName : kArgReadOnlyAttrName, builder.getUnitAttr()); newSpec->setAttr(kTuningSpecEntrypointAttrName, builder.getUnitAttr()); + // As the newSpec is a named sequence operation with the symbol name + // '__kernel_config', the module should add the unit attribute + // 'iree_codegen.tuning_spec_with_default_entrypoint' to indicate this change. + module->setAttr(kTuningDefaultSpecAttrName, 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..9d6a44e10575 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(kTuningDefaultSpecAttrName)) { + clonedSpec->removeAttr(kTuningDefaultSpecAttrName); + } // 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/Dialect/Codegen/IR/IREECodegenDialect.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp index afd1d6232b04..39bf64333905 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp @@ -68,9 +68,8 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op, if (auto moduleOp = dyn_cast(op)) { if (!llvm::any_of(moduleOp.getOps(), [](auto &op) { if (auto namedSeqOp = dyn_cast(&op)) { - return SymbolTable::getSymbolName(namedSeqOp) - .getValue() - .contains(kKernelConfigSpecName); + return SymbolTable::getSymbolName(namedSeqOp).getValue() == + kKernelConfigSpecName; } return false; })) { From 20d858b99ebe1cb8b0fd53c65f7fb2122eaed565 Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Thu, 2 Jan 2025 16:25:01 -0600 Subject: [PATCH 5/6] address comments Signed-off-by: Bangtian Liu --- .../target/ROCM/test/default_tuning_specs_amdgpu.mlir | 8 ++++++-- .../iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp | 4 +++- .../Codegen/Common/MaterializeTuningSpecsPass.cpp | 4 ++-- .../Codegen/Common/test/materialize_tuning_specs.mlir | 4 +++- .../Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp | 3 +-- 5 files changed, 15 insertions(+), 8 deletions(-) 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 a23d32319bd3..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 {iree_codegen.tuning_spec_with_default_entrypoint, 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 {iree_codegen.tuning_spec_with_default_entrypoint, 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 dab6ebd29294..5011a9bc3a8e 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{}, @@ -84,7 +86,7 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef specsToLink) { // As the newSpec is a named sequence operation with the symbol name // '__kernel_config', the module should add the unit attribute // 'iree_codegen.tuning_spec_with_default_entrypoint' to indicate this change. - module->setAttr(kTuningDefaultSpecAttrName, 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 9d6a44e10575..c1fa52e94db0 100644 --- a/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp @@ -236,8 +236,8 @@ struct MaterializeTuningSpecsPass final ModuleOp clonedSpec = spec.clone(); // Drop the module-level attribute due to renamed entrypoints during // linking. - if (clonedSpec->hasAttr(kTuningDefaultSpecAttrName)) { - clonedSpec->removeAttr(kTuningDefaultSpecAttrName); + if (clonedSpec->hasAttr(kTuningSpecDefaultEntrypointAttrName)) { + clonedSpec->removeAttr(kTuningSpecDefaultEntrypointAttrName); } // Make sure there are no symbol name collisions. clonedSpec.setSymName( 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/Dialect/Codegen/IR/IREECodegenDialect.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp index 39bf64333905..2a44da83d03e 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp @@ -68,8 +68,7 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op, if (auto moduleOp = dyn_cast(op)) { if (!llvm::any_of(moduleOp.getOps(), [](auto &op) { if (auto namedSeqOp = dyn_cast(&op)) { - return SymbolTable::getSymbolName(namedSeqOp).getValue() == - kKernelConfigSpecName; + return namedSeqOp.getName() == kKernelConfigSpecName; } return false; })) { From f679b93b8203d3dc16519ad8a34412bdc1b62a79 Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Thu, 2 Jan 2025 18:00:28 -0600 Subject: [PATCH 6/6] remove code comment and format the code Signed-off-by: Bangtian Liu --- .../compiler/Codegen/Common/LinkTuningSpecsPass.cpp | 3 --- .../Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp | 10 ++++------ docs/website/docs/reference/tuning.md | 2 +- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp index 5011a9bc3a8e..6b1502c23aac 100644 --- a/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp @@ -83,9 +83,6 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef specsToLink) { 0, hasConsumedSequences ? kArgConsumedAttrName : kArgReadOnlyAttrName, builder.getUnitAttr()); newSpec->setAttr(kTuningSpecEntrypointAttrName, builder.getUnitAttr()); - // As the newSpec is a named sequence operation with the symbol name - // '__kernel_config', the module should add the unit attribute - // 'iree_codegen.tuning_spec_with_default_entrypoint' to indicate this change. module->setAttr(kTuningSpecDefaultEntrypointAttrName, builder.getUnitAttr()); Region ®ion = newSpec.getRegion(); 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 2a44da83d03e..8b02d9dbf6fd 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.cpp @@ -66,12 +66,10 @@ IREECodegenDialect::verifyOperationAttribute(Operation *op, if (symbol == kTuningSpecDefaultEntrypointAttrName) { if (auto moduleOp = dyn_cast(op)) { - if (!llvm::any_of(moduleOp.getOps(), [](auto &op) { - if (auto namedSeqOp = dyn_cast(&op)) { - return namedSeqOp.getName() == kKernelConfigSpecName; - } - return false; - })) { + 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 '" diff --git a/docs/website/docs/reference/tuning.md b/docs/website/docs/reference/tuning.md index 2679afbc18e5..4a598b5e9c57 100644 --- a/docs/website/docs/reference/tuning.md +++ b/docs/website/docs/reference/tuning.md @@ -124,7 +124,7 @@ that conform to the following format: * 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 marked with `__kernel_config`. + 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