[Codegen][Tuner] Add support for default tuning specs (#19394)
These default specs are target architecture-specific and will be shipped
with the compiler.
* Default specs belong to target plugins and get embedded in
`libIREECompiler.so`, just like ukernels.
* Plugins then register their default tuning specs with the default
embedded directory.
* We store them as mlir text. We can't easily assemble them as mlir
bytecode without taking a circular dependency on iree-opt. We can
revisit this in the future and add a new tool `iree-as` that will only
link with dialects.
* After the initial loading, we cache the default specs in the IREE
codegen dialect transform library manager.
* Add a placeholder spec for gfx942.
* Document and test the inclusion order. User specs come before default
specs.
Issue: https://github.com/iree-org/iree/issues/19214
diff --git a/compiler/plugins/target/ROCM/BUILD.bazel b/compiler/plugins/target/ROCM/BUILD.bazel
index 682806e..ef7aeec 100644
--- a/compiler/plugins/target/ROCM/BUILD.bazel
+++ b/compiler/plugins/target/ROCM/BUILD.bazel
@@ -27,6 +27,7 @@
"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 69204ab..a9f4dc1 100644
--- a/compiler/plugins/target/ROCM/CMakeLists.txt
+++ b/compiler/plugins/target/ROCM/CMakeLists.txt
@@ -65,6 +65,7 @@
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 2036457..851b6ed 100644
--- a/compiler/plugins/target/ROCM/ROCMTarget.cpp
+++ b/compiler/plugins/target/ROCM/ROCMTarget.cpp
@@ -8,6 +8,7 @@
#include <cstdint>
+#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 @@
}
};
-} // 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 @@
});
}
+// 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<mlir::iree_compiler::IREE::HAL::ROCMSession>(
"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 0000000..00fc739
--- /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 0000000..a04c528
--- /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 0000000..f0a8ca5
--- /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 0000000..ab6ada2
--- /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 0000000..0ed3614
--- /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 0000000..51c0e01
--- /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<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+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 ebf4dfd..f0521a0 100644
--- a/compiler/plugins/target/ROCM/test/BUILD.bazel
+++ b/compiler/plugins/target/ROCM/test/BUILD.bazel
@@ -15,6 +15,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/CMakeLists.txt b/compiler/plugins/target/ROCM/test/CMakeLists.txt
index 38158aa..36d9ba6 100644
--- a/compiler/plugins/target/ROCM/test/CMakeLists.txt
+++ b/compiler/plugins/target/ROCM/test/CMakeLists.txt
@@ -14,6 +14,7 @@
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 0000000..f049f79
--- /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 f14aa92..7495f43 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 @@
"dialect library)."),
llvm::cl::init(""));
+llvm::cl::opt<bool> 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<std::string> clCodegenTuningSpecDumpDir(
"iree-codegen-dump-tuning-specs-to",
llvm::cl::desc(
@@ -50,8 +61,11 @@
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 @@
return success();
}
+static FailureOr<ModuleOp>
+getUserTuningSpec(ModuleOp module, IREE::Codegen::IREECodegenDialect &dialect) {
+ if (clCodegenTuningSpecPath.empty()) {
+ return failure();
+ }
+
+ FailureOr<ModuleOp> maybeTransformLibrary =
+ dialect.getOrLoadTransformLibraryModule(clCodegenTuningSpecPath);
+ if (failed(maybeTransformLibrary)) {
+ return module->emitError()
+ << "Failed to load tuning spec transform dialect library from "
+ << clCodegenTuningSpecPath;
+ }
+
+ return *maybeTransformLibrary;
+}
+
+static FailureOr<ModuleOp>
+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<StringRef> 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<DenseElementsAttr>
serializeTuningSpecToAttr(ModuleOp tuningSpec) {
std::string buffer;
@@ -101,29 +164,55 @@
}
void runOnOperation() override {
- if (clCodegenTuningSpecPath.empty()) {
- return;
- }
-
ModuleOp module = getOperation();
MLIRContext *ctx = &getContext();
auto dialect = ctx->getOrLoadDialect<IREE::Codegen::IREECodegenDialect>();
- auto maybeTransformLibrary =
- dialect->getOrLoadTransformLibraryModule(clCodegenTuningSpecPath);
- if (failed(maybeTransformLibrary)) {
- module->emitError()
- << "Failed to load tuning spec transform dialect library from "
- << clCodegenTuningSpecPath;
+ assert(dialect);
+
+ FailureOr<ModuleOp> 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<ModuleOp> 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<DenseElementsAttr> 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<ModuleOp, 2> 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<ModuleOp> linkedTuningSpec(
@@ -131,7 +220,13 @@
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 @@
return signalPassFailure();
}
- if (!clCodegenTuningSpecDumpDir.empty()) {
- if (failed(dumpFinalTuningSpecToDir(linkedTuningSpec.get(),
- clCodegenTuningSpecDumpDir))) {
- return signalPassFailure();
- }
+ if (failed(dumpFinalTuningSpecToDir(linkedTuningSpec.get()))) {
+ return signalPassFailure();
}
FailureOr<DenseElementsAttr> serializedSpec =
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
index 5644f48..5de2e3d 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 @@
"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 325e728..4dc774c 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 @@
"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 d28cd68..8e2a081 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 0000000..424dacd
--- /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 9e5a408..35775d0 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 @@
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 437a80b..7380230 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<ModuleOp>
-IREECodegenDialect::getOrLoadTransformLibraryModule(std::string libraryPath) {
- // Acquire a lock on the map that will release once out of scope.
- std::lock_guard<std::mutex> guard(libraryMutex);
+/// Helper function that implements the module lookup and validation.
+/// The caller is responsible for acquiring the mutex for `libraryModules`.
+static FailureOr<ModuleOp> getOrParseTransformLibraryModuleImpl(
+ StringRef libraryPath,
+ llvm::StringMap<OwningOpRef<ModuleOp>> &libraryModules,
+ llvm::function_ref<LogicalResult(OwningOpRef<ModuleOp> &)> loadModuleFn) {
auto loadedLibrary = libraryModules.find(libraryPath);
if (loadedLibrary != libraryModules.end()) {
@@ -28,9 +31,7 @@
// succeeds so that other threads don't have to retry.
OwningOpRef<ModuleOp> &parsedLibrary = libraryModules[libraryPath];
- MLIRContext *ctx = getContext();
- if (failed(transform::detail::parseTransformModuleFromFile(ctx, libraryPath,
- parsedLibrary))) {
+ if (failed(loadModuleFn(parsedLibrary))) {
return failure();
}
@@ -54,4 +55,32 @@
return parsedLibrary.get();
}
+FailureOr<ModuleOp>
+IREECodegenDialect::getOrLoadTransformLibraryModule(StringRef libraryPath) {
+ // Acquire a lock on the map that will release once out of scope.
+ std::lock_guard<std::mutex> guard(libraryMutex);
+ MLIRContext *ctx = getContext();
+
+ return getOrParseTransformLibraryModuleImpl(
+ libraryPath, libraryModules, [=](OwningOpRef<ModuleOp> &parsedLibrary) {
+ return transform::detail::parseTransformModuleFromFile(ctx, libraryPath,
+ parsedLibrary);
+ });
+}
+
+FailureOr<ModuleOp> IREECodegenDialect::getOrParseTransformLibraryModule(
+ StringRef libraryPath, StringRef libraryMLIRSource) {
+ // Acquire a lock on the map that will release once out of scope.
+ std::lock_guard<std::mutex> guard(libraryMutex);
+ MLIRContext *ctx = getContext();
+
+ return getOrParseTransformLibraryModuleImpl(
+ libraryPath, libraryModules, [=](OwningOpRef<ModuleOp> &parsedLibrary) {
+ ParserConfig config(ctx);
+ parsedLibrary =
+ parseSourceString<ModuleOp>(libraryMLIRSource, ctx, libraryPath);
+ return success(*parsedLibrary != nullptr);
+ });
+}
+
} // namespace mlir::iree_compiler::IREE::Codegen
