From 7875df9433faff36bc91d12036153c77d00c34f4 Mon Sep 17 00:00:00 2001 From: Jakub Kuderski Date: Tue, 26 Nov 2024 20:23:26 -0500 Subject: [PATCH] [Codegen] Add pass to materialize tuning specs ... and update 'Materialize User Configs' to pick up those tuning specs. The overall flow is as follows: * We pick up any user-specified tuning specs in `materialize tuning specs` and link them into a single transform dialect library module. * We serialize that linked tuning spec as MLIR bytecode. * We embed this MLIR bytecode as a module attribute. This is so that none of the subsequent passes will accidentally `walk` or otherwise modify it. * In `materilize user configs`, we first check if there are any transform libraries provided. If not, then we check if the tuning spec is present. * We deserialize the tuning spec attribute into a transform dialect library module and execute it. * We remove the serialized tuning spec from the module, as it's no longer needed. Signed-off-by: Jakub Kuderski --- .../iree/compiler/Codegen/Common/BUILD.bazel | 4 + .../compiler/Codegen/Common/CMakeLists.txt | 4 + .../Codegen/Common/LinkTuningSpecsPass.cpp | 47 ++--- .../Common/MaterializeTuningSpecsPass.cpp | 166 ++++++++++++++++++ .../Codegen/Common/MaterializeUserConfigs.cpp | 57 +++++- .../src/iree/compiler/Codegen/Common/Passes.h | 7 + .../iree/compiler/Codegen/Common/Passes.td | 19 ++ .../compiler/Codegen/Common/test/BUILD.bazel | 5 + .../Codegen/Common/test/CMakeLists.txt | 4 + .../Common/test/materialize_tuning_specs.mlir | 23 +++ ...materialize_tuning_specs_invalid_spec.mlir | 12 ++ ...erialize_user_config_from_tuning_spec.mlir | 42 +++++ .../Codegen/Common/test/tuning_spec.mlir | 9 + .../Dialect/Codegen/IR/IREECodegenAttrs.h | 4 +- .../Dialect/Codegen/IR/IREECodegenDialect.td | 2 +- .../Codegen/IR/IREECodegenLibraryManager.cpp | 42 +++-- 16 files changed, 411 insertions(+), 36 deletions(-) create mode 100644 compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp create mode 100644 compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs.mlir create mode 100644 compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs_invalid_spec.mlir create mode 100644 compiler/src/iree/compiler/Codegen/Common/test/materialize_user_config_from_tuning_spec.mlir create mode 100644 compiler/src/iree/compiler/Codegen/Common/test/tuning_spec.mlir diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel index 776b03d91a99..e3513ba69d29 100644 --- a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel @@ -127,6 +127,7 @@ iree_compiler_cc_library( "LowerUKernelsToCalls.cpp", "MaterializeEncodingIntoNop.cpp", "MaterializeEncodingIntoPackUnPack.cpp", + "MaterializeTuningSpecsPass.cpp", "MemrefCopyToLinalg.cpp", "NormalizeLoopBounds.cpp", "OptimizeTensorInsertExtractSlices.cpp", @@ -201,6 +202,7 @@ iree_compiler_cc_library( "@llvm-project//mlir:BufferizationDialect", "@llvm-project//mlir:BufferizationInterfaces", "@llvm-project//mlir:BufferizationTransforms", + "@llvm-project//mlir:BytecodeWriter", "@llvm-project//mlir:DestinationStyleOpInterface", "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:FuncDialect", @@ -219,6 +221,7 @@ iree_compiler_cc_library( "@llvm-project//mlir:MemRefDialect", "@llvm-project//mlir:MemRefTransforms", "@llvm-project//mlir:MemRefUtils", + "@llvm-project//mlir:Parser", "@llvm-project//mlir:Pass", "@llvm-project//mlir:SCFDialect", "@llvm-project//mlir:SCFToControlFlow", @@ -284,6 +287,7 @@ iree_compiler_cc_library( "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:LinalgDialect", "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Parser", "@llvm-project//mlir:PDLDialect", "@llvm-project//mlir:PDLInterpDialect", "@llvm-project//mlir:SCFDialect", diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt index f7ed254c31f3..adec8aad7583 100644 --- a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt @@ -119,6 +119,7 @@ iree_cc_library( "LowerUKernelsToCalls.cpp" "MaterializeEncodingIntoNop.cpp" "MaterializeEncodingIntoPackUnPack.cpp" + "MaterializeTuningSpecsPass.cpp" "MemrefCopyToLinalg.cpp" "NormalizeLoopBounds.cpp" "OptimizeTensorInsertExtractSlices.cpp" @@ -163,6 +164,7 @@ iree_cc_library( MLIRArithUtils MLIRBufferizationDialect MLIRBufferizationTransforms + MLIRBytecodeWriter MLIRDestinationStyleOpInterface MLIRFuncDialect MLIRFuncTransforms @@ -180,6 +182,7 @@ iree_cc_library( MLIRMemRefDialect MLIRMemRefTransforms MLIRMemRefUtils + MLIRParser MLIRPass MLIRSCFDialect MLIRSCFToControlFlow @@ -257,6 +260,7 @@ iree_cc_library( MLIRMemRefTransformOps MLIRPDLDialect MLIRPDLInterpDialect + MLIRParser MLIRPass MLIRRewrite MLIRSCFDialect diff --git a/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp index ab9ddce82dd0..8f57104d6f07 100644 --- a/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp @@ -44,8 +44,9 @@ findNestedModulesWithNamedSequences(ModuleOp module) { static SmallVector findTuningSpecs(ModuleOp module) { Block *body = module.getBody(); return llvm::filter_to_vector( - body->getOps(), - [](NamedSequenceOp op) { return op->hasAttr(kTuningSpecAttrName); }); + body->getOps(), [](NamedSequenceOp op) { + return op->hasAttr(kTuningSpecEntrypointAttrName); + }); } static LogicalResult validateTuningSpec(NamedSequenceOp op) { @@ -85,7 +86,7 @@ emitLinkedTuningSpec(ModuleOp module, ArrayRef specsToLink) { /*res_attrs*/ ArrayAttr{}); newSpec.setArgAttr(0, transform::TransformDialect::kArgReadOnlyAttrName, builder.getUnitAttr()); - newSpec->setAttr(kTuningSpecAttrName, builder.getUnitAttr()); + newSpec->setAttr(kTuningSpecEntrypointAttrName, builder.getUnitAttr()); Region ®ion = newSpec.getRegion(); Block *body = builder.createBlock(®ion, region.begin(), @@ -122,28 +123,34 @@ struct LinkTuningSpecsPass final } void runOnOperation() override { - ModuleOp module = getOperation(); - SmallVector tuningSpecs; - - for (ModuleOp nested : findNestedModulesWithNamedSequences(module)) { - llvm::append_range(tuningSpecs, findTuningSpecs(nested)); + if (failed(linkTuningSpecs(getOperation()))) { + signalPassFailure(); } + } +}; - for (NamedSequenceOp spec : tuningSpecs) { - LDBG("Found tuning spec: " << spec.getSymName()); - if (failed(validateTuningSpec(spec))) { - return signalPassFailure(); - } - } +} // namespace + +FailureOr linkTuningSpecs(ModuleOp module) { + SmallVector tuningSpecs; - if (tuningSpecs.empty()) { - LDBG("No tuning specs found, exiting without linking"); - return; + for (ModuleOp nested : findNestedModulesWithNamedSequences(module)) { + llvm::append_range(tuningSpecs, findTuningSpecs(nested)); + } + + for (NamedSequenceOp spec : tuningSpecs) { + LDBG("Found tuning spec: " << spec.getSymName()); + if (failed(validateTuningSpec(spec))) { + return failure(); } + } - emitLinkedTuningSpec(module, tuningSpecs); + if (tuningSpecs.empty()) { + LDBG("No tuning specs found, exiting without linking"); + return NamedSequenceOp{}; } -}; -} // namespace + return emitLinkedTuningSpec(module, tuningSpecs); +} + } // namespace mlir::iree_compiler diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp new file mode 100644 index 000000000000..f14aa92d5a66 --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp @@ -0,0 +1,166 @@ +// Copyright 2024 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include +#include "iree/compiler/Codegen/Common/Passes.h" +#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h" +#include "llvm/ADT/SmallString.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/ToolOutputFile.h" +#include "llvm/Support/raw_ostream.h" +#include "mlir/Bytecode/BytecodeWriter.h" +#include "mlir/Dialect/Transform/IR/TransformDialect.h" +#include "mlir/Dialect/Transform/IR/TransformOps.h" +#include "mlir/Dialect/Transform/IR/TransformTypes.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/Support/FileUtilities.h" + +#define DEBUG_TYPE "iree-codegen-materialize-tuning-specs" +#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") +#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") + +namespace mlir::iree_compiler { + +#define GEN_PASS_DEF_MATERIALIZETUNINGSPECSPASS +#include "iree/compiler/Codegen/Common/Passes.h.inc" + +namespace { + +llvm::cl::opt clCodegenTuningSpecPath( + "iree-codegen-tuning-spec-path", + llvm::cl::desc("File path to a module containing a tuning spec (transform " + "dialect library)."), + llvm::cl::init("")); + +llvm::cl::opt clCodegenTuningSpecDumpDir( + "iree-codegen-dump-tuning-specs-to", + llvm::cl::desc( + "Dump the final tuning spec modules to the specified directory. When " + "set to '-', prints the tuning spec to stdout."), + llvm::cl::init("")); + +using mlir::transform::NamedSequenceOp; + +static LogicalResult dumpFinalTuningSpecToDir(ModuleOp tuningSpec, + StringRef dir) { + if (dir == "-") { + tuningSpec->print(llvm::outs()); + return success(); + } + + llvm::sys::fs::create_directories(dir); + llvm::SmallString<64> dumpPath; + auto dumpFileEC = llvm::sys::fs::createUniqueFile( + Twine(dir) + "/iree_tuning_spec_%%.mlir", dumpPath); + if (dumpFileEC) { + return tuningSpec->emitError() + << "Failed to create a unique file in " << dir << "\n"; + } + LDBG("Linked tuning spec file path: " << dumpPath); + + std::string error; + auto file = mlir::openOutputFile(dumpPath, &error); + if (!file) { + return tuningSpec->emitError() + << "Failed to open a tuning spec dump file " << dumpPath << "\n"; + } + + tuningSpec->print(file->os()); + file->keep(); + return success(); +} + +static FailureOr +serializeTuningSpecToAttr(ModuleOp tuningSpec) { + std::string buffer; + llvm::raw_string_ostream os(buffer); + if (failed(writeBytecodeToFile(tuningSpec, os))) { + return failure(); + } + + auto bufferSize = static_cast(buffer.size()); + auto bufferShape = VectorType::get( + bufferSize, IntegerType::get(tuningSpec->getContext(), 8)); + return DenseElementsAttr::getFromRawBuffer( + bufferShape, ArrayRef(buffer.data(), buffer.data() + bufferSize)); +} + +struct MaterializeTuningSpecsPass final + : impl::MaterializeTuningSpecsPassBase { + void getDependentDialects(DialectRegistry ®istry) const override { + registerTransformDialectTranslationDependentDialects(registry); + } + + void runOnOperation() override { + if (clCodegenTuningSpecPath.empty()) { + return; + } + + ModuleOp module = getOperation(); + MLIRContext *ctx = &getContext(); + auto dialect = ctx->getOrLoadDialect(); + auto maybeTransformLibrary = + dialect->getOrLoadTransformLibraryModule(clCodegenTuningSpecPath); + if (failed(maybeTransformLibrary)) { + module->emitError() + << "Failed to load tuning spec transform dialect library from " + << clCodegenTuningSpecPath; + return signalPassFailure(); + } + + ModuleOp userTuningSpec = *maybeTransformLibrary; + if (!userTuningSpec.getSymName()) { + // Set a module name so that we can refer to its nested symbols. + userTuningSpec.setSymName("iree_user_tuning_spec"); + } + + Location loc = userTuningSpec.getLoc(); + + // This module will always be released at the end of the pass. + OwningOpRef linkedTuningSpec( + ModuleOp::create(loc, "iree_linked_tuning_spec")); + linkedTuningSpec.get()->setAttr( + transform::TransformDialect::kWithNamedSequenceAttrName, + UnitAttr::get(ctx)); + linkedTuningSpec->insert(linkedTuningSpec->begin(), userTuningSpec.clone()); + + // TODO(https://github.com/iree-org/iree/issues/19214): Add linked tuning + // spec memoization to IREECodegenDialect. We should be able to provide a + // list of input libraries that may have already been linked and ask the + // dialect to return it to us, or invoke a callback that will insert it if + // not found. + FailureOr newEntrypoint = + linkTuningSpecs(linkedTuningSpec.get()); + if (failed(newEntrypoint)) { + module->emitError("Failed to link tuning specs"); + return signalPassFailure(); + } + + if (!clCodegenTuningSpecDumpDir.empty()) { + if (failed(dumpFinalTuningSpecToDir(linkedTuningSpec.get(), + clCodegenTuningSpecDumpDir))) { + return signalPassFailure(); + } + } + + FailureOr serializedSpec = + serializeTuningSpecToAttr(linkedTuningSpec.get()); + if (failed(serializedSpec)) { + module->emitError("Failed to serialize linked tuning specs"); + return signalPassFailure(); + } + module->setAttr(kSerializedTuningSpecAttrName, *serializedSpec); + } +}; + +} // namespace +} // namespace mlir::iree_compiler diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeUserConfigs.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeUserConfigs.cpp index c4c97925eefe..21fee4a3f065 100644 --- a/compiler/src/iree/compiler/Codegen/Common/MaterializeUserConfigs.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeUserConfigs.cpp @@ -4,15 +4,17 @@ // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +#include #include "iree/compiler/Codegen/Common/Passes.h" #include "iree/compiler/Codegen/Common/UserConfig.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h" #include "llvm/ADT/StringRef.h" -#include "llvm/Support/LogicalResult.h" #include "mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/Parser/Parser.h" #define DEBUG_TYPE "iree-codegen-materialize-user-configs" #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") @@ -110,6 +112,40 @@ getTransformLibraryFromPath(ModuleOp compiledModule, StringRef path) { entrySequenceName.str()}; } +/// Look up the tuning spec in the given module or any of its parents. +static LogicalResult getModuleTuningSpec(ModuleOp compiledModule, + OwningOpRef &tuningSpec) { + IREE::Util::SerializableAttrInterface serializedTuningSpec; + Operation *op = compiledModule; + while (!serializedTuningSpec && op) { + serializedTuningSpec = + op->getAttrOfType( + kSerializedTuningSpecAttrName); + op = op->getParentOp(); + } + + if (!serializedTuningSpec) { + return failure(); + } + + SmallVector bytecode; + if (failed(serializedTuningSpec.serializeToVector( + compiledModule->getLoc(), llvm::endianness::native, bytecode))) { + return compiledModule.emitError() + << "Failed to read attribute " << kSerializedTuningSpecAttrName; + } + + ParserConfig config(compiledModule.getContext()); + tuningSpec = parseSourceString( + StringRef(bytecode.data(), bytecode.size()), config); + if (!tuningSpec) { + return compiledModule.emitError() << "Failed to parse tuning spec in " + << kSerializedTuningSpecAttrName; + } + LDBG("--loaded tuning spec"); + return success(); +} + struct MaterializeUserConfigsPass final : impl::MaterializeUserConfigsPassBase { void getDependentDialects(DialectRegistry ®istry) const override { @@ -119,9 +155,28 @@ struct MaterializeUserConfigsPass final void runOnOperation() override { ModuleOp moduleOp = getOperation(); + // Try to load the transform library from the user flag first. If none is + // specified, fall back to using the module tuning spec. FailureOr userTransformLibrary = getTransformLibraryFromPath(moduleOp, clCodegenTransformDialectLibraryFileName); + OwningOpRef tuningSpec; + if (failed(userTransformLibrary)) { + if (succeeded(getModuleTuningSpec(moduleOp, tuningSpec))) { + assert(tuningSpec); + userTransformLibrary = TransformLibraryWithEntrypoint{ + tuningSpec.get(), kKernelConfigSpecName.str()}; + } + } + + // Remove the tuning spec, if any, from the current module. If the tuning + // spec is attached to some other parent op, we conservatively keep it + // as-is, as we are not sure who the producer is and if they want it + // removed. + if (moduleOp->hasAttr(kSerializedTuningSpecAttrName)) { + moduleOp->removeAttr(kSerializedTuningSpecAttrName); + LDBG("--dropped the serialized tuning spec from the module"); + } for (auto funcOp : moduleOp.getOps()) { diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.h b/compiler/src/iree/compiler/Codegen/Common/Passes.h index eac457dc6280..2938bdd87da5 100644 --- a/compiler/src/iree/compiler/Codegen/Common/Passes.h +++ b/compiler/src/iree/compiler/Codegen/Common/Passes.h @@ -21,6 +21,7 @@ #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/Transform/IR/TransformDialect.h" +#include "mlir/Dialect/Transform/IR/TransformOps.h" #include "mlir/Pass/Pass.h" namespace mlir::iree_compiler { @@ -52,6 +53,12 @@ void addConstantBufferizePasses(OpPassManager &funcPassManager); /// Populate Encoding to Nop pass and canonicalizer pass to the pipeline void addEncodingToNopPasses(FunctionLikeNest &passManager); +/// Links nested transform dialect tuning specs named sequences into a single +/// entry point. Returns the new named sequence op (inserted into the `module`) +/// that includes the nested tuning specs, or a null op when no nested named +/// sequences were found. +FailureOr linkTuningSpecs(ModuleOp module); + //------------------------------------------------------------------------------ // Wrappers that not use tablegen options. See Passes.td for details. //------------------------------------------------------------------------------ diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.td b/compiler/src/iree/compiler/Codegen/Common/Passes.td index 9852cc693e5c..5471c95b0cad 100644 --- a/compiler/src/iree/compiler/Codegen/Common/Passes.td +++ b/compiler/src/iree/compiler/Codegen/Common/Passes.td @@ -436,6 +436,25 @@ def MaterializeEncodingIntoNopPass : let summary = "Drop the encodings from tensor types with encodings."; } +def MaterializeTuningSpecsPass : Pass<"iree-codegen-materialize-tuning-specs", "ModuleOp"> { + let summary = + "Load tuning spec transform dialect libraries and encode them in the module"; + let description = [{ + Links all available tuning spec transform dialect modules into a single + tuning spec. Next, serializes this tuning spec to bytecode and attaches it + as a module attribute. We do this so that the full tuning spec is always + encoded in the program IR and can be checked with `--mlir-print-ir-after-all` + (or equivalent). The alternative would be to add the tuning spec as a + submodule in the compiled program, but this may result in the tuning spec + being inadvertently visited by other passes that attempt to `walk` the outer + module. Serialization makes the tuning specs opaque and prevents it from + happening. + + This attribute is expected to be short-lived and removed by + `iree-codegen-materialize-user-configs`. + }]; +} + def MaterializeUserConfigsPass : Pass<"iree-codegen-materialize-user-configs", "ModuleOp"> { let summary = "Sets the lowering configs and translation info from user configs"; let dependentDialects = [ diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel index 4a89365196c6..5644f4855ab7 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel @@ -55,6 +55,9 @@ iree_lit_test_suite( "link_tuning_specs.mlir", "lower_ukernel_to_calls.mlir", "materialize_encoding_into_nop.mlir", + "materialize_tuning_specs.mlir", + "materialize_tuning_specs_invalid_spec.mlir", + "materialize_user_config_from_tuning_spec.mlir", "materialize_user_configs.mlir", "normalize_loop_bounds.mlir", "optimize_tensor_insert_extract_slices.mlir", @@ -95,6 +98,7 @@ iree_lit_test_suite( "convolution_match_spec.mlir", "reductions_codegen_spec.mlir", "reductions_match_spec.mlir", + "tuning_spec.mlir", ], ), cfg = "//compiler:lit.cfg.py", @@ -105,6 +109,7 @@ iree_lit_test_suite( "convolution_match_spec.mlir", "reductions_codegen_spec.mlir", "reductions_match_spec.mlir", + "tuning_spec.mlir", ], tools = [ "//tools:iree-opt", diff --git a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt index ae563af014fd..325e72828c60 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt @@ -51,6 +51,9 @@ iree_lit_test_suite( "link_tuning_specs.mlir" "lower_ukernel_to_calls.mlir" "materialize_encoding_into_nop.mlir" + "materialize_tuning_specs.mlir" + "materialize_tuning_specs_invalid_spec.mlir" + "materialize_user_config_from_tuning_spec.mlir" "materialize_user_configs.mlir" "normalize_loop_bounds.mlir" "optimize_tensor_insert_extract_slices.mlir" @@ -92,6 +95,7 @@ iree_lit_test_suite( convolution_match_spec.mlir reductions_codegen_spec.mlir reductions_match_spec.mlir + tuning_spec.mlir ) ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ### 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 new file mode 100644 index 000000000000..d28cd6874f3b --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs.mlir @@ -0,0 +1,23 @@ +// RUN: iree-opt --pass-pipeline='builtin.module(iree-codegen-materialize-tuning-specs)' \ +// RUN: --iree-codegen-tuning-spec-path=%p/tuning_spec.mlir \ +// RUN: --iree-codegen-dump-tuning-specs-to=- \ +// RUN: --mlir-disable-threading --no-implicit-module %s | FileCheck %s + +// Check that the final tuning spec is as expected. +// CHECK-LABEL: module @iree_linked_tuning_spec attributes {transform.with_named_sequence} +// CHECK-LABEL: module @user_spec attributes {transform.with_named_sequence} +// CHECK-LABEL: transform.named_sequence @hello +// CHECK-SAME: attributes {iree_codegen.tuning_spec_entrypoint} +// CHECK-LABEL: transform.named_sequence @__kernel_config +// CHECK: @user_spec::@hello + +// Check that the transform spec gets materialized as a module attribute. +// CHECK: module attributes +// CHECK-SAME: iree_codegen.tuning_spec_mlirbc = dense<{{.+}}> : vector<{{[0-9]+}}xi8> +// CHECK-LABEL: func.func @main_0 + +module { + func.func @main_0() { + return + } +} diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs_invalid_spec.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs_invalid_spec.mlir new file mode 100644 index 000000000000..f5f80e1b50ba --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs_invalid_spec.mlir @@ -0,0 +1,12 @@ +// RUN: iree-opt --pass-pipeline='builtin.module(iree-codegen-materialize-tuning-specs)' \ +// RUN: --iree-codegen-tuning-spec-path=%s --no-implicit-module --verify-diagnostics %s + +// Check that we error out on mlir inputs that are not tuning specs (e.g., the input itself). + +// expected-error@+2 {{Module without the 'transform.with_named_sequence' attribute is not a transform dialect library}} +// expected-error@+1 {{Failed to load tuning spec transform dialect library from}} +module { + func.func @main_0() { + return + } +} diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_user_config_from_tuning_spec.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_user_config_from_tuning_spec.mlir new file mode 100644 index 000000000000..08f52791de3f --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_user_config_from_tuning_spec.mlir @@ -0,0 +1,42 @@ +// RUN: iree-opt --pass-pipeline='builtin.module(builtin.module(iree-codegen-materialize-tuning-specs,iree-codegen-materialize-user-configs))' \ +// RUN: --iree-codegen-tuning-spec-path=%p/tuning_spec.mlir \ +// RUN: --mlir-disable-threading --no-implicit-module %s | FileCheck %s + +// RUN: iree-opt --pass-pipeline='builtin.module(iree-codegen-materialize-tuning-specs,builtin.module(iree-codegen-materialize-user-configs))' \ +// RUN: --iree-codegen-tuning-spec-path=%p/tuning_spec.mlir \ +// RUN: --mlir-disable-threading --no-implicit-module %s | FileCheck %s --check-prefix=PARENT + +// (1) We start by running the `Materialize Tuning Specs` pass to embed the +// transform dialect library into the module. Doing it by hand hand is not +// possible, because we serialize it as MLIR bytecode. +// +// Check that the transform spec gets executed and that it does not remain as +// a module attribute after `Materialize User Configs`. + +// CHECK-LABEL: [ IR printer: Hello Tuning Spec top-level ] +// CHECK-NEXT: func.func @main_0 +// +// CHECK-LABEL: module @parent { +// CHECK-LABEL: module @child { +// CHECK: func.func @main_0 + +// (2) Check that the transform spec gets picked up from the **parent** module. +// The tuning spec attribute should remain on the parent module as we +// (conservatively) only remove tuning spec from the module passed +// to the `materialize-user-configs` pass. + +// PARENT-LABEL: [ IR printer: Hello Tuning Spec top-level ] +// PARENT-NEXT: func.func @main_0 +// +// PARENT-LABEL: module @parent attributes { +// PARENT-SAME: iree_codegen.tuning_spec_mlirbc = dense< +// PARENT-LABEL: module @child { +// PARENT: func.func @main_0 + +module @parent { + module @child { + func.func @main_0() { + return + } + } +} diff --git a/compiler/src/iree/compiler/Codegen/Common/test/tuning_spec.mlir b/compiler/src/iree/compiler/Codegen/Common/test/tuning_spec.mlir new file mode 100644 index 000000000000..24af07386c2e --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/test/tuning_spec.mlir @@ -0,0 +1,9 @@ +// RUN: iree-opt %s + +module @user_spec attributes { transform.with_named_sequence } { + transform.named_sequence @hello(%arg0: !transform.any_op {transform.readonly}) -> () + attributes { iree_codegen.tuning_spec_entrypoint } { + transform.print {name = "Hello Tuning Spec", skip_regions} + transform.yield + } +} 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 bb2d747f0a5f..989d7bda3441 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h @@ -39,8 +39,10 @@ namespace mlir::iree_compiler { // Constant names. //===----------------------------------------------------------------------===// constexpr StringLiteral kConfigAttrName = "lowering_config"; -constexpr StringLiteral kTuningSpecAttrName = +constexpr StringLiteral kTuningSpecEntrypointAttrName = "iree_codegen.tuning_spec_entrypoint"; +constexpr StringLiteral kSerializedTuningSpecAttrName = + "iree_codegen.tuning_spec_mlirbc"; constexpr StringLiteral kKernelConfigSpecName = "__kernel_config"; //===----------------------------------------------------------------------===// diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.td b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.td index 7f47fbe10be4..9e5a408f5957 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.td +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.td @@ -23,7 +23,7 @@ def IREECodegen_Dialect : Dialect { let description = [{ This dialect is primarily meant to hold attributes that carry the state of the compilation when lowered to scalar code for an - architecture. Typically, a backend starts by analysing the entry + architecture. Typically, a backend starts by analyzing the entry point functions within the `hal.executable.variant` and deciding which compilation pipeline to chose. During this, even the values for parameters such as tile sizes, etc. are also decided. The rest diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenLibraryManager.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenLibraryManager.cpp index 93176084fd82..437a80b04cb7 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenLibraryManager.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenLibraryManager.cpp @@ -5,6 +5,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h" +#include "mlir/Dialect/Transform/IR/TransformDialect.h" #include "mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h" namespace mlir::iree_compiler::IREE::Codegen { @@ -17,25 +18,40 @@ IREECodegenDialect::getOrLoadTransformLibraryModule(std::string libraryPath) { auto loadedLibrary = libraryModules.find(libraryPath); if (loadedLibrary != libraryModules.end()) { // Check whether the library already failed to load. - if (!(loadedLibrary->second) || !(*(loadedLibrary->second))) { - return failure(); + if (ModuleOp module = loadedLibrary->second.get()) { + return module; } - return *(loadedLibrary->second); + return failure(); + } + + // We update the storage for the library regardless of whether parsing + // succeeds so that other threads don't have to retry. + OwningOpRef &parsedLibrary = libraryModules[libraryPath]; + + MLIRContext *ctx = getContext(); + if (failed(transform::detail::parseTransformModuleFromFile(ctx, libraryPath, + parsedLibrary))) { + return failure(); } - OwningOpRef mergedParsedLibraries; - if (failed(transform::detail::assembleTransformLibraryFromPaths( - getContext(), SmallVector{libraryPath}, - mergedParsedLibraries))) { - // We update the storage for the library regardless of whether parsing - // succeeds so that other threads don't have to retry. - OwningOpRef emptyLibrary; - libraryModules[libraryPath] = std::move(emptyLibrary); + if (!parsedLibrary.get()->hasAttr( + transform::TransformDialect::kWithNamedSequenceAttrName)) { + parsedLibrary->emitError() + << "Module without the '" + << transform::TransformDialect::kWithNamedSequenceAttrName + << "' attribute is not a transform dialect library"; + + // Invalidate the module stored in the library so that this does not + // succeed on a retry. + parsedLibrary = nullptr; return failure(); } - libraryModules[libraryPath] = std::move(mergedParsedLibraries); - return *libraryModules[libraryPath]; + if (!parsedLibrary->getSymName()) { + parsedLibrary->setSymName("__transform"); + } + + return parsedLibrary.get(); } } // namespace mlir::iree_compiler::IREE::Codegen