[LLVMCPU] strip per-function target attributes to allow inlining ukernels (#24833)
The ukernel bitcode carries per-function
target-cpu/tune-cpu/target-features attributes that conflict with the
attribute-less dispatch (which inherits the module-level TargetMachine)
and stop LLVM from inlining the ukernel. Strip them (and vscale_range,
so RVV bitcode inlines) at link time.
Fixes #24755.
Assisted-by: Claude Code
Signed-off-by: Ege Beysel <beyselege@gmail.com>
Co-authored-by: Stefan Schuermans <schuermans@roofline.ai>
Co-authored-by: Federico Bruzzone <federico.bruzzone.i@gmail.com>
diff --git a/compiler/plugins/target/LLVMCPU/builtins/UKernel.cpp b/compiler/plugins/target/LLVMCPU/builtins/UKernel.cpp
index f216c0f..01da2f4 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/UKernel.cpp
+++ b/compiler/plugins/target/LLVMCPU/builtins/UKernel.cpp
@@ -8,6 +8,7 @@
#include "iree/builtins/ukernel/ukernel_bitcode.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/HAL/Utils/LLVMLinkerUtils.h"
#include "llvm/Bitcode/BitcodeReader.h"
#include "llvm/Support/MemoryBufferRef.h"
#include "mlir/Support/LLVM.h"
@@ -55,8 +56,13 @@
// combination of data types, a specific SIMD ISA variant, etc. Then all the
// unused code paths can get DCE'd. That's why failure to inline a ukernel
// can result in a large penalty in both performance and code size.
+ // Strip the per-function target attributes so they can't conflict with the
+ // attribute-less dispatch and block inlining; drop vscale_range too so RVV
+ // bitcode inlines.
for (auto &func : module.get()->functions()) {
func.addFnAttr(llvm::Attribute::AlwaysInline);
+ stripFunctionTargetAttrs(func);
+ func.removeFnAttr(llvm::Attribute::VScaleRange);
}
return module;
}
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
index 5293261..21c7a28 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/BUILD.bazel
@@ -22,6 +22,7 @@
"lower_inner_tiled_to_bitcode_ukernel.mlir",
"lower_inner_tiled_to_bitcode_ukernel_i8_vnni.mlir",
"select_ukernel.mlir",
+ "ukernel_inlining.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 9373360..5d316e7 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/CMakeLists.txt
@@ -19,6 +19,7 @@
"lower_inner_tiled_to_bitcode_ukernel.mlir"
"lower_inner_tiled_to_bitcode_ukernel_i8_vnni.mlir"
"select_ukernel.mlir"
+ "ukernel_inlining.mlir"
TOOLS
FileCheck
iree-compile
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/ukernel_inlining.mlir b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/ukernel_inlining.mlir
new file mode 100644
index 0000000..2d5f564
--- /dev/null
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/test/ukernel_inlining.mlir
@@ -0,0 +1,60 @@
+// RUN: rm -rf %t && mkdir -p %t
+
+// x86-64, AVX-512 VNNI:
+// RUN: iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=llvm-cpu \
+// RUN: --iree-llvmcpu-target-triple=x86_64-unknown-unknown-eabi-elf \
+// RUN: --iree-llvmcpu-target-cpu=znver4 \
+// RUN: --iree-opt-data-tiling --iree-llvmcpu-enable-ukernels=all \
+// RUN: --iree-hal-dump-executable-intermediates-to=%t/x86 %s -o /dev/null
+// RUN: cat %t/x86/*.linked.ll | FileCheck %s --check-prefix=LINKED
+// RUN: cat %t/x86/*.optimized.ll | FileCheck %s --check-prefix=OPT
+
+// aarch64, i8mm:
+// RUN: iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=llvm-cpu \
+// RUN: --iree-llvmcpu-target-triple=aarch64-none-elf \
+// RUN: --iree-llvmcpu-target-cpu-features=+i8mm \
+// RUN: --iree-opt-data-tiling --iree-llvmcpu-enable-ukernels=mmt4d \
+// RUN: --iree-hal-dump-executable-intermediates-to=%t/i8mm %s -o /dev/null
+// RUN: cat %t/i8mm/*.linked.ll | FileCheck %s --check-prefix=LINKED
+// RUN: cat %t/i8mm/*.optimized.ll | FileCheck %s --check-prefix=OPT
+
+// aarch64, dotprod:
+// RUN: iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=llvm-cpu \
+// RUN: --iree-llvmcpu-target-triple=aarch64-none-elf \
+// RUN: --iree-llvmcpu-target-cpu-features=+dotprod \
+// RUN: --iree-opt-data-tiling --iree-llvmcpu-enable-ukernels=mmt4d \
+// RUN: --iree-hal-dump-executable-intermediates-to=%t/dotprod %s -o /dev/null
+// RUN: cat %t/dotprod/*.linked.ll | FileCheck %s --check-prefix=LINKED
+// RUN: cat %t/dotprod/*.optimized.ll | FileCheck %s --check-prefix=OPT
+
+// riscv64, RVV + xsmtvdot:
+// RUN: iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=llvm-cpu \
+// RUN: --iree-llvmcpu-target-triple=riscv64-unknown-elf \
+// RUN: --iree-llvmcpu-target-cpu-features=+m,+a,+f,+d,+c,+v,+xsmtvdot,+zvl256b \
+// RUN: --iree-opt-data-tiling --iree-llvmcpu-enable-ukernels=mmt4d \
+// RUN: --iree-hal-dump-executable-intermediates-to=%t/rvv %s -o /dev/null
+// RUN: cat %t/rvv/*.linked.ll | FileCheck %s --check-prefix=LINKED
+// RUN: cat %t/rvv/*.optimized.ll | FileCheck %s --check-prefix=OPT
+
+// Ukernels must inline into the dispatch so their compile-time specialization
+// can be DCE'd. For each target above, check that a ukernel is linked in and
+// that, after optimization, no call to a ukernel symbol survives in the dispatch.
+
+// LINKED: call {{.*}}@iree_uk_mmt4d(
+// OPT: define {{.*}}_dispatch_0
+// OPT-NOT: call {{.*}}@iree_uk_
+
+func.func @matmul(%a: tensor<128x128xi8>, %b: tensor<128x128xi8>)
+ -> tensor<128x128xi32> {
+ %c0 = arith.constant 0 : i32
+ %init = tensor.empty() : tensor<128x128xi32>
+ %fill = linalg.fill ins(%c0 : i32) outs(%init : tensor<128x128xi32>)
+ -> tensor<128x128xi32>
+ %result = linalg.matmul ins(%a, %b : tensor<128x128xi8>, tensor<128x128xi8>)
+ outs(%fill : tensor<128x128xi32>) -> tensor<128x128xi32>
+ return %result : tensor<128x128xi32>
+}
