[Codegen][CPU] End-to-end-through-codegen lit test for ukernel framework. (#24571)
Closes the loop on the new C-bitcode ukernel framework with an
`iree-opt` lit test that drives the full LLVMCPU codegen pipeline on an
`inner_tiled` rooted dispatch wrapped in the standard
`hal.executable.variant` + binding-subspan structure, and asserts the
whole chain end-to-end:
- Bitcode reaches the dispatch's `hal.executable.variant` as a
first-class `objects(...)` attribute (the long-lived home).
- `llvm.func @iree_uk_mma_…` is declared with `hal.import.bitcode =
true`, the flag that tells `RewriteExternCallOpToDynamicImportCallOp` to
skip the runtime-import-table indirection.
- A direct `llvm.call @iree_uk_mma_…` replaces the original
`inner_tiled` — no `__import_ordinal_iree_uk_mma_…` (which would mean we
accidentally re-used the legacy runtime-resolved path), no leftover
`inner_tiled`, no leftover `ukernel.generic`.
To make that chain work, two pieces had to land alongside the test:
* `UKernelProviderAttr::createAndReplaceWithUkernelOp` now actually
builds the `UKernelGenericOp` (returning `success`, not `nullopt`),
setting `fn_def_attrs = {hal.import.bitcode = true}` on it. That flag
propagates onto the `func.func` declaration that
`LowerUKernelOpsToCalls` synthesizes, and through it onto the
`llvm.func` that `ConvertToLLVM` produces. Without this, the
"self-contained bitcode" property of the new framework would not hold —
every external call would be routed through the runtime import table
just like the legacy mmt4d ukernels.
* Bitcode lookup-and-attach is factored into a public helper
`IREE::CPU::attachUKernelBitcodeOnOp`, called both from
`LLVMCPUSelectUKernels` (at kernel-config time, mirroring GPU's
`ensureUKernelBitcodeAndFinalizeConfig`) and from the provider itself
(as a backstop for the BYO / hand-written-IR paths). It attaches the
bitcode both as a `hal.executable.objects` discardable attribute on the
op (the GPU-aligned shape) *and* directly on the enclosing
`hal.executable.variant`'s `objects` operand attribute (the durable
home, sidestepping passes that strip discardable attrs through
`ConvertToLLVM`).
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/builtins/ukernel/test/BUILD.bazel b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
index 6f5095f..a710d6a 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
[
+ "e2e_inner_tiled_pipeline.mlir",
"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 409e1e1..63f4c15 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
+ "e2e_inner_tiled_pipeline.mlir"
"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/e2e_inner_tiled_pipeline.mlir b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/e2e_inner_tiled_pipeline.mlir
new file mode 100644
index 0000000..9944696
--- /dev/null
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/e2e_inner_tiled_pipeline.mlir
@@ -0,0 +1,114 @@
+// RUN: iree-opt --split-input-file \
+// RUN: --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(iree-codegen-configuration-preprocessing-pipeline, builtin.module(iree-codegen-llvmcpu-configuration-pipeline, iree-codegen-llvmcpu-lowering-pipeline), iree-codegen-translation-postprocessing-pipeline)))" %s \
+// RUN: | FileCheck %s
+
+// End-to-end-through-codegen test for the new C-bitcode ukernel framework.
+// Drives the LLVMCPU configuration + lowering pipelines on an
+// `iree_codegen.inner_tiled` rooted dispatch wrapped in the standard
+// `hal.executable.variant` + binding-subspan structure (the same wrapper
+// shape `iree-compile` produces after dispatch creation). Asserts the
+// whole chain:
+//
+// - SelectUKernels saw the `llvm_ukernels = "inner_tiled"` flag in the
+// target config and the matching MMA intrinsic, so it annotated the
+// `inner_tiled` and attached the matching bitcode to the parent
+// `hal.executable.variant`'s `objects(...)` (the long-lived home; not
+// a discardable attr that intermediate passes might strip).
+// - LowerBitcodeUKernels rewrote the `inner_tiled` into a
+// `ukernel.generic` with `fn_def_attrs = {hal.import.bitcode = true}`.
+// - LowerUKernelOpsToCalls turned that into a `func.call` against a
+// `func.func` declaration carrying `hal.import.bitcode`.
+// - Final LLVM lowering emits `llvm.func @iree_uk_mma_…` with
+// `hal.import.bitcode = true` and a direct `llvm.call` to it. The
+// `hal.import.bitcode` flag is what makes
+// `RewriteExternCallOpToDynamicImportCallOp` in `ConvertToLLVM` skip
+// the runtime-import-table indirection it applies to every other
+// external call — letting our call resolve directly against the
+// linked bitcode at LLVM optimization time.
+
+#executable_target = #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",
+ iree_codegen.ukernel_provider = #iree_cpu.ukernel_provider,
+ llvm_ukernels = "inner_tiled",
+ native_vector_size = 64 : index,
+ target_triple = "x86_64-unknown-unknown-eabi-elf"
+}>
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer, ReadOnly>,
+ #hal.pipeline.binding<storage_buffer, ReadOnly>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+
+hal.executable private @bf16_inner_tiled_ukernel {
+ hal.executable.variant public @variant target(#executable_target) {
+ hal.executable.export public @bf16_inner_tiled_ukernel ordinal(0) layout(#pipeline_layout)
+ count(%device: !hal.device) -> (index, index, index) {
+ %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice()
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @bf16_inner_tiled_ukernel() {
+ %c0 = arith.constant 0 : index
+ %lhs = hal.interface.binding.subspan layout(#pipeline_layout) binding(0)
+ alignment(64) offset(%c0) flags(ReadOnly)
+ : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2x4x1x2xbf16>>
+ %rhs = hal.interface.binding.subspan layout(#pipeline_layout) binding(1)
+ alignment(64) offset(%c0) flags(ReadOnly)
+ : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2x4x16x2xbf16>>
+ %out = hal.interface.binding.subspan layout(#pipeline_layout) binding(2)
+ alignment(64) offset(%c0)
+ : !iree_tensor_ext.dispatch.tensor<readwrite:tensor<2x2x1x16xf32>>
+ %lhs_t = iree_tensor_ext.dispatch.tensor.load %lhs,
+ offsets = [0, 0, 0, 0], sizes = [2, 4, 1, 2], strides = [1, 1, 1, 1]
+ : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2x4x1x2xbf16>> -> tensor<2x4x1x2xbf16>
+ %rhs_t = iree_tensor_ext.dispatch.tensor.load %rhs,
+ offsets = [0, 0, 0, 0], sizes = [2, 4, 16, 2], strides = [1, 1, 1, 1]
+ : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2x4x16x2xbf16>> -> tensor<2x4x16x2xbf16>
+ %acc_t = iree_tensor_ext.dispatch.tensor.load %out,
+ offsets = [0, 0, 0, 0], sizes = [2, 2, 1, 16], strides = [1, 1, 1, 1]
+ : !iree_tensor_ext.dispatch.tensor<readwrite:tensor<2x2x1x16xf32>> -> tensor<2x2x1x16xf32>
+ %res = iree_codegen.inner_tiled ins(%lhs_t, %rhs_t) outs(%acc_t) {
+ 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>
+ iree_tensor_ext.dispatch.tensor.store %res, %out,
+ offsets = [0, 0, 0, 0], sizes = [2, 2, 1, 16], strides = [1, 1, 1, 1]
+ : tensor<2x2x1x16xf32> -> !iree_tensor_ext.dispatch.tensor<readwrite:tensor<2x2x1x16xf32>>
+ return
+ }
+ }
+ }
+}
+
+// CHECK-LABEL: hal.executable private @bf16_inner_tiled_ukernel
+//
+// The bitcode for the matched ukernel is attached as `objects(...)` on
+// the variant — the self-contained-IR property of the new framework.
+// CHECK: hal.executable.variant public @variant
+// CHECK-SAME: objects([
+// CHECK-SAME: path = "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.x86_64_avx512bf16.bc"
+//
+// The ukernel function is declared as `llvm.func` carrying
+// `hal.import.bitcode = true` — the flag that tells `ConvertToLLVM` to
+// emit a direct call instead of routing through the runtime import table.
+// CHECK: llvm.func @iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16
+// CHECK-SAME: attributes {hal.import.bitcode = true
+//
+// And the dispatch function emits a direct `llvm.call` to that symbol.
+// No `inner_tiled`, no `ukernel.generic`, and no `__import_ordinal_`
+// indirection (which would mean we accidentally re-used the legacy
+// runtime-resolved-import path).
+// CHECK: llvm.call @iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16
+// CHECK-NOT: iree_codegen.inner_tiled
+// CHECK-NOT: iree_codegen.ukernel.generic
+// CHECK-NOT: __import_ordinal_iree_uk_mma
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 e1c867b..2eb1f43 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
@@ -18,6 +18,8 @@
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/Attributes.h"
@@ -1192,51 +1194,204 @@
cast<IREE::Util::SerializableAttrInterface>(bitcodeDenseAttr));
}
+// Idempotently appends `bitcodeObject` to the `hal.executable.objects`
+// array on `op`. Returns true if the array already contained an
+// equivalent entry (i.e. nothing was changed).
+static bool addBitcodeObjectIfMissing(Operation *op, StringAttr attrId,
+ Attribute bitcodeObject) {
+ ArrayAttr existing = op->getAttrOfType<ArrayAttr>(attrId);
+ if (existing) {
+ for (Attribute a : existing) {
+ if (a == bitcodeObject) {
+ return true;
+ }
+ }
+ }
+ SmallVector<Attribute> objects;
+ if (existing) {
+ objects.append(existing.begin(), existing.end());
+ }
+ objects.push_back(bitcodeObject);
+ op->setAttr(attrId, ArrayAttr::get(op->getContext(), objects));
+ return false;
+}
+
+bool attachUKernelBitcodeOnOp(Operation *op, StringRef name) {
+ MLIRContext *context = op->getContext();
+ ArrayAttr sourceExecutableObjects = lookUpExecutableObjects(op);
+ IREE::HAL::ExecutableObjectAttr bitcodeObject =
+ getUKernelBitcode(context, sourceExecutableObjects, name);
+ if (!bitcodeObject) {
+ return false;
+ }
+ auto attrId = StringAttr::get(context, kHalExecutableObjectsAttrName);
+ addBitcodeObjectIfMissing(op, attrId, bitcodeObject);
+
+ // Also attach to the enclosing `hal.executable.variant`'s `objects` (a
+ // first-class operand attribute, not a discardable attr). This is the
+ // long-lived home for the bitcode: it survives codegen passes that
+ // would otherwise strip discardable attributes off intermediate ops,
+ // and matches the destination that `iree-hal-hoist-executable-objects`
+ // would eventually hoist to. Doing it eagerly here just shortens the
+ // distance the bitcode has to travel.
+ Operation *parent = op->getParentOp();
+ while (parent) {
+ if (auto variantOp = dyn_cast<IREE::HAL::ExecutableVariantOp>(parent)) {
+ SmallVector<Attribute> variantObjects;
+ if (std::optional<ArrayAttr> existing = variantOp.getObjects()) {
+ for (Attribute a : *existing) {
+ if (a == bitcodeObject) {
+ return true;
+ }
+ variantObjects.push_back(a);
+ }
+ }
+ variantObjects.push_back(bitcodeObject);
+ variantOp.setObjectsAttr(ArrayAttr::get(context, variantObjects));
+ return true;
+ }
+ parent = parent->getParentOp();
+ }
+ return true;
+}
+
+// Returns the index, in the ACC operand's shape, of the innermost
+// CrossIntrinsic dimension (the N cross-intrinsic dim for our layouts), or
+// nullopt if it is dynamic / absent. The ukernel needs this dim's stride to
+// address each unrolled intrinsic's ACC fragment.
+//
+// Example: with `MMA_X86_AVX512BF16_1x16x2` (each intrinsic produces a 1x16
+// f32 fragment) and intrinsics_m = intrinsics_n = 2, the ACC tile holds a 2x2
+// grid of such fragments. The two CrossIntrinsic dims are that grid's M and N;
+// the innermost is N, so this returns the index of the N-grid dim in the ACC
+// result shape, whose stride is the element distance from fragment (m, n) to
+// (m, n+1).
+static std::optional<unsigned>
+getAccInnermostCrossIntrinsicDim(IREE::Codegen::InnerTiledOp op,
+ DataTiledMMAAttr mma) {
+ auto outputType = dyn_cast<ShapedType>(op.getResultTypes()[0]);
+ if (!outputType) {
+ return std::nullopt;
+ }
+ Codegen::TileSwizzle accSwizzle = getSwizzle(mma, /*operandIdx=*/2);
+ SmallVector<Codegen::TileSwizzle::Dim> swizzleDims;
+ for (const Codegen::TileSwizzle::ExpandShapeDimVectorType &group :
+ accSwizzle.expandShape()) {
+ swizzleDims.append(group.begin(), group.end());
+ }
+ applyPermutationToVector(swizzleDims, accSwizzle.permutation());
+ int rankDiff = outputType.getRank() - static_cast<int>(swizzleDims.size());
+ auto crossIntrinsic = Codegen::TileSwizzle::Dim::Kind::CrossIntrinsic;
+ for (size_t i = swizzleDims.size(); i-- > 0;) {
+ if (swizzleDims[i].kind() != crossIntrinsic) {
+ continue;
+ }
+ int outputIdx = i + rankDiff;
+ if (outputType.isDynamicDim(outputIdx)) {
+ return std::nullopt;
+ }
+ return outputIdx;
+ }
+ // No CrossIntrinsic dims (intrinsics_m == intrinsics_n == 1): the single
+ // fragment sits at the start of the inner tile.
+ if (!swizzleDims.empty()) {
+ return rankDiff;
+ }
+ return std::nullopt;
+}
+
+// Rewrites an `inner_tiled` carrying a CPU `DataTiledMMAAttr` to a
+// `ukernel.generic`, threading the data-tiled-MMA scalar parameters as
+// operands so the ukernel can loop over arbitrary `intrinsics_{m,n,k}`:
+// ins(lhs, rhs) outs(acc) (k_outer, intrinsics_m, intrinsics_n, intrinsics_k)
+// plus a `strided_dims` entry giving the ACC's innermost cross-intrinsic
+// stride.
+static LogicalResult
+handleInnerTiledMmaUkernel(RewriterBase &rewriter, StringRef name,
+ IREE::Codegen::InnerTiledOp op, DataTiledMMAAttr mma,
+ ArrayRef<Value> inputs, ArrayRef<Value> outputs,
+ DictionaryAttr fnDefAttrs) {
+ std::optional<unsigned> accInnerDim =
+ getAccInnermostCrossIntrinsicDim(op, mma);
+ if (!accInnerDim) {
+ return rewriter.notifyMatchFailure(
+ op, "ACC innermost cross-intrinsic dim is dynamic or absent");
+ }
+ Location loc = op.getLoc();
+ Type i32 = rewriter.getI32Type();
+ auto constI32 = [&](int64_t v) {
+ return arith::ConstantIntOp::create(rewriter, loc, i32, v);
+ };
+ // Outer-K tile count: the LHS operand's outer dims are (m_outer, k_outer),
+ // so dim 1 is the K-tile count the ukernel loops over.
+ Value kOuter = arith::IndexCastOp::create(
+ rewriter, loc, i32,
+ tensor::DimOp::create(rewriter, loc, op.getInputs()[0], 1));
+ // `strided_dims` is the `ukernel.generic` ABI knob for which dims' strides
+ // are passed to the C function: a list per shaped operand (here LHS, RHS,
+ // ACC), each naming the dims whose stride follows that operand's
+ // `(base, offset)` in the call. An empty list passes `base, offset` only; a
+ // null attribute (the default, used by simpler ukernels) passes all strides.
+ // Here only ACC needs one — the innermost cross-intrinsic (N) dim — so the
+ // ukernel can address each unrolled intrinsic's ACC fragment. LHS/RHS are
+ // contiguous in the data-tiled layout, so they take no stride argument.
+ SmallVector<SmallVector<int64_t>> stridedDims(3, {});
+ stridedDims[2].push_back(*accInnerDim);
+ DictionaryAttr discardableAttrs = op->getDiscardableAttrDictionary();
+ auto newOp = rewriter.replaceOpWithNewOp<IREE::Codegen::UKernelGenericOp>(
+ op, op.getOutputs().getTypes(), name, inputs, outputs,
+ ValueRange{kOuter, constI32(mma.getIntrinsicsM()),
+ constI32(mma.getIntrinsicsN()),
+ constI32(mma.getIntrinsicsK())},
+ fnDefAttrs, stridedDims);
+ newOp->setDiscardableAttrs(discardableAttrs);
+ return success();
+}
+
std::optional<LogicalResult> UKernelProviderAttr::createAndReplaceWithUkernelOp(
RewriterBase &rewriter, StringRef name, DictionaryAttr targetConfiguration,
Operation *contextualOp, ArrayRef<Value> inputs, ArrayRef<Value> outputs,
SmallVectorImpl<Value> &otherOperands) const {
- // 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.
+ // Idempotent: in the normal flow `LLVMCPUSelectUKernels` attached the
+ // bitcode at kernel-config time, so this is a no-op. Still defensive
+ // for tests / future paths that bypass SelectUKernels.
+ rewriter.modifyOpInPlace(
+ contextualOp, [&] { attachUKernelBitcodeOnOp(contextualOp, name); });
+
+ // We build the `ukernel.generic` ourselves (rather than returning nullopt
+ // and letting the default fallback in `LowerBitcodeUKernelsPass` handle it)
+ // for two reasons: (1) to set `fn_def_attrs = {hal.import.bitcode = true}`,
+ // and (2) for `inner_tiled` ops, to thread the `DataTiledMMAAttr` scalar
+ // parameters as operands. The `hal.import.bitcode` flag propagates onto the
+ // `func.func` declaration that `LowerUKernelOpsToCalls` synthesizes;
+ // `RewriteExternCallOpToDynamicImportCallOp` (in LLVMCPU's ConvertToLLVM)
+ // keys on it to *skip* the import-table indirection it otherwise applies to
+ // every external call, letting the call resolve directly against the linked
+ // bitcode at LLVM optimization time. Without it the new framework would
+ // accidentally re-use the legacy runtime-resolved-import path.
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));
- });
+ auto fnDefAttrs = DictionaryAttr::get(
+ context, {{rewriter.getStringAttr("hal.import.bitcode"),
+ rewriter.getBoolAttr(true)}});
+
+ // `inner_tiled` with a CPU `DataTiledMMAAttr`: thread the unrolling factors,
+ // outer-K count and ACC stride so the ukernel can loop over the intrinsics.
+ if (auto innerTiled = dyn_cast<IREE::Codegen::InnerTiledOp>(contextualOp)) {
+ if (auto mma = dyn_cast<DataTiledMMAAttr>(innerTiled.getKind())) {
+ return handleInnerTiledMmaUkernel(rewriter, name, innerTiled, mma, inputs,
+ outputs, fnDefAttrs);
}
}
- return std::nullopt;
+
+ // Any other op: a plain `ukernel.generic` with no extra operands.
+ DictionaryAttr discardableAttrs =
+ contextualOp->getDiscardableAttrDictionary();
+ auto newOp = rewriter.replaceOpWithNewOp<IREE::Codegen::UKernelGenericOp>(
+ contextualOp, contextualOp->getResults().getTypes(), name, inputs,
+ outputs, otherOperands, fnDefAttrs,
+ /*num_strided_outer_dims=*/0);
+ newOp->setDiscardableAttrs(discardableAttrs);
+ return success();
}
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
index e78face..c0a1425 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
@@ -98,6 +98,19 @@
std::optional<std::tuple<int64_t, int64_t, int64_t>>
getRowMajorTilesMNKShape(MMAIntrinsic intrinsic);
+// Idempotently attaches the bitcode for ukernel `name` as
+// `hal.executable.objects` on `op`, looking it up first in any
+// user-supplied `hal.executable.objects` ancestor, then falling back to
+// the global `EmbeddedDataDirectory` (populated at LLVMCPU plugin init).
+// No-op if `op` already carries a matching entry. Returns true if a matching
+// bitcode was found (and is now attached), false if none exists for `name` —
+// callers use the bool to decide whether a ukernel is actually available.
+// Used by both `LLVMCPUSelectUKernels` (at kernel-config time, so the bitcode
+// is on the op for the rest of codegen) and by
+// `UKernelProviderAttr::createAndReplaceWithUkernelOp` (as a backstop for
+// tests / future paths that bypass SelectUKernels).
+bool attachUKernelBitcodeOnOp(Operation *op, StringRef name);
+
} // namespace mlir::iree_compiler::IREE::CPU
// clang-format on
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.cpp
index e60748b..bb97786 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.cpp
@@ -68,6 +68,14 @@
return {};
}
+ // Resolve and attach the bitcode on the op now (at kernel-config time),
+ // so it survives all the way down to `LowerUKernelOpsToCalls` as a
+ // discardable attribute. Mirrors the GPU side
+ // (`ensureUKernelBitcodeAndFinalizeConfig` in
+ // `compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUSelectUKernels.cpp`)
+ // and makes the configuration-pass output self-contained for lit tests.
+ IREE::CPU::attachUKernelBitcodeOnOp(op, name);
+
MLIRContext *context = op->getContext();
return IREE::Codegen::UKernelDescriptorAttr::get(
context, StringAttr::get(context, name),
