[Codegen][CPU] SelectUKernels for inner_tiled + --iree-llvmcpu-enable-llvm-ukernels. (#24569)
Wires the new C-bitcode ukernel framework into the LLVMCPU kernel-config
pipeline. Three pieces:
* `LLVMCPUSelectUKernels.{h,cpp}` — when the target config has
`llvm_ukernels = "inner_tiled"`, `selectUKernel` matches
`iree_codegen.inner_tiled` ops carrying a
`#iree_cpu.data_tiled_mma_layout` and returns a `UKernelDescriptorAttr`
whose name is the corresponding `MMAIntrinsic` enum value lowercased
with the `iree_uk_mma_` prefix (e.g.
`MMA_X86_AVX512BF16_1x16x2_F32_BF16` →
`iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16`). One ukernel per
intrinsic; `intrinsics_{m,n,k}` flow as function arguments and
specialize via bitcode LTO at the call site. Filters out `None` and the
`MMA_GENERIC_*` family — neither has a corresponding ukernel.
* `--iree-llvmcpu-enable-llvm-ukernels=<csv>` CL flag — bound to a new
`LLVMTarget.llvmUkernels` field, serialized into the target config as
`llvm_ukernels = "<csv>"`. Empty (the default) means off; passing
`inner_tiled` opts dispatches into the new framework. The flag name
leaves room for a future `mlir_ukernels`-style additional category (and
for now `inner_tiled` is also the only value). New helper
`hasLlvmUkernel(config, category)` in Codegen/Utils — distinct from the
existing `hasUkernel` to keep the legacy mmt4d ukernels and the new
LLVM-bitcode ukernels on independent enable flags.
* `KernelDispatch.cpp::setRootConfig(InnerTiledOp)` — calls
`selectUKernel` at the start and, on a match, sets the descriptor via
`setUKernelDescriptor`. Tile sizes and translation_info are unchanged:
the ukernel substitutes for the *innermost* intrinsic execution only,
the surrounding tile sizes and pipeline still hold. Bitcode resolution
is deferred to the `#iree_cpu.ukernel_provider` during
`LowerBitcodeUKernelsPass`, matching the GPU split between "match"
(here) and "attach" (provider).
Lit test (`select_ukernel.mlir`) runs
`iree-llvmcpu-select-lowering-strategy` on two `inner_tiled` ops sharing
the same MMA intrinsic but differing in target config: the one with
`llvm_ukernels = "inner_tiled"` picks up the
`iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16` descriptor; the other does
not, locking in the "off by default" property.
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/LLVMTargetOptions.cpp b/compiler/plugins/target/LLVMCPU/LLVMTargetOptions.cpp
index 21eaa4a..c58cb4e 100644
--- a/compiler/plugins/target/LLVMCPU/LLVMTargetOptions.cpp
+++ b/compiler/plugins/target/LLVMCPU/LLVMTargetOptions.cpp
@@ -196,6 +196,9 @@
if (ukernels.compare(DEFAULT_ENABLE_UKERNELS) != 0) {
addString("ukernels", ukernels);
}
+ if (llvmUkernels.compare(DEFAULT_ENABLE_LLVM_UKERNELS) != 0) {
+ addString("llvm_ukernels", llvmUkernels);
+ }
if (linkUkernelBitcode != DEFAULT_LINK_UKERNEL_BITCODE) {
addBool("link_ukernel_bitcode", linkUkernelBitcode);
}
@@ -327,6 +330,7 @@
}
target.ukernels = getString("ukernels", target.ukernels, false);
+ target.llvmUkernels = getString("llvm_ukernels", target.llvmUkernels, false);
target.linkUkernelBitcode =
getBool("link_ukernel_bitcode", target.linkUkernelBitcode);
target.enableInnerTiled =
@@ -602,6 +606,20 @@
llvm::cl::desc("Enables ukernels in the llvmcpu backend. May be "
"`default`, `none`, `all`, or a comma-separated list of "
"specific unprefixed ukernels to enable, e.g. `mmt4d`."));
+ // TODO(ukernels): this separate flag from `iree-llvmcpu-enable-ukernels`
+ // is a stopgap while the LLVM-bitcode C ukernels are the only new-style
+ // category. Once the MLIR-ukernel category exists, fold both into a single
+ // unified ukernel-enable scheme and remove this flag.
+ binder.opt<std::string>(
+ "iree-llvmcpu-enable-llvm-ukernels", enableLlvmUkernels,
+ llvm::cl::cat(category),
+ llvm::cl::desc(
+ "Enables the new-style LLVM-bitcode C ukernels under "
+ "compiler/plugins/target/LLVMCPU/builtins/ukernel/. "
+ "Comma-separated list of categories to enable, e.g. "
+ "`inner_tiled`. Empty (the default) means none. The companion "
+ "MLIR-ukernel category will share this flag in a future "
+ "extension."));
binder.opt<bool>(
"iree-llvmcpu-link-ukernel-bitcode", linkUKernelBitcode,
llvm::cl::cat(category),
@@ -657,6 +675,7 @@
target.vectorWidthInBytes = targetVectorWidthInBytes;
target.maxStackAllocSizeInBytes = targetMaxStackAllocSizeInBytes.value;
target.ukernels = enableUkernels;
+ target.llvmUkernels = enableLlvmUkernels;
target.linkUkernelBitcode = linkUKernelBitcode;
target.enableInnerTiled = enableInnerTiled;
diff --git a/compiler/plugins/target/LLVMCPU/LLVMTargetOptions.h b/compiler/plugins/target/LLVMCPU/LLVMTargetOptions.h
index d741272..cb39c88 100644
--- a/compiler/plugins/target/LLVMCPU/LLVMTargetOptions.h
+++ b/compiler/plugins/target/LLVMCPU/LLVMTargetOptions.h
@@ -47,6 +47,12 @@
static constexpr const char *DEFAULT_ENABLE_UKERNELS = "default";
static constexpr bool DEFAULT_LINK_UKERNEL_BITCODE = true;
static constexpr bool DEFAULT_ENABLE_INNER_TILED = false;
+ // The "LLVM ukernels" — i.e. the new-style C-bitcode ukernels under
+ // `compiler/plugins/target/LLVMCPU/builtins/ukernel/`. Empty default
+ // means off; values are a comma-separated list of categories to enable,
+ // e.g. `inner_tiled`. Distinct from the legacy `ukernels` field above,
+ // which controls the runtime-side `mmt4d`-style ukernels.
+ static constexpr const char *DEFAULT_ENABLE_LLVM_UKERNELS = "";
// Default initialize all fields.
LLVMTarget();
@@ -129,6 +135,11 @@
// enable, e.g. `mmt4d`.
std::string ukernels = DEFAULT_ENABLE_UKERNELS;
+ // Enables LLVM-bitcode ukernels (the new-style C ukernels under
+ // `compiler/plugins/target/LLVMCPU/builtins/ukernel/`). Comma-separated
+ // list of categories to enable, e.g. `inner_tiled`. Empty = off.
+ std::string llvmUkernels = DEFAULT_ENABLE_LLVM_UKERNELS;
+
// Link built-in ukernel bitcode libraries into generated executables.
bool linkUkernelBitcode = DEFAULT_LINK_UKERNEL_BITCODE;
@@ -205,6 +216,7 @@
llvm::cl::PowerOf2ByteSize targetMaxStackAllocSizeInBytes =
LLVMTarget::DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES;
std::string enableUkernels = LLVMTarget::DEFAULT_ENABLE_UKERNELS;
+ std::string enableLlvmUkernels = LLVMTarget::DEFAULT_ENABLE_LLVM_UKERNELS;
bool linkUKernelBitcode = LLVMTarget::DEFAULT_LINK_UKERNEL_BITCODE;
bool enableInnerTiled = LLVMTarget::DEFAULT_ENABLE_INNER_TILED;
bool listTargets; // Ignored - used with llvm::cl::ValueDisallowed.
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
index ff9a4b3..6f5095f 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
@@ -20,6 +20,7 @@
"lookup_builtin_and_byo.mlir",
"lower_inner_tiled_to_bitcode_ukernel.mlir",
"lower_inner_tiled_to_bitcode_ukernel_i8_vnni.mlir",
+ "select_ukernel.mlir",
],
include = ["*.mlir"],
),
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
index 9da2113..409e1e1 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
@@ -17,6 +17,7 @@
"lookup_builtin_and_byo.mlir"
"lower_inner_tiled_to_bitcode_ukernel.mlir"
"lower_inner_tiled_to_bitcode_ukernel_i8_vnni.mlir"
+ "select_ukernel.mlir"
TOOLS
FileCheck
iree-opt
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/select_ukernel.mlir b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/select_ukernel.mlir
new file mode 100644
index 0000000..6fc0834
--- /dev/null
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/select_ukernel.mlir
@@ -0,0 +1,74 @@
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' \
+// RUN: --split-input-file %s | FileCheck %s
+
+// Drives the LLVMCPU `LLVMCPUSelectLoweringStrategy` pass on
+// `inner_tiled` rooted dispatches carrying a CPU `DataTiledMMAAttr`. The
+// new-style C-bitcode ukernel framework is selected when the target
+// config has `llvm_ukernels = "inner_tiled"` (set by the
+// `--iree-llvmcpu-enable-llvm-ukernels=inner_tiled` CL flag at
+// `iree-compile` time; we set it directly in the test since `iree-opt`
+// doesn't run target-option serialization).
+
+// Enabled case: `selectCPUUKernel` matches the BF16 1x16x2 intrinsic and
+// annotates the op with `iree_codegen.ukernel = …`.
+#executable_target_enabled = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {
+ cpu_features = "+avx512f,+avx512bf16",
+ data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128",
+ llvm_ukernels = "inner_tiled",
+ native_vector_size = 64 : index,
+ target_triple = "x86_64-unknown-unknown-eabi-elf"
+}>
+
+func.func @bf16_inner_tiled_ukernel_enabled(
+ %lhs: tensor<2x4x1x2xbf16>, %rhs: tensor<2x4x16x2xbf16>, %acc: tensor<2x2x1x16xf32>
+ ) -> tensor<2x2x1x16xf32> attributes {hal.executable.target = #executable_target_enabled} {
+ %0 = iree_codegen.inner_tiled ins(%lhs, %rhs) outs(%acc) {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d2)>,
+ affine_map<(d0, d1, d2) -> (d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>
+ ],
+ iterator_types = [#linalg.iterator_type<parallel>,
+ #linalg.iterator_type<parallel>,
+ #linalg.iterator_type<reduction>],
+ kind = #iree_cpu.data_tiled_mma_layout<intrinsic = MMA_X86_AVX512BF16_1x16x2_F32_BF16>,
+ semantics = #iree_cpu.mma_semantics<>
+ } : tensor<2x4x1x2xbf16>, tensor<2x4x16x2xbf16> into tensor<2x2x1x16xf32>
+ return %0 : tensor<2x2x1x16xf32>
+}
+// CHECK-LABEL: func.func @bf16_inner_tiled_ukernel_enabled
+// CHECK: iree_codegen.inner_tiled
+// CHECK-SAME: iree_codegen.ukernel = #iree_codegen.ukernel_descriptor<"iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16", bitcode>
+
+// -----
+
+// Disabled case: target config lacks `llvm_ukernels`, so `selectCPUUKernel`
+// returns null and no descriptor is added — the default, keeping the new
+// framework off in plain builds.
+#executable_target_disabled = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {
+ cpu_features = "+avx512f,+avx512bf16",
+ data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128",
+ native_vector_size = 64 : index,
+ target_triple = "x86_64-unknown-unknown-eabi-elf"
+}>
+
+func.func @bf16_inner_tiled_ukernel_disabled(
+ %lhs: tensor<2x4x1x2xbf16>, %rhs: tensor<2x4x16x2xbf16>, %acc: tensor<2x2x1x16xf32>
+ ) -> tensor<2x2x1x16xf32> attributes {hal.executable.target = #executable_target_disabled} {
+ %0 = iree_codegen.inner_tiled ins(%lhs, %rhs) outs(%acc) {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d2)>,
+ affine_map<(d0, d1, d2) -> (d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>
+ ],
+ iterator_types = [#linalg.iterator_type<parallel>,
+ #linalg.iterator_type<parallel>,
+ #linalg.iterator_type<reduction>],
+ kind = #iree_cpu.data_tiled_mma_layout<intrinsic = MMA_X86_AVX512BF16_1x16x2_F32_BF16>,
+ semantics = #iree_cpu.mma_semantics<>
+ } : tensor<2x4x1x2xbf16>, tensor<2x4x16x2xbf16> into tensor<2x2x1x16xf32>
+ return %0 : tensor<2x2x1x16xf32>
+}
+// CHECK-LABEL: func.func @bf16_inner_tiled_ukernel_disabled
+// CHECK: iree_codegen.inner_tiled
+// CHECK-NOT: iree_codegen.ukernel = #iree_codegen.ukernel_descriptor
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
index 2e82813..a186358 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
@@ -68,6 +68,7 @@
"LLVMCPUMmt4dVectorLowering.cpp",
"LLVMCPUPeel.cpp",
"LLVMCPUSelectLoweringStrategy.cpp",
+ "LLVMCPUSelectUKernels.cpp",
"LLVMCPUSplitReduction.cpp",
"LLVMCPUSynchronizeSymbolVisibility.cpp",
"LLVMCPUTile.cpp",
@@ -87,6 +88,7 @@
hdrs = [
"DispatchABI.h",
"KernelDispatch.h",
+ "LLVMCPUSelectUKernels.h",
"Passes.h",
"TargetMLTransformInfo.h",
"Utils.h",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
index 21ee960..5c0ccd3 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
@@ -45,6 +45,7 @@
HDRS
"DispatchABI.h"
"KernelDispatch.h"
+ "LLVMCPUSelectUKernels.h"
"Passes.h"
"TargetMLTransformInfo.h"
"Utils.h"
@@ -64,6 +65,7 @@
"LLVMCPUMmt4dVectorLowering.cpp"
"LLVMCPUPeel.cpp"
"LLVMCPUSelectLoweringStrategy.cpp"
+ "LLVMCPUSelectUKernels.cpp"
"LLVMCPUSplitReduction.cpp"
"LLVMCPUSynchronizeSymbolVisibility.cpp"
"LLVMCPUTile.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index 6d843e0..8c4791b 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -13,6 +13,7 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
#include "iree/compiler/Codegen/Interfaces/PartitionableLoopsInterface.h"
+#include "iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.h"
#include "iree/compiler/Codegen/LLVMCPU/TargetMLTransformInfo.h"
#include "iree/compiler/Codegen/LLVMCPU/Utils.h"
#include "iree/compiler/Codegen/Utils/CPUUtils.h"
@@ -3048,6 +3049,17 @@
static LogicalResult setRootConfig(mlir::FunctionOpInterface entryPointFn,
IREE::Codegen::InnerTiledOp op) {
assert(!getLoweringConfig(op) && "expected lowering_config is not set");
+ // Annotate the op with a ukernel descriptor if `selectUKernel` matches one.
+ // This is gated by `--iree-llvmcpu-enable-llvm-ukernels=inner_tiled` on the
+ // target config; off by default. Tile sizes and translation_info below stay
+ // the same — the ukernel handles only the innermost intrinsic execution,
+ // and the surrounding tiling continues to do its job. The bitcode is
+ // resolved and attached as `hal.executable.objects` later by the
+ // `#iree_cpu.ukernel_provider` during `LowerBitcodeUKernelsPass`.
+ if (IREE::Codegen::UKernelDescriptorAttr ukernelDescriptor =
+ selectCPUUKernel(op)) {
+ setUKernelDescriptor(op, ukernelDescriptor);
+ }
SmallVector<int64_t> bounds;
op.getIterationBounds(bounds);
unsigned numLoops = bounds.size();
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.cpp
new file mode 100644
index 0000000..e60748b
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.cpp
@@ -0,0 +1,77 @@
+// Copyright 2026 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 "iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
+#include "iree/compiler/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
+
+namespace mlir::iree_compiler {
+
+namespace {
+
+// Maps an `IREE::CPU::MMAIntrinsic` to the ukernel function/file name that
+// implements it. The convention (see the README under
+// compiler/plugins/target/LLVMCPU/builtins/ukernel/) is: the enum value,
+// lowercased, with the `iree_uk_` prefix (replacing the `MMA_` in the
+// enum). Returns empty string for `None` and for the type-polymorphic
+// generic-scalar family.
+static std::string getUkernelNameForIntrinsic(IREE::CPU::MMAIntrinsic intr) {
+ StringRef enumName = IREE::CPU::stringifyMMAIntrinsic(intr);
+ // Reject the `None` marker and the generic-scalar family — neither has
+ // a corresponding ukernel.
+ if (enumName == "None" || enumName.starts_with("MMA_GENERIC_")) {
+ return {};
+ }
+ // Strip the `MMA_` prefix and lowercase the rest.
+ if (!enumName.consume_front("MMA_")) {
+ return {};
+ }
+ std::string out = "iree_uk_mma_";
+ out.reserve(out.size() + enumName.size());
+ for (char c : enumName) {
+ out.push_back(
+ static_cast<char>(std::tolower(static_cast<unsigned char>(c))));
+ }
+ return out;
+}
+
+} // namespace
+
+IREE::Codegen::UKernelDescriptorAttr selectCPUUKernel(Operation *op) {
+ auto innerTiled = dyn_cast<IREE::Codegen::InnerTiledOp>(op);
+ if (!innerTiled) {
+ return {};
+ }
+ auto mma = dyn_cast<IREE::CPU::DataTiledMMAAttr>(innerTiled.getKind());
+ if (!mma) {
+ return {};
+ }
+
+ auto execTarget = IREE::HAL::ExecutableTargetAttr::lookup(op);
+ if (!execTarget) {
+ return {};
+ }
+ DictionaryAttr config = execTarget.getConfiguration();
+ if (!hasLlvmUkernel(config, "inner_tiled")) {
+ return {};
+ }
+
+ std::string name = getUkernelNameForIntrinsic(mma.getIntrinsic());
+ if (name.empty()) {
+ return {};
+ }
+
+ MLIRContext *context = op->getContext();
+ return IREE::Codegen::UKernelDescriptorAttr::get(
+ context, StringAttr::get(context, name),
+ IREE::Codegen::UKernelArgumentKind::Bitcode);
+}
+
+} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.h b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.h
new file mode 100644
index 0000000..c486eb2
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.h
@@ -0,0 +1,30 @@
+// Copyright 2026 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
+
+#ifndef IREE_COMPILER_CODEGEN_LLVMCPU_LLVMCPUSELECTUKERNELS_H_
+#define IREE_COMPILER_CODEGEN_LLVMCPU_LLVMCPUSELECTUKERNELS_H_
+
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "mlir/IR/Operation.h"
+
+namespace mlir::iree_compiler {
+
+// Returns a `UKernelDescriptorAttr` if a built-in LLVMCPU C ukernel applies
+// to `op` AND `--iree-llvmcpu-enable-llvm-ukernels` enables the matching
+// category, otherwise returns null. The caller is expected to attach the
+// returned descriptor (e.g. via `setUKernelDescriptor`); attaching the
+// bitcode itself is deferred to the `#iree_cpu.ukernel_provider`'s
+// `createAndReplaceWithUkernelOp` at `LowerBitcodeUKernelsPass` time.
+//
+// The CPU side mirrors `compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/
+// LLVMGPUSelectUKernels.cpp` but is simpler: there's no per-target-arch
+// bitcode filename (the build features go in the filename via the
+// embedded-TOC entry name) and no shared-memory sizing.
+IREE::Codegen::UKernelDescriptorAttr selectCPUUKernel(Operation *op);
+
+} // namespace mlir::iree_compiler
+
+#endif // IREE_COMPILER_CODEGEN_LLVMCPU_LLVMCPUSELECTUKERNELS_H_
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
index 2ce7367..1e30292 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
@@ -248,6 +248,29 @@
return false;
}
+bool hasLlvmUkernel(DictionaryAttr targetConfig, StringRef ukernelCategory) {
+ auto enabled = targetConfig.getAs<StringAttr>("llvm_ukernels");
+ StringRef enabledStr = enabled ? enabled.getValue() : StringRef{};
+ // No `default`/`none`/`all` resolution: empty means off, anything else is
+ // a comma-separated list of categories. If the caller passes an empty
+ // `ukernelCategory`, the question is "is any LLVM-ukernel category
+ // enabled?".
+ if (enabledStr.empty()) {
+ return false;
+ }
+ if (ukernelCategory.empty()) {
+ return true;
+ }
+ while (!enabledStr.empty()) {
+ auto split = enabledStr.split(',');
+ if (split.first == ukernelCategory) {
+ return true;
+ }
+ enabledStr = split.second;
+ }
+ return false;
+}
+
// TODO(dcaballe): If we have to check for a significantly large number of
// features in the future, we may want to consider a persistent state to carry
// over processed HAL information or keeping the TTI instance alive and query
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.h b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
index 8a9ab91..a6a32f1 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
@@ -94,6 +94,13 @@
// is enabled at all.
bool hasUkernel(DictionaryAttr attr, StringRef ukernelName = "");
+// Like `hasUkernel`, but for the LLVM-bitcode ukernels under
+// `compiler/plugins/target/LLVMCPU/builtins/ukernel/` (gated by the
+// `iree-llvmcpu-enable-llvm-ukernels` CL flag / `llvm_ukernels` target
+// config string). Distinct from `hasUkernel` because the two families have
+// independent enable flags and overlapping category names.
+bool hasLlvmUkernel(DictionaryAttr attr, StringRef ukernelCategory = "");
+
/// Returns true if `attr` has `feature` in its CPU features.
bool hasFeature(DictionaryAttr targetConfig, StringRef feature);
