[Codegen][CPU] Wire LowerBitcodeUKernels into the Mmt4dTilingExpert pipeline. (#24570)
Completes the new-style C-ukernel pipeline integration:
* `LLVMCPUTarget.cpp`: install `iree_codegen.ukernel_provider =
#iree_cpu.ukernel_provider` on every LLVMCPU executable target config
whenever `--iree-llvmcpu-enable-llvm-ukernels` is non-empty. This
mirrors what ROCM does and is what lets the generic
`LowerBitcodeUKernelsPass` find the CPU provider for the
bitcode-lookup-and-attach.
* `KernelDispatch.cpp::setRootConfig(InnerTiledOp)`: when
`selectUKernel` returns a descriptor, set the reduction-dim vector tile
size to 0 (untiled) instead of 1 (unrolled). The ukernel's C function
takes the outer K extent as a runtime argument (`k_outer`) and runs the
K loop itself; pre-unrolling here would both defeat that and bloat the
IR. Matches the equivalent split in
`setDataTiledMmaInnerTiledLoweringConfig` on the GPU side.
* `Passes.cpp::addMmt4dTilingExpertPassPipeline`: insert
`createLowerBitcodeUKernelsPass()` between the parallel-tile and
reduction-tile passes. A nop when no op carries a `iree_codegen.ukernel`
descriptor, so it stays in the default pipeline. The companion
`LowerUKernelOpsToCalls` already runs after bufferization (in
`addLowerToLLVMPasses`); together they implement the user's "early
`inner_tiled` -> `ukernel.generic` on tensors; late `ukernel.generic` ->
`func.call` on memrefs" split.
Lit test (`select_ukernel.mlir`) gains a second RUN line that chains
`iree-llvmcpu-select-lowering-strategy` with
`iree-codegen-lower-bitcode-ukernels` and asserts:
- the `inner_tiled` is rewritten to `iree_codegen.ukernel.generic`
carrying the right `iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16` name;
- the matching bitcode appears as `hal.executable.objects` (resolved
from the global `EmbeddedDataDirectory` populated at LLVMCPU plugin
init);
- and the original `inner_tiled` is gone, guarding against the rewrite
silently no-op'ing on a regression.
An end-to-end `iree-compile` test would also be desirable, but the
current pipeline state needs more wiring (workgroup-distribution
verifier complains about the bufferization-introduced copy ops around
the workgroup loop, on both the ukernel and non-ukernel paths) — that
investigation belongs with the follow-up commit that fills in real
ukernel bodies and an actual numerical matmul test.
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/LLVMCPUTarget.cpp b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
index ad04206..fda3035 100644
--- a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
+++ b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
@@ -180,6 +180,14 @@
configItems.emplace_back(
b.getStringAttr(IREE::Encoding::kEncodingResolverAttrName),
IREE::CPU::CPUEncodingResolverAttr::get(context, {}));
+ // When at least one LLVM-bitcode ukernel category is enabled, install
+ // the CPU ukernel provider attribute so that the subsequent
+ // `LowerBitcodeUKernelsPass` can find and attach the matching
+ // bitcode.
+ if (!target.llvmUkernels.empty()) {
+ configItems.emplace_back(b.getStringAttr(kUKernelProviderName),
+ IREE::CPU::UKernelProviderAttr::get(context));
+ }
// Compute the format used at runtime to select the executable loader.
std::string format;
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/select_ukernel.mlir b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/select_ukernel.mlir
index 6fc0834..2952f41 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/select_ukernel.mlir
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/select_ukernel.mlir
@@ -1,5 +1,7 @@
// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' \
// RUN: --split-input-file %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy,func.func(iree-codegen-lower-bitcode-ukernels))' \
+// RUN: --split-input-file %s | FileCheck %s --check-prefix=CHAIN
// Drives the LLVMCPU `LLVMCPUSelectLoweringStrategy` pass on
// `inner_tiled` rooted dispatches carrying a CPU `DataTiledMMAAttr`. The
@@ -10,10 +12,12 @@
// doesn't run target-option serialization).
// Enabled case: `selectCPUUKernel` matches the BF16 1x16x2 intrinsic and
-// annotates the op with `iree_codegen.ukernel = …`.
+// annotates the op with `iree_codegen.ukernel = …`. The chained CHAIN run
+// additionally lowers that descriptor to a `ukernel.generic`.
#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",
+ 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"
@@ -36,10 +40,22 @@
} : tensor<2x4x1x2xbf16>, tensor<2x4x16x2xbf16> into tensor<2x2x1x16xf32>
return %0 : tensor<2x2x1x16xf32>
}
+// The reduction (K) dim is left untiled: the ukernel owns the K loop, so the
+// lowering_config carries no `vector_reduction` entry (contrast the disabled
+// case below, which tiles it to `vector_reduction = [0, 0, 1]`).
+// CHECK: #iree_cpu.lowering_config<distribution = [1, 1, 0], vector_common_parallel = [1, 1, 0]>
// 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>
+// The next pass lowers the descriptor to a `ukernel.generic` carrying the
+// resolved bitcode, and the original `inner_tiled` is gone.
+// CHAIN-LABEL: func.func @bf16_inner_tiled_ukernel_enabled
+// CHAIN: iree_codegen.ukernel.generic
+// CHAIN-SAME: hal.executable.objects = [{{.*}}"iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.x86_64_avx512bf16.bc"
+// CHAIN-SAME: "iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16"
+// CHAIN-NOT: iree_codegen.inner_tiled
+
// -----
// Disabled case: target config lacks `llvm_ukernels`, so `selectCPUUKernel`
@@ -69,6 +85,16 @@
} : tensor<2x4x1x2xbf16>, tensor<2x4x16x2xbf16> into tensor<2x2x1x16xf32>
return %0 : tensor<2x2x1x16xf32>
}
+// Without a ukernel, the reduction (K) dim is tiled normally.
+// CHECK: #iree_cpu.lowering_config<distribution = [1, 1, 0], vector_common_parallel = [1, 1, 0], vector_reduction = [0, 0, 1]>
// CHECK-LABEL: func.func @bf16_inner_tiled_ukernel_disabled
// CHECK: iree_codegen.inner_tiled
// CHECK-NOT: iree_codegen.ukernel = #iree_codegen.ukernel_descriptor
+
+// Same with the chained pipeline: SelectLoweringStrategy doesn't attach a
+// descriptor, so LowerBitcodeUKernels has nothing to rewrite and the
+// `inner_tiled` survives untouched.
+// CHAIN-LABEL: func.func @bf16_inner_tiled_ukernel_disabled
+// CHAIN: iree_codegen.inner_tiled
+// CHAIN-NOT: iree_codegen.ukernel.generic
+// CHAIN-NOT: hal.executable.objects
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index 8c4791b..c075923 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -3051,13 +3051,11 @@
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
+ // target config; off by default. When a ukernel is selected, 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)) {
+ IREE::Codegen::UKernelDescriptorAttr ukernelDescriptor = selectCPUUKernel(op);
+ if (ukernelDescriptor) {
setUKernelDescriptor(op, ukernelDescriptor);
}
SmallVector<int64_t> bounds;
@@ -3085,11 +3083,21 @@
}
}
- // Vector tiling: every iter dim down to a single inner tile. The
- // generator's `splitParallelAndReductionTiles` lifts the parallel entries
- // into `vector_common_parallel` and the reduction entries into
- // `vector_reduction`.
+ // Vector tiling: every parallel iter dim down to a single inner tile.
+ // Reduction dims tile to 1 by default — codegen unrolls the K loop. But
+ // when a ukernel is in play, leave reduction dims untiled (size 0): the
+ // ukernel's C function takes the outer K count as a runtime argument
+ // (`k_outer`) and runs the K loop itself; pre-unrolling it here would
+ // both defeat that and bloat the IR. This mirrors the equivalent split
+ // in `setDataTiledMmaInnerTiledLoweringConfig` on the GPU side.
SmallVector<int64_t> vecTileSizes(numLoops, 1);
+ if (ukernelDescriptor) {
+ for (auto [i, kind] : llvm::enumerate(iteratorTypes)) {
+ if (kind == utils::IteratorType::reduction) {
+ vecTileSizes[i] = 0;
+ }
+ }
+ }
LoweringConfigGenerator generator(op);
generator.setDistributionTileSizes(distTileSizes);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
index e9f4fb6..16020e8 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -348,6 +348,15 @@
funcPassManager.addPass(createCPUPrepareUkernelsPass());
funcPassManager.addPass(createCPULowerToUKernelsPass(
pipelineOpt.cpuOpts.skipIntermediateRoundings));
+ // Rewrite `inner_tiled` ops carrying an `iree_codegen.ukernel` descriptor
+ // to `iree_codegen.ukernel.generic` calls (set earlier by `selectUKernel`
+ // in `KernelDispatch`, gated by `--iree-llvmcpu-enable-llvm-ukernels`).
+ // Must run after the parallel-tile pass above so the inner_tiled has its
+ // M/N outer extents tiled to single tiles, and before bufferization since
+ // the rewrite happens on tensors (the eventual `ukernel.generic` -> func
+ // call lowering happens after bufferize). A nop when no op carries the
+ // descriptor, so safe to leave in the default pipeline.
+ funcPassManager.addPass(createLowerBitcodeUKernelsPass());
funcPassManager.addPass(createLLVMCPUTileRootAndFuseInputOperandsPass(
IREE::CPU::TilingLevel::VectorReductionTiles));
// `VectorInnerParallelTiles` level models the tiling and fusion for the
