diff --git a/compiler/plugins/target/ROCM/BUILD.bazel b/compiler/plugins/target/ROCM/BUILD.bazel index 682806e23539..ef7aeec8353a 100644 --- a/compiler/plugins/target/ROCM/BUILD.bazel +++ b/compiler/plugins/target/ROCM/BUILD.bazel @@ -27,6 +27,7 @@ iree_compiler_cc_library( "ROCMTargetUtils.h", ], deps = [ + "//compiler/plugins/target/ROCM/builtins/tuning:iree_default_tuning_specs_amdgpu", "//compiler/plugins/target/ROCM/builtins/ukernel:iree_uk_amdgpu_bitcode", "//compiler/src/iree/compiler/Codegen/Common", "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect", diff --git a/compiler/plugins/target/ROCM/CMakeLists.txt b/compiler/plugins/target/ROCM/CMakeLists.txt index 69204abd3d13..a9f4dc17c89e 100644 --- a/compiler/plugins/target/ROCM/CMakeLists.txt +++ b/compiler/plugins/target/ROCM/CMakeLists.txt @@ -65,6 +65,7 @@ iree_cc_library( iree::compiler::Dialect::HAL::Utils::LLVMLinkerUtils iree::compiler::PluginAPI iree::compiler::Utils + iree::compiler::plugins::target::ROCM::builtins::tuning::iree_default_tuning_specs_amdgpu iree::compiler::plugins::target::ROCM::builtins::ukernel::iree_uk_amdgpu_bitcode iree::schemas::amdgpu_executable_def_c_fbs iree::schemas::executable_debug_info_c_fbs diff --git a/compiler/plugins/target/ROCM/ROCMTarget.cpp b/compiler/plugins/target/ROCM/ROCMTarget.cpp index 20364577f41f..851b6ed79a75 100644 --- a/compiler/plugins/target/ROCM/ROCMTarget.cpp +++ b/compiler/plugins/target/ROCM/ROCMTarget.cpp @@ -8,6 +8,7 @@ #include +#include "compiler/plugins/target/ROCM/builtins/tuning/iree_default_tuning_specs_amdgpu.h" #include "compiler/plugins/target/ROCM/builtins/ukernel/iree_uk_amdgpu_bitcode.h" #include "iree/compiler/Codegen/Common/Passes.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h" @@ -900,14 +901,9 @@ struct ROCMSession final } }; -} // namespace - -} // namespace mlir::iree_compiler::IREE::HAL - // Iterate over ukernel bitcode embedded-data files, and insert them into the // EmbeddedDataDirectory singleton. static void addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory() { - using mlir::iree_compiler::EmbeddedDataDirectory; EmbeddedDataDirectory::withGlobal([](EmbeddedDataDirectory &dir) { const iree_file_toc_t *toc = iree_uk_amdgpu_bitcode_create(); for (size_t i = 0; i < iree_uk_amdgpu_bitcode_size(); ++i) { @@ -916,11 +912,30 @@ static void addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory() { }); } +// Iterate over default tuning spec embedded-data files, and insert them into +// the EmbeddedDataDirectory singleton. +static void addAMDGPUDefaultTuningSpecsToGlobalEmbeddedDataDirectory() { + EmbeddedDataDirectory::withGlobal([](EmbeddedDataDirectory &dir) { + const iree_file_toc_t *toc = iree_default_tuning_specs_amdgpu_create(); + for (size_t i = 0, e = iree_default_tuning_specs_amdgpu_size(); i != e; + ++i) { + dir.addFile(toc[i].name, llvm::StringRef{toc[i].data, toc[i].size}); + } + }); +} + +} // namespace + +} // namespace mlir::iree_compiler::IREE::HAL + extern "C" bool iree_register_compiler_plugin_hal_target_rocm( mlir::iree_compiler::PluginRegistrar *registrar) { registrar->registerPlugin( "hal_target_rocm"); - addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory(); + mlir::iree_compiler::IREE::HAL:: + addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory(); + mlir::iree_compiler::IREE::HAL:: + addAMDGPUDefaultTuningSpecsToGlobalEmbeddedDataDirectory(); return true; } diff --git a/compiler/plugins/target/ROCM/builtins/tuning/BUILD.bazel b/compiler/plugins/target/ROCM/builtins/tuning/BUILD.bazel new file mode 100644 index 000000000000..00fc739578db --- /dev/null +++ b/compiler/plugins/target/ROCM/builtins/tuning/BUILD.bazel @@ -0,0 +1,53 @@ +# 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 + +load("//build_tools/bazel:build_defs.oss.bzl", "iree_cmake_extra_content") +load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite") +load("//build_tools/embed_data:build_defs.bzl", "iree_c_embed_data") + +package( + default_visibility = ["//visibility:public"], + features = ["layering_check"], + licenses = ["notice"], # Apache 2.0 +) + +iree_cmake_extra_content( + content = """ +if(NOT IREE_TARGET_BACKEND_ROCM) + return() +endif() +""", + inline = True, +) + +# Target archs for tuning specs. https://llvm.org/docs/AMDGPUUsage.html#processors +gpu_archs = [ + "gfx942", +] + +tuning_spec_mlir_files = [ + "iree_default_tuning_spec_%s.mlir" % gpu_arch + for gpu_arch in gpu_archs +] + +iree_c_embed_data( + name = "iree_default_tuning_specs_amdgpu", + srcs = tuning_spec_mlir_files, + c_file_output = "iree_default_tuning_specs_amdgpu.c", + flatten = True, + h_file_output = "iree_default_tuning_specs_amdgpu.h", +) + +# Verify that the tuning specs are valid. We need this here because we do not +# assemble the tuning spec mlir and invalid specs do not lead to build errors. +iree_lit_test_suite( + name = "verify_default_tuning_specs_amdgpu", + srcs = tuning_spec_mlir_files, + cfg = "//compiler:lit.cfg.py", + tools = [ + "//tools:iree-opt", + ], +) diff --git a/compiler/plugins/target/ROCM/builtins/tuning/CMakeLists.txt b/compiler/plugins/target/ROCM/builtins/tuning/CMakeLists.txt new file mode 100644 index 000000000000..a04c528edc6c --- /dev/null +++ b/compiler/plugins/target/ROCM/builtins/tuning/CMakeLists.txt @@ -0,0 +1,39 @@ +################################################################################ +# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from # +# compiler/plugins/target/ROCM/builtins/tuning/BUILD.bazel # +# # +# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary # +# CMake-only content. # +# # +# To disable autogeneration for this file entirely, delete this header. # +################################################################################ + +iree_add_all_subdirs() + +if(NOT IREE_TARGET_BACKEND_ROCM) + return() +endif() + +iree_c_embed_data( + NAME + iree_default_tuning_specs_amdgpu + SRCS + "iree_default_tuning_spec_gfx942.mlir" + C_FILE_OUTPUT + "iree_default_tuning_specs_amdgpu.c" + H_FILE_OUTPUT + "iree_default_tuning_specs_amdgpu.h" + FLATTEN + PUBLIC +) + +iree_lit_test_suite( + NAME + verify_default_tuning_specs_amdgpu + SRCS + "iree_default_tuning_spec_gfx942.mlir" + TOOLS + iree-opt +) + +### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ### 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 new file mode 100644 index 000000000000..f0a8ca58f33a --- /dev/null +++ b/compiler/plugins/target/ROCM/builtins/tuning/iree_default_tuning_spec_gfx942.mlir @@ -0,0 +1,15 @@ +// RUN: iree-opt %s + +// This is just an initial tuning spec for gfx942 and is not intended for +// production use. +// 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 } { + +transform.named_sequence @__kernel_config(%variant_op: !transform.any_op {transform.readonly}) -> () + attributes { iree_codegen.tuning_spec_entrypoint } { + transform.yield +} + +} diff --git a/compiler/plugins/target/ROCM/builtins/tuning/test/BUILD.bazel b/compiler/plugins/target/ROCM/builtins/tuning/test/BUILD.bazel new file mode 100644 index 000000000000..ab6ada2de4b1 --- /dev/null +++ b/compiler/plugins/target/ROCM/builtins/tuning/test/BUILD.bazel @@ -0,0 +1,35 @@ +# 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 + +load("//build_tools/bazel:build_defs.oss.bzl", "iree_cmake_extra_content") +load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite") + +package( + default_visibility = ["//visibility:public"], + features = ["layering_check"], + licenses = ["notice"], # Apache 2.0 +) + +iree_cmake_extra_content( + content = """ +if(NOT IREE_TARGET_BACKEND_ROCM) + return() +endif() +""", + inline = True, +) + +iree_lit_test_suite( + name = "lit", + srcs = [ + "spec_gfx942.mlir", + ], + cfg = "//compiler:lit.cfg.py", + tools = [ + "//tools:iree-opt", + "@llvm-project//llvm:FileCheck", + ], +) diff --git a/compiler/plugins/target/ROCM/builtins/tuning/test/CMakeLists.txt b/compiler/plugins/target/ROCM/builtins/tuning/test/CMakeLists.txt new file mode 100644 index 000000000000..0ed361461c04 --- /dev/null +++ b/compiler/plugins/target/ROCM/builtins/tuning/test/CMakeLists.txt @@ -0,0 +1,27 @@ +################################################################################ +# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from # +# compiler/plugins/target/ROCM/builtins/tuning/test/BUILD.bazel # +# # +# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary # +# CMake-only content. # +# # +# To disable autogeneration for this file entirely, delete this header. # +################################################################################ + +iree_add_all_subdirs() + +if(NOT IREE_TARGET_BACKEND_ROCM) + return() +endif() + +iree_lit_test_suite( + NAME + lit + SRCS + "spec_gfx942.mlir" + TOOLS + FileCheck + iree-opt +) + +### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ### diff --git a/compiler/plugins/target/ROCM/builtins/tuning/test/spec_gfx942.mlir b/compiler/plugins/target/ROCM/builtins/tuning/test/spec_gfx942.mlir new file mode 100644 index 000000000000..51c0e0109fc5 --- /dev/null +++ b/compiler/plugins/target/ROCM/builtins/tuning/test/spec_gfx942.mlir @@ -0,0 +1,28 @@ +// RUN: iree-opt --split-input-file --iree-gpu-test-target=gfx942 \ +// RUN: --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(iree-hal-configure-target-executable-variants{target=rocm})))" \ +// RUN: --iree-codegen-enable-default-tuning-specs \ +// RUN: --iree-codegen-notify-transform-strategy-application \ +// RUN: --verify-diagnostics %s | FileCheck %s + +// CHECK-LABEL: func.func @placeholder + +#pipeline_layout = #hal.pipeline.layout, + #hal.pipeline.binding, + #hal.pipeline.binding +]> +hal.executable public @main { + hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb">) { + hal.executable.export public @placeholder ordinal(0) layout(#pipeline_layout) { + ^bb0(%arg0: !hal.device): + %x, %y, %z = flow.dispatch.workgroup_count_from_slice + hal.return %x, %y, %z : index, index, index + } + builtin.module { + // expected-remark@+1 {{Applied transform configuration strategy @iree_default_tuning_spec_gfx942::@__kernel_config}} + func.func @placeholder() { + return + } + } + } +} diff --git a/compiler/plugins/target/ROCM/test/BUILD.bazel b/compiler/plugins/target/ROCM/test/BUILD.bazel index ebf4dfd7463e..f0521a0e8c50 100644 --- a/compiler/plugins/target/ROCM/test/BUILD.bazel +++ b/compiler/plugins/target/ROCM/test/BUILD.bazel @@ -15,6 +15,7 @@ package( iree_lit_test_suite( name = "lit", srcs = [ + "default_tuning_specs_amdgpu.mlir", "gpu_lower_to_ukernels.mlir", "lowering_strategy_from_tuning_spec.mlir", "ukernel_pipeline_transform.mlir", diff --git a/compiler/plugins/target/ROCM/test/CMakeLists.txt b/compiler/plugins/target/ROCM/test/CMakeLists.txt index 38158aac8c5b..36d9ba6db31d 100644 --- a/compiler/plugins/target/ROCM/test/CMakeLists.txt +++ b/compiler/plugins/target/ROCM/test/CMakeLists.txt @@ -14,6 +14,7 @@ iree_lit_test_suite( NAME lit SRCS + "default_tuning_specs_amdgpu.mlir" "gpu_lower_to_ukernels.mlir" "lowering_strategy_from_tuning_spec.mlir" "ukernel_pipeline_transform.mlir" diff --git a/compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir b/compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir new file mode 100644 index 000000000000..f049f79aede3 --- /dev/null +++ b/compiler/plugins/target/ROCM/test/default_tuning_specs_amdgpu.mlir @@ -0,0 +1,55 @@ +// RUN: iree-opt --pass-pipeline='builtin.module(iree-codegen-materialize-tuning-specs)' \ +// RUN: --iree-codegen-enable-default-tuning-specs \ +// RUN: --iree-codegen-dump-tuning-specs-to=- \ +// RUN: --iree-gpu-test-target=gfx942 --mlir-disable-threading \ +// RUN: --no-implicit-module %s | FileCheck %s --check-prefix=DEFAULT + +// RUN: iree-opt --pass-pipeline='builtin.module(iree-codegen-materialize-tuning-specs)' \ +// RUN: --iree-codegen-tuning-spec-path=%p/tuning_spec_mmt_tile_and_fuse.mlir \ +// RUN: --iree-codegen-enable-default-tuning-specs \ +// RUN: --iree-codegen-dump-tuning-specs-to=- \ +// RUN: --iree-gpu-test-target=gfx942 --mlir-disable-threading \ +// RUN: --no-implicit-module %s | FileCheck %s --check-prefix=BOTH + +// Note: This test needs to be in the plugin subdirectory because it depends +// on the default spec that's only embedded in the compiler library when the +// ROCM plugin is built. + +// ============================================================================ + +// 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: transform.named_sequence @__kernel_config +// DEFAULT-SAME: attributes {iree_codegen.tuning_spec_entrypoint} + +// Check that the default tuning spec gets materialized as a module attribute. +// DEFAULT: module attributes +// DEFAULT-SAME: iree_codegen.tuning_spec_mlirbc = dense<{{.+}}> : vector<{{[0-9]+}}xi8> +// DEFAULT-LABEL: func.func @main_0 + +// ============================================================================ + +// 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 @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: transform.named_sequence @__kernel_config +// BOTH-SAME: attributes {iree_codegen.tuning_spec_entrypoint} +// BOTH: transform.named_sequence @__kernel_config +// BOTH: @mmt_tile_and_fuse_spec_0::@main +// BOTH: @iree_default_tuning_spec_gfx942_1::@__kernel_config + +// BOTH: module attributes +// BOTH-SAME: iree_codegen.tuning_spec_mlirbc = dense<{{.+}}> : vector<{{[0-9]+}}xi8> +// BOTH-LABEL: func.func @main_0 + +module { + func.func @main_0() { + return + } +} diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp index f14aa92d5a66..7495f4354215 100644 --- a/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp @@ -8,14 +8,19 @@ #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 "iree/compiler/Codegen/Utils/GPUUtils.h" +#include "iree/compiler/Utils/EmbeddedDataDirectory.h" #include "llvm/ADT/SmallString.h" +#include "llvm/ADT/SmallVectorExtras.h" #include "llvm/Support/FileSystem.h" +#include "llvm/Support/FormatVariadic.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/AsmState.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinOps.h" @@ -41,6 +46,12 @@ llvm::cl::opt clCodegenTuningSpecPath( "dialect library)."), llvm::cl::init("")); +llvm::cl::opt clCodegenEnableDefaultTuningSpecs( + "iree-codegen-enable-default-tuning-specs", + llvm::cl::desc("Whether to enable default tuning spec transform libraries " + "shipped with the compiler"), + llvm::cl::init(false)); + llvm::cl::opt clCodegenTuningSpecDumpDir( "iree-codegen-dump-tuning-specs-to", llvm::cl::desc( @@ -50,8 +61,11 @@ llvm::cl::opt clCodegenTuningSpecDumpDir( using mlir::transform::NamedSequenceOp; -static LogicalResult dumpFinalTuningSpecToDir(ModuleOp tuningSpec, - StringRef dir) { +static LogicalResult dumpFinalTuningSpecToDir(ModuleOp tuningSpec) { + StringRef dir = clCodegenTuningSpecDumpDir; + if (dir.empty()) { + return success(); + } if (dir == "-") { tuningSpec->print(llvm::outs()); return success(); @@ -79,6 +93,55 @@ static LogicalResult dumpFinalTuningSpecToDir(ModuleOp tuningSpec, return success(); } +static FailureOr +getUserTuningSpec(ModuleOp module, IREE::Codegen::IREECodegenDialect &dialect) { + if (clCodegenTuningSpecPath.empty()) { + return failure(); + } + + FailureOr maybeTransformLibrary = + dialect.getOrLoadTransformLibraryModule(clCodegenTuningSpecPath); + if (failed(maybeTransformLibrary)) { + return module->emitError() + << "Failed to load tuning spec transform dialect library from " + << clCodegenTuningSpecPath; + } + + return *maybeTransformLibrary; +} + +static FailureOr +getDefaultTuningSpec(ModuleOp module, + IREE::Codegen::IREECodegenDialect &dialect) { + if (!clCodegenEnableDefaultTuningSpecs) { + return failure(); + } + + IREE::GPU::TargetAttr gpuTarget = getGPUTargetAttr(module); + if (!gpuTarget) { + return failure(); + } + + // Try to look up the default tuning spec for this architecture, if any. + StringRef arch = gpuTarget.getArch(); + std::string defaultTuningSpecName = + llvm::formatv("iree_default_tuning_spec_{}.mlir", arch); + std::optional defaultTuningSpecSource; + EmbeddedDataDirectory::withGlobal([&](EmbeddedDataDirectory &dir) { + defaultTuningSpecSource = dir.getFile(defaultTuningSpecName); + }); + if (!defaultTuningSpecSource) { + // Not all architectures are expected to provide default tuning specs, so + // this shouldn't be considered a hard error (but that's up to the caller). + return failure(); + } + + // Load the library through the codegen dialect so that we cache the parsed + // module. + return dialect.getOrParseTransformLibraryModule(defaultTuningSpecName, + *defaultTuningSpecSource); +} + static FailureOr serializeTuningSpecToAttr(ModuleOp tuningSpec) { std::string buffer; @@ -101,29 +164,55 @@ struct MaterializeTuningSpecsPass final } 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; + assert(dialect); + + FailureOr userTuningSpec = getUserTuningSpec(module, *dialect); + const bool hasUserTuningSpec = succeeded(userTuningSpec); + if (!hasUserTuningSpec && !clCodegenTuningSpecPath.empty()) { + // When a user spec is requested but fails to load, this is a hard + // failure. 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"); + FailureOr defaultTuningSpec = + getDefaultTuningSpec(module, *dialect); + const bool hasDefaultTuningSpec = succeeded(defaultTuningSpec); + if (!hasUserTuningSpec && !hasDefaultTuningSpec) { + // No specs available, nothing to do. + return; } - Location loc = userTuningSpec.getLoc(); + // If only the default tuning spec is available, use it directly and skip + // the linking stage. + if (!hasUserTuningSpec) { + if (failed(dumpFinalTuningSpecToDir(*defaultTuningSpec))) { + return signalPassFailure(); + } + FailureOr serializedSpec = + serializeTuningSpecToAttr(*defaultTuningSpec); + if (failed(serializedSpec)) { + module->emitError("Failed to serialize default tuning specs"); + return signalPassFailure(); + } + module->setAttr(kSerializedTuningSpecAttrName, *serializedSpec); + return; + } + + // When the user tuning spec is available, link all available libraries into + // a single module. We insert the default tuning spec last, so that any + // user-specified tuning configurations take precedence. + SmallVector allSpecs = {*userTuningSpec}; + if (hasDefaultTuningSpec) { + allSpecs.push_back(*defaultTuningSpec); + } + + Location loc = + FusedLoc::get(ctx, llvm::map_to_vector<2>(allSpecs, [](ModuleOp m) { + return m.getLoc(); + })); // This module will always be released at the end of the pass. OwningOpRef linkedTuningSpec( @@ -131,7 +220,13 @@ struct MaterializeTuningSpecsPass final linkedTuningSpec.get()->setAttr( transform::TransformDialect::kWithNamedSequenceAttrName, UnitAttr::get(ctx)); - linkedTuningSpec->insert(linkedTuningSpec->begin(), userTuningSpec.clone()); + for (auto [idx, spec] : llvm::enumerate(allSpecs)) { + ModuleOp clonedSpec = spec.clone(); + // Make sure there are no symbol name collisions. + clonedSpec.setSymName( + llvm::formatv("{}_{}", clonedSpec.getSymName().value(), idx).str()); + linkedTuningSpec->push_back(clonedSpec); + } // TODO(https://github.com/iree-org/iree/issues/19214): Add linked tuning // spec memoization to IREECodegenDialect. We should be able to provide a @@ -145,11 +240,8 @@ struct MaterializeTuningSpecsPass final return signalPassFailure(); } - if (!clCodegenTuningSpecDumpDir.empty()) { - if (failed(dumpFinalTuningSpecToDir(linkedTuningSpec.get(), - clCodegenTuningSpecDumpDir))) { - return signalPassFailure(); - } + if (failed(dumpFinalTuningSpecToDir(linkedTuningSpec.get()))) { + return signalPassFailure(); } FailureOr serializedSpec = diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel index 5644f4855ab7..5de2e3d6b95e 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel @@ -56,6 +56,7 @@ iree_lit_test_suite( "lower_ukernel_to_calls.mlir", "materialize_encoding_into_nop.mlir", "materialize_tuning_specs.mlir", + "materialize_tuning_specs_default_missing.mlir", "materialize_tuning_specs_invalid_spec.mlir", "materialize_user_config_from_tuning_spec.mlir", "materialize_user_configs.mlir", diff --git a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt index 325e72828c60..4dc774caa54a 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt @@ -52,6 +52,7 @@ iree_lit_test_suite( "lower_ukernel_to_calls.mlir" "materialize_encoding_into_nop.mlir" "materialize_tuning_specs.mlir" + "materialize_tuning_specs_default_missing.mlir" "materialize_tuning_specs_invalid_spec.mlir" "materialize_user_config_from_tuning_spec.mlir" "materialize_user_configs.mlir" 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 d28cd6874f3b..8e2a081d614b 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 @@ -3,13 +3,14 @@ // 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 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 @user_spec attributes {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} // CHECK-LABEL: transform.named_sequence @__kernel_config -// CHECK: @user_spec::@hello +// CHECK: @user_spec_0::@hello // Check that the transform spec gets materialized as a module attribute. // CHECK: module attributes diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs_default_missing.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs_default_missing.mlir new file mode 100644 index 000000000000..424dacdea361 --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs_default_missing.mlir @@ -0,0 +1,14 @@ +// RUN: iree-opt --pass-pipeline='builtin.module(iree-codegen-materialize-tuning-specs)' \ +// RUN: --iree-codegen-enable-default-tuning-specs --no-implicit-module %s \ +// RUN: | FileCheck %s + +// Check that we do not materialize any tuning spec when there's no default spec for the given +// target (since we do not set any target) +// CHECK: module { +// CHECK-LABEL: func.func @main_0 + +module { + func.func @main_0() { + return + } +} 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 9e5a408f5957..35775d08ccf2 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.td +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.td @@ -36,8 +36,23 @@ def IREECodegen_Dialect : Dialect { let extraClassDeclaration = [{ void initializeCodegenAttrs(); + /// Returns the loaded module for the `libraryPath` transform dialect + /// library MLIR input. If the module has already been loaded in the + /// past, returns the memoized module without parsing the file again. + /// This function is thead-safe. FailureOr<::mlir::ModuleOp> - getOrLoadTransformLibraryModule(std::string libraryPath); + getOrLoadTransformLibraryModule(StringRef libraryPath); + + /// Returns the parsed module for the `libraryPath` transform dialect + /// library MLIR input `libraryMLIRSource`. If the module has already + /// been parsed in the past, returns the memoized module without parsing + /// the source again. Note that this function does not access + /// `libraryPath` and only uses it as the key associated with the + /// matching parsed module. + /// This function is thead-safe. + FailureOr<::mlir::ModuleOp> + getOrParseTransformLibraryModule(StringRef libraryPath, + StringRef libraryMLIRSource); private: 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 437a80b04cb7..7380230f1d15 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenLibraryManager.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenLibraryManager.cpp @@ -7,13 +7,16 @@ #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h" #include "mlir/Dialect/Transform/IR/TransformDialect.h" #include "mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h" +#include "mlir/Parser/Parser.h" namespace mlir::iree_compiler::IREE::Codegen { -FailureOr -IREECodegenDialect::getOrLoadTransformLibraryModule(std::string libraryPath) { - // Acquire a lock on the map that will release once out of scope. - std::lock_guard guard(libraryMutex); +/// Helper function that implements the module lookup and validation. +/// The caller is responsible for acquiring the mutex for `libraryModules`. +static FailureOr getOrParseTransformLibraryModuleImpl( + StringRef libraryPath, + llvm::StringMap> &libraryModules, + llvm::function_ref &)> loadModuleFn) { auto loadedLibrary = libraryModules.find(libraryPath); if (loadedLibrary != libraryModules.end()) { @@ -28,9 +31,7 @@ IREECodegenDialect::getOrLoadTransformLibraryModule(std::string libraryPath) { // 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))) { + if (failed(loadModuleFn(parsedLibrary))) { return failure(); } @@ -54,4 +55,32 @@ IREECodegenDialect::getOrLoadTransformLibraryModule(std::string libraryPath) { return parsedLibrary.get(); } +FailureOr +IREECodegenDialect::getOrLoadTransformLibraryModule(StringRef libraryPath) { + // Acquire a lock on the map that will release once out of scope. + std::lock_guard guard(libraryMutex); + MLIRContext *ctx = getContext(); + + return getOrParseTransformLibraryModuleImpl( + libraryPath, libraryModules, [=](OwningOpRef &parsedLibrary) { + return transform::detail::parseTransformModuleFromFile(ctx, libraryPath, + parsedLibrary); + }); +} + +FailureOr IREECodegenDialect::getOrParseTransformLibraryModule( + StringRef libraryPath, StringRef libraryMLIRSource) { + // Acquire a lock on the map that will release once out of scope. + std::lock_guard guard(libraryMutex); + MLIRContext *ctx = getContext(); + + return getOrParseTransformLibraryModuleImpl( + libraryPath, libraryModules, [=](OwningOpRef &parsedLibrary) { + ParserConfig config(ctx); + parsedLibrary = + parseSourceString(libraryMLIRSource, ctx, libraryPath); + return success(*parsedLibrary != nullptr); + }); +} + } // namespace mlir::iree_compiler::IREE::Codegen