[Codegen][CPU] Register embedded ukernel bitcode + provider lookup. (#24568)
Wires the LLVMCPU plugin and the `#iree_cpu.ukernel_provider` attribute
to actually resolve ukernel bitcode:
* On plugin init, the embedded `iree_uk_cpu_bitcode` TOC is iterated and
each `.bc` is added to the global `EmbeddedDataDirectory` (mirrors
`addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory` in
`ROCMTarget.cpp`).
* `UKernelProviderAttr::createAndReplaceWithUkernelOp` now resolves the
bitcode by name: user-supplied `hal.executable.objects` ancestors first
(BYO and BYO-override paths), then the embedded directory. If found and
not already on the op, the attribute is attached to the source op; the
default `LowerBitcodeUKernelsPass` rewrite then preserves it on the
resulting `ukernel.generic`. The provider still returns `std::nullopt`
for the rewrite itself — specialized `inner_tiled` handling (scalar
`intrinsics_{m,n,k}` + outer K count + ACC inner stride) lands with
SelectUKernels in a follow-up.
* Filename match is `<ukernelName>.<features>.bc` (or bare
`<ukernelName>.bc`), matching the build rules under
`compiler/plugins/target/LLVMCPU/builtins/ukernel/`. Multi-variant
disambiguation by target feature set also lands with SelectUKernels.
BUILD.bazel: LLVMCPU plugin gains a dep on
`//compiler/plugins/target/LLVMCPU/builtins/ukernel:iree_uk_cpu_bitcode`,
and the `iree_cpu` dialect picks up HAL IR + `EmbeddedDataDirectory` so
the provider implementation can live with the attribute (consistent with
how `ROCMUkernelBitcodeSupport.cpp` sits inside the ROCM dialect).
Tests: a three-way lit test (`lookup_builtin_and_byo.mlir`) covers each
branch of the lookup:
1. Built-in only — no user objects; expect a `dense_resource<…>` with
the embedded bf16 bitcode (2732 bytes) attached.
2. BYO override — user supplies bitcode under the built-in's filename;
expect the user's bytes to flow through unchanged (and CHECK-NOT against
`dense_resource` to catch a regression where the provider silently swaps
in the embedded copy).
3. Pure BYO — user-only name with no matching built-in; expect the
user's bitcode to surface on the rewritten op.
Progress towards https://github.com/iree-org/iree/issues/24574.
Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>diff --git a/compiler/plugins/target/LLVMCPU/BUILD.bazel b/compiler/plugins/target/LLVMCPU/BUILD.bazel
index 8181ef5..c2bb725 100644
--- a/compiler/plugins/target/LLVMCPU/BUILD.bazel
+++ b/compiler/plugins/target/LLVMCPU/BUILD.bazel
@@ -32,6 +32,7 @@
":LinkerTool",
":StaticLibraryGenerator",
"//compiler/plugins/target/LLVMCPU/builtins:Builtins",
+ "//compiler/plugins/target/LLVMCPU/builtins/ukernel:iree_uk_cpu_bitcode",
"//compiler/src/iree/compiler/Codegen/Common",
"//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
diff --git a/compiler/plugins/target/LLVMCPU/CMakeLists.txt b/compiler/plugins/target/LLVMCPU/CMakeLists.txt
index dc9e1d6..345f389 100644
--- a/compiler/plugins/target/LLVMCPU/CMakeLists.txt
+++ b/compiler/plugins/target/LLVMCPU/CMakeLists.txt
@@ -69,6 +69,7 @@
iree::compiler::PluginAPI
iree::compiler::Utils
iree::compiler::plugins::target::LLVMCPU::builtins::Builtins
+ iree::compiler::plugins::target::LLVMCPU::builtins::ukernel::iree_uk_cpu_bitcode
PUBLIC
)
diff --git a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
index e5334da..ad04206 100644
--- a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
+++ b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
@@ -14,6 +14,7 @@
#include "compiler/plugins/target/LLVMCPU/builtins/Device.h"
#include "compiler/plugins/target/LLVMCPU/builtins/Musl.h"
#include "compiler/plugins/target/LLVMCPU/builtins/UKernel.h"
+#include "compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_cpu_bitcode.h"
#include "iree/compiler/Codegen/Common/Passes.h"
#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
@@ -28,6 +29,7 @@
#include "iree/compiler/Dialect/HAL/Utils/LLVMLinkerUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/PluginAPI/Client.h"
+#include "iree/compiler/Utils/EmbeddedDataDirectory.h"
#include "iree/compiler/Utils/ModuleUtils.h"
#include "llvm/Bitcode/BitcodeReader.h"
#include "llvm/Bitcode/BitcodeWriter.h"
@@ -910,6 +912,23 @@
CPUCodegenOptions codegenOptions;
};
+// Iterates over the embedded LLVMCPU C-ukernel bitcode TOC and inserts each
+// entry into the global `EmbeddedDataDirectory`. The CPU `UKernelProvider`
+// looks files up there by filename (after first checking any user-supplied
+// `hal.executable.objects` above the op), so this is what makes built-in
+// ukernels available to `iree-compile` for ops carrying a
+// `iree_codegen.ukernel` descriptor that does not come with its own bitcode.
+// Mirrors `addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory` in
+// `compiler/plugins/target/ROCM/ROCMTarget.cpp`.
+static void addLLVMCPUUkernelBitcodeToGlobalEmbeddedDataDirectory() {
+ EmbeddedDataDirectory::withGlobal([](EmbeddedDataDirectory &dir) {
+ const iree_file_toc_t *toc = iree_uk_cpu_bitcode_create();
+ for (size_t i = 0; i < iree_uk_cpu_bitcode_size(); ++i) {
+ dir.addFile(toc[i].name, llvm::StringRef{toc[i].data, toc[i].size});
+ }
+ });
+}
+
} // namespace
} // namespace mlir::iree_compiler::IREE::HAL
@@ -920,5 +939,7 @@
mlir::iree_compiler::PluginRegistrar *registrar) {
registrar->registerPlugin<mlir::iree_compiler::IREE::HAL::LLVMCPUSession>(
"hal_target_llvm_cpu");
+ mlir::iree_compiler::IREE::HAL::
+ addLLVMCPUUkernelBitcodeToGlobalEmbeddedDataDirectory();
return true;
}
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
index 5ff5b44..ff9a4b3 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
@@ -17,6 +17,7 @@
srcs = enforce_glob(
# keep sorted
[
+ "lookup_builtin_and_byo.mlir",
"lower_inner_tiled_to_bitcode_ukernel.mlir",
"lower_inner_tiled_to_bitcode_ukernel_i8_vnni.mlir",
],
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
index 1671a1e..9da2113 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
@@ -14,6 +14,7 @@
NAME
lit
SRCS
+ "lookup_builtin_and_byo.mlir"
"lower_inner_tiled_to_bitcode_ukernel.mlir"
"lower_inner_tiled_to_bitcode_ukernel_i8_vnni.mlir"
TOOLS
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/lookup_builtin_and_byo.mlir b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/lookup_builtin_and_byo.mlir
new file mode 100644
index 0000000..586c697
--- /dev/null
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/lookup_builtin_and_byo.mlir
@@ -0,0 +1,161 @@
+// RUN: iree-opt --iree-codegen-lower-bitcode-ukernels --split-input-file %s | FileCheck %s
+
+// Three-way test of the `#iree_cpu.ukernel_provider` bitcode lookup:
+//
+// (1) Built-in only. Op carries a ukernel descriptor naming a built-in but
+// no `hal.executable.objects`; the provider must find the matching
+// bitcode in the global `EmbeddedDataDirectory` (populated at LLVMCPU
+// plugin init from the embedded TOC) and attach it.
+//
+// (2) Bring-your-own override. Op carries a user-supplied
+// `hal.executable.objects` with a matching filename. The provider must
+// prefer the user's bytes over the built-in — never silently swap in
+// the embedded copy. The user bytes here (`dense<[0xAA, ...]>`) are
+// intentionally distinguishable from the real built-in bitcode.
+//
+// (3) Pure BYO with a name that does not match any built-in. The provider
+// must still attach the user's bitcode and the rewrite must succeed.
+//
+// Together these guard the three branches of `getUKernelBitcode` against
+// silent regressions.
+
+#executable_target = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {
+ iree_codegen.ukernel_provider = #iree_cpu.ukernel_provider
+}>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// (1) Built-in only: no user-supplied hal.executable.objects → expect the
+// embedded bf16 ukernel bitcode to be attached as a `dense_resource`.
+// CHECK-LABEL: @builtin_lookup
+// CHECK: %[[UK:.+]] = iree_codegen.ukernel.generic
+// CHECK-SAME: {hal.executable.objects = [
+// CHECK-SAME: #hal.executable.object<{
+// CHECK-SAME: path = "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.x86_64_avx512bf16.bc"
+// CHECK-SAME: data = dense_resource<iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.x86_64_avx512bf16.bc>
+// CHECK-SAME: iree_codegen.ukernel = #iree_codegen.ukernel_descriptor<"iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16", bitcode>}
+// CHECK-SAME: "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16"
+module attributes {hal.executable.target = #executable_target} {
+ func.func @builtin_lookup(
+ %arg0: tensor<16x2xbf16>, %arg1: tensor<16x2xbf16>) -> tensor<16x16xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = tensor.empty() : tensor<16x16xf32>
+ %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x16xf32>) -> tensor<16x16xf32>
+ %2 = linalg.generic {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"]
+ } ins(%arg0, %arg1 : tensor<16x2xbf16>, tensor<16x2xbf16>)
+ outs(%1 : tensor<16x16xf32>) attrs = {
+ iree_codegen.ukernel = #iree_codegen.ukernel_descriptor<
+ "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16", bitcode>
+ } {
+ ^bb0(%in: bf16, %in_0: bf16, %out: f32):
+ %3 = arith.extf %in : bf16 to f32
+ %4 = arith.extf %in_0 : bf16 to f32
+ %5 = arith.mulf %3, %4 : f32
+ %6 = arith.addf %out, %5 : f32
+ linalg.yield %6 : f32
+ } -> tensor<16x16xf32>
+ return %2 : tensor<16x16xf32>
+ }
+}
+
+// -----
+
+#executable_target = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {
+ iree_codegen.ukernel_provider = #iree_cpu.ukernel_provider
+}>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// (2) Override: user bytes (0xAA-pattern, all four lanes equal — prints as
+// the signless-i8 collapsed splat `dense<-86>`) must reach the
+// rewritten op unchanged. If the provider silently replaced them with
+// the embedded bitcode, the splat below would be swapped for a
+// `dense_resource<…>` and the CHECK-NOT would fail.
+// CHECK-LABEL: @byo_overrides_builtin
+// CHECK: %[[UK:.+]] = iree_codegen.ukernel.generic
+// CHECK-SAME: {hal.executable.objects = [
+// CHECK-SAME: #hal.executable.object<{
+// CHECK-SAME: path = "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.x86_64_avx512bf16.bc"
+// CHECK-SAME: data = dense<-86>
+// CHECK-NOT: dense_resource
+// CHECK-SAME: "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16"
+module attributes {hal.executable.target = #executable_target} {
+ func.func @byo_overrides_builtin(
+ %arg0: tensor<16x2xbf16>, %arg1: tensor<16x2xbf16>) -> tensor<16x16xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = tensor.empty() : tensor<16x16xf32>
+ %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x16xf32>) -> tensor<16x16xf32>
+ %2 = linalg.generic {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"]
+ } ins(%arg0, %arg1 : tensor<16x2xbf16>, tensor<16x2xbf16>)
+ outs(%1 : tensor<16x16xf32>) attrs = {
+ iree_codegen.ukernel = #iree_codegen.ukernel_descriptor<
+ "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16", bitcode>,
+ hal.executable.objects = [
+ #hal.executable.object<{
+ path = "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.x86_64_avx512bf16.bc",
+ data = dense<[170, 170, 170, 170]> : vector<4xi8>}>
+ ]
+ } {
+ ^bb0(%in: bf16, %in_0: bf16, %out: f32):
+ %3 = arith.extf %in : bf16 to f32
+ %4 = arith.extf %in_0 : bf16 to f32
+ %5 = arith.mulf %3, %4 : f32
+ %6 = arith.addf %out, %5 : f32
+ linalg.yield %6 : f32
+ } -> tensor<16x16xf32>
+ return %2 : tensor<16x16xf32>
+ }
+}
+
+// -----
+
+#executable_target = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {
+ iree_codegen.ukernel_provider = #iree_cpu.ukernel_provider
+}>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// (3) Pure BYO: user-only name with no matching built-in. The provider must
+// still surface the user's bitcode on the rewritten op.
+// CHECK-LABEL: @byo_pure_with_custom_name
+// CHECK: %[[UK:.+]] = iree_codegen.ukernel.generic
+// CHECK-SAME: {hal.executable.objects = [
+// CHECK-SAME: path = "my_custom_external_ukernel.bc"
+// CHECK-SAME: data = dense<[10, 20, 30, 40]>
+// CHECK-SAME: "my_custom_external_ukernel"
+module attributes {hal.executable.target = #executable_target} {
+ func.func @byo_pure_with_custom_name(
+ %arg0: tensor<16x2xbf16>, %arg1: tensor<16x2xbf16>) -> tensor<16x16xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = tensor.empty() : tensor<16x16xf32>
+ %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x16xf32>) -> tensor<16x16xf32>
+ %2 = linalg.generic {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"]
+ } ins(%arg0, %arg1 : tensor<16x2xbf16>, tensor<16x2xbf16>)
+ outs(%1 : tensor<16x16xf32>) attrs = {
+ iree_codegen.ukernel = #iree_codegen.ukernel_descriptor<
+ "my_custom_external_ukernel", bitcode>,
+ hal.executable.objects = [
+ #hal.executable.object<{
+ path = "my_custom_external_ukernel.bc",
+ data = dense<[10, 20, 30, 40]> : vector<4xi8>}>
+ ]
+ } {
+ ^bb0(%in: bf16, %in_0: bf16, %out: f32):
+ %3 = arith.extf %in : bf16 to f32
+ %4 = arith.extf %in_0 : bf16 to f32
+ %5 = arith.mulf %3, %4 : f32
+ %6 = arith.addf %out, %5 : f32
+ linalg.yield %6 : f32
+ } -> tensor<16x16xf32>
+ return %2 : tensor<16x16xf32>
+ }
+}
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
index efa42af..8ef2115 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
@@ -63,6 +63,9 @@
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils:MMAUtils",
"//compiler/src/iree/compiler/Dialect/Encoding/IR",
+ "//compiler/src/iree/compiler/Dialect/HAL/IR",
+ "//compiler/src/iree/compiler/Dialect/Util/IR",
+ "//compiler/src/iree/compiler/Utils",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:AffineDialect",
"@llvm-project//mlir:ArithDialect",
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
index e3fe554..6ff9835 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
@@ -45,6 +45,9 @@
iree::compiler::Codegen::Dialect::Codegen::Utils
iree::compiler::Codegen::Dialect::Codegen::Utils::MMAUtils
iree::compiler::Dialect::Encoding::IR
+ iree::compiler::Dialect::HAL::IR
+ iree::compiler::Dialect::Util::IR
+ iree::compiler::Utils
PUBLIC
)
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
index 60ed072..e1c867b 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
@@ -10,12 +10,16 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"
#include "iree/compiler/Codegen/Dialect/Codegen/Utils/MMAUtils.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
+#include "iree/compiler/Utils/EmbeddedDataDirectory.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/AsmState.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
@@ -1099,15 +1103,139 @@
// UKernelProviderAttr
//===----------------------------------------------------------------------===//
+constexpr StringLiteral kHalExecutableObjectsAttrName =
+ "hal.executable.objects";
+
+// Walks parents from `op` to return the nearest `hal.executable.objects`
+// array attribute. If a parent `hal.executable.variant` is reached, its
+// `objects` attribute is returned. Adapted from
+// `ExecutableTargetAttr::lookup`. This is how user-supplied bitcode reaches
+// the ukernel provider: a source MLIR program can attach
+// `hal.executable.objects` at any ancestor of the op (commonly the
+// module, the executable variant, or the op itself).
+static ArrayAttr lookUpExecutableObjects(Operation *op) {
+ MLIRContext *context = op->getContext();
+ auto attrId = StringAttr::get(context, kHalExecutableObjectsAttrName);
+ while (op) {
+ if (auto variantOp = dyn_cast<IREE::HAL::ExecutableVariantOp>(op)) {
+ if (std::optional<ArrayAttr> objects = variantOp.getObjects()) {
+ return *objects;
+ }
+ }
+ if (auto attr = op->getAttrOfType<ArrayAttr>(attrId)) {
+ return attr;
+ }
+ op = op->getParentOp();
+ }
+ return {};
+}
+
+// Returns true if `path` matches `<ukernelName>.<features>.bc`, the
+// filename convention used by the bitcode build rules under
+// `compiler/plugins/target/LLVMCPU/builtins/ukernel/`. The `<features>`
+// part is opaque to the lookup; this CPU framework currently picks
+// whichever built-in (or user-supplied) variant exists for the given
+// algorithm name. Multi-feature-variant disambiguation will arrive with
+// the SelectUKernels pass.
+static bool isMatchingUKernelFile(StringRef path, StringRef ukernelName) {
+ if (!path.ends_with(".bc")) {
+ return false;
+ }
+ if (!path.starts_with(ukernelName)) {
+ return false;
+ }
+ // Either exact name (`<name>.bc`) or features-suffixed
+ // (`<name>.<features>.bc`).
+ size_t after = ukernelName.size();
+ return after == path.size() - 3 ||
+ (after < path.size() - 3 && path[after] == '.');
+}
+
+// Returns an `ExecutableObjectAttr` carrying the bitcode for `ukernelName`.
+// First searches `sourceExecutableObjects` (the user-supplied
+// `hal.executable.objects` ancestor — the bring-your-own-ukernel path),
+// then falls back to the global `EmbeddedDataDirectory` (populated at
+// LLVMCPU plugin init from the embedded TOC). Returns null if no match.
+static IREE::HAL::ExecutableObjectAttr
+getUKernelBitcode(MLIRContext *context, ArrayAttr sourceExecutableObjects,
+ StringRef ukernelName) {
+ if (sourceExecutableObjects) {
+ for (Attribute a : sourceExecutableObjects) {
+ auto object = dyn_cast<IREE::HAL::ExecutableObjectAttr>(a);
+ if (object && isMatchingUKernelFile(object.getPath(), ukernelName)) {
+ return object;
+ }
+ }
+ }
+ std::optional<StringRef> matchedFilename;
+ std::optional<StringRef> matchedBytes;
+ EmbeddedDataDirectory::withGlobal([&](EmbeddedDataDirectory &dir) {
+ for (auto &entry : dir.getMap()) {
+ if (isMatchingUKernelFile(entry.getKey(), ukernelName)) {
+ matchedFilename = entry.getKey();
+ matchedBytes = entry.getValue();
+ return;
+ }
+ }
+ });
+ if (!matchedFilename) {
+ return {};
+ }
+ AsmResourceBlob blob = HeapAsmResourceBlob::allocateAndCopyInferAlign(
+ ArrayRef<char>(matchedBytes->data(), matchedBytes->size()));
+ auto bitcodeDenseAttr = DenseI8ResourceElementsAttr::get(
+ VectorType::get({static_cast<int64_t>(matchedBytes->size())},
+ IntegerType::get(context, 8)),
+ *matchedFilename, std::move(blob));
+ return IREE::HAL::ExecutableObjectAttr::get(
+ context, StringAttr::get(context, *matchedFilename),
+ cast<IREE::Util::SerializableAttrInterface>(bitcodeDenseAttr));
+}
+
std::optional<LogicalResult> UKernelProviderAttr::createAndReplaceWithUkernelOp(
RewriterBase &rewriter, StringRef name, DictionaryAttr targetConfiguration,
Operation *contextualOp, ArrayRef<Value> inputs, ArrayRef<Value> outputs,
SmallVectorImpl<Value> &otherOperands) const {
- // Fall through to the default UKernelGenericOp construction in
- // LowerBitcodeUKernelsPass. Specialized handling of `inner_tiled` ops
- // (threading `intrinsics_{m,n,k}` and the outer K count as scalar operands)
- // will be added in a follow-up commit alongside the SelectUKernels pass that
- // sets the `iree_codegen.ukernel` descriptor and attaches the bitcode.
+ // Resolve the bitcode for this ukernel: user-supplied
+ // `hal.executable.objects` first (BYO and BYO-override paths), falling
+ // back to the built-in bitcode embedded into `iree-compile` at LLVMCPU
+ // plugin init. If found, attach it as `hal.executable.objects` on the
+ // source op so the default `LowerBitcodeUKernelsPass` rewrite preserves
+ // it on the resulting `ukernel.generic`. If the source op already
+ // carries a matching object, this is a no-op (the BYO case).
+ //
+ // We return `std::nullopt` to let the pass run its default
+ // `UKernelGenericOp` construction. Specialized `inner_tiled` handling
+ // (threading `intrinsics_{m,n,k}` and the outer K count as scalar
+ // operands, and computing the ACC inner stride) lands in a follow-up
+ // alongside SelectUKernels.
+ MLIRContext *context = rewriter.getContext();
+ ArrayAttr sourceExecutableObjects = lookUpExecutableObjects(contextualOp);
+ IREE::HAL::ExecutableObjectAttr bitcodeObject =
+ getUKernelBitcode(context, sourceExecutableObjects, name);
+ if (bitcodeObject) {
+ auto attrId = StringAttr::get(context, kHalExecutableObjectsAttrName);
+ ArrayAttr existing = contextualOp->getAttrOfType<ArrayAttr>(attrId);
+ bool alreadyOnOp = false;
+ if (existing) {
+ for (Attribute a : existing) {
+ if (a == bitcodeObject) {
+ alreadyOnOp = true;
+ break;
+ }
+ }
+ }
+ if (!alreadyOnOp) {
+ SmallVector<Attribute> objects;
+ if (existing) {
+ objects.append(existing.begin(), existing.end());
+ }
+ objects.push_back(bitcodeObject);
+ rewriter.modifyOpInPlace(contextualOp, [&] {
+ contextualOp->setAttr(attrId, ArrayAttr::get(context, objects));
+ });
+ }
+ }
return std::nullopt;
}
