[LLVMCPU] Make cpu_data feature field a compile-time constant (#24811)
The processor_data passed to a ukernel call is built with the CPU
feature bitmask set to a compile-time constant taken from the executable
target's cpu_features. The remaining fields of environment's cpu_data
(currently unused) are copied from the runtime processor data.
A constant feature field lets each ukernel's feature check fold at
compile time, so the matching tile is selected and inlined and the
off-target tile variants are DCE'd before codegen.
Assisted-by: Claude Code
---------
Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp
index d6e993a..13ccfd7 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Codegen/LLVMCPU/DispatchABI.h"
+#include <cstdint>
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/schemas/cpu_data.h"
@@ -12,6 +13,7 @@
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/Path.h"
+#include "llvm/TargetParser/Triple.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
@@ -893,38 +895,35 @@
di.getBasicType(resultValue.getType()), builder);
}
-Value HALDispatchABI::updateProcessorDataFromTargetAttr(
- Operation *forOp, Value processorDataPtrValue, OpBuilder &builder) {
- // Get the target attr.
- IREE::HAL::ExecutableTargetAttr targetAttr =
- IREE::HAL::ExecutableTargetAttr::lookup(forOp);
+static uint64_t getFeatureBitPattern(Operation *forOp) {
+ uint64_t specifiedCpuDataField0 = 0;
+ auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(forOp);
if (!targetAttr) {
- return processorDataPtrValue;
+ return specifiedCpuDataField0;
}
DictionaryAttr targetConfig = targetAttr.getConfiguration();
-
- // Lookup CPU features.
std::optional<StringRef> cpuFeatures = getConfigCpuFeatures(targetConfig);
if (!cpuFeatures) {
- return processorDataPtrValue;
+ return specifiedCpuDataField0;
}
-
- // Currently requiring all CPU feature bits to be in field 0. Generalize as
- // needed when other CPU feature fields start to be used.
- uint64_t specifiedCpuDataField0 = 0;
- {
- // Map llvm feature-name to bit used to represent it in IREE_CPUDATA_FIELD0.
- //
- // TODO(ravishankarm): This link to the runtime schemas needs to be broken.
- // Instead we should use a reflection callback to resolve arch guarded
- // features directly in the compiler.
- llvm::StringMap<uint64_t> featureToBitPattern;
- auto targetTriple = getTargetTriple(targetConfig);
- if (!targetTriple) {
- return processorDataPtrValue;
- }
- std::string targetArchUppercase =
- StringRef(getIreeArchNameForTargetTriple(targetTriple.value())).upper();
+ std::optional<llvm::Triple> targetTriple = getTargetTriple(targetConfig);
+ if (!targetTriple) {
+ return specifiedCpuDataField0;
+ }
+ // Currently requiring all CPU feature bits to be in field 0. Generalize
+ // as needed when other CPU feature fields start to be used.
+ // The remaining fields _can_ carry architecture-defined runtime processor
+ // data and are passed through unchanged.
+ //
+ // Map llvm feature-name to bit used to represent it in
+ // IREE_CPUDATA_FIELD0.
+ //
+ // TODO(ravishankarm): This link to the runtime schemas needs to be
+ // broken. Instead we should use a reflection callback to resolve arch
+ // guarded features directly in the compiler.
+ llvm::StringMap<uint64_t> featureToBitPattern;
+ std::string targetArchUppercase =
+ StringRef(getIreeArchNameForTargetTriple(targetTriple.value())).upper();
#define IREE_CPU_FEATURE_BIT(arch, field_index, bit_pos, bit_name, llvm_name) \
if (targetArchUppercase == #arch) { \
assert(field_index == 0); \
@@ -933,36 +932,36 @@
#include "iree/schemas/cpu_feature_bits.inl"
#undef IREE_CPU_FEATURE_BIT
- // Find CPU features in featureToBitPattern
- SmallVector<StringRef> cpuFeatureStrings;
- cpuFeatures.value().split(cpuFeatureStrings, ',', /*MakeSplit=*/-1,
- /*KeepEmpty=*/false);
- for (auto featureString : cpuFeatureStrings) {
- // CPU features are typically prefixed with a +, e.g. +avx,+avx2,+fma.
- featureString.consume_front("+");
- // Silently skip unknown CPU features, more flexible for now. Note that
- // some features occurring here are not standard CPU features but internal
- // things such as the "+reserve-x18" that we add on arm64.
- if (featureToBitPattern.count(featureString)) {
- specifiedCpuDataField0 |= featureToBitPattern.lookup(featureString);
- }
+ // Find CPU features in featureToBitPattern.
+ SmallVector<StringRef> cpuFeatureStrings;
+ cpuFeatures.value().split(cpuFeatureStrings, ',', /*MakeSplit=*/-1,
+ /*KeepEmpty=*/false);
+ for (auto featureString : cpuFeatureStrings) {
+ // CPU features are typically prefixed with a +, e.g. +avx,+avx2,+fma.
+ featureString.consume_front("+");
+ // Silently skip unknown CPU features, more flexible for now. Note that
+ // some features occurring here are not standard CPU features but
+ // internal things such as the "+reserve-x18" that we add on arm64.
+ if (featureToBitPattern.count(featureString)) {
+ specifiedCpuDataField0 |= featureToBitPattern.lookup(featureString);
}
}
- if (specifiedCpuDataField0 == 0) {
- return processorDataPtrValue;
- }
+ return specifiedCpuDataField0;
+}
- // Create a new stack allocation for the bit pattern.
+Value HALDispatchABI::updateProcessorDataFromTargetAttr(
+ Operation *forOp, Value processorDataPtrValue, OpBuilder &builder) {
+ uint64_t specifiedCpuDataField0 = getFeatureBitPattern(forOp);
Location loc = forOp->getLoc();
MLIRContext *context = forOp->getContext();
auto ptrType = LLVM::LLVMPointerType::get(context);
auto i64Ty = builder.getI64Type();
// The stack allocation goes into the entry point of the block.
- // The compile-time cpu features are patched onto `cpu_data` from the target
- // environment. That should happen in the loop body, so that the compile-time
- // cpu features are visible to the post-link LLVM optimizations that would
- // fold the microkernel tile size selection logic that checks these, and the
- // selection happens at compile-time.
+ // The compile-time cpu features overwrite the field 0 of `cpu_data` from the
+ // target environment. That should happen in the loop body, so that the
+ // compile-time cpu features are visible to the post-link LLVM optimizations
+ // that would fold the microkernel tile size selection logic that checks
+ // these, and the selection happens at compile-time.
Value alloca;
{
auto funcOp = forOp->getParentOfType<LLVM::LLVMFuncOp>();
@@ -974,15 +973,11 @@
alloca = LLVM::AllocaOp::create(builder, loc, ptrType, i64Ty, arraySize,
/*alignment=*/sizeof(uint64_t));
}
- // Load the 0-th value.
- Value srcData0 =
- LLVM::LoadOp::create(builder, loc, i64Ty, processorDataPtrValue);
- // Set the specified CPU arch data.
- Value bitPatternVal = LLVM::ConstantOp::create(
+ // Field 0: the compile-time CPU feature bitmask.
+ Value field0 = LLVM::ConstantOp::create(
builder, loc, i64Ty, builder.getI64IntegerAttr(specifiedCpuDataField0));
- srcData0 = LLVM::OrOp::create(builder, loc, srcData0, bitPatternVal);
- LLVM::StoreOp::create(builder, loc, srcData0, alloca);
- // Copy over the rest.
+ LLVM::StoreOp::create(builder, loc, field0, alloca);
+ // Fields 1..N: the runtime processor data, copied through.
for (int64_t i = 1, e = ProcessorDataCapacity; i < e; ++i) {
Value loadPtr = LLVM::GEPOp::create(
builder, loc, processorDataPtrValue.getType(), i64Ty,
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir
index 501ca10..f6b2b48 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir
@@ -31,15 +31,25 @@
}
}
// CHECK: llvm.func @default_cconv_with_extra_fields(!llvm.ptr, i32, f64, !llvm.ptr, i32) -> f32
-// CHECK: llvm.func @bar
+// CHECK: llvm.func @bar(%[[ENV:[^:]+]]:
+// CHECK-DAG: %[[C0:.+]] = llvm.mlir.constant(0 : i64) : i64
// CHECK-DAG: %[[Ci32:.+]] = llvm.mlir.constant(42 : i32) : i32
// CHECK-DAG: %[[Cf64:.+]] = llvm.mlir.constant(4.200000e+01 : f64) : f64
-// CHECK-DAG: %[[ALLOCA:.+]] = llvm.alloca
-// CHECK-DAG: %[[DATA:.+]] = llvm.getelementptr inbounds %arg0[4]
+// CHECK: %[[CPU_DATA:.+]] = llvm.alloca %{{.+}} x i64
+// CHECK: %[[ALLOCA:.+]] = llvm.alloca %{{.+}} x f32
+// CHECK: %[[ENV_DATA:.+]] = llvm.getelementptr inbounds %[[ENV]][4]
+// CHECK: llvm.store %[[C0]], %[[CPU_DATA]]
+// Remaining fields are copied from the runtime processor data.
+// CHECK: %[[SRC1:.+]] = llvm.getelementptr inbounds %[[ENV_DATA]][1]
+// CHECK: %[[VAL1:.+]] = llvm.load %[[SRC1]]
+// CHECK: %[[DST1:.+]] = llvm.getelementptr inbounds %[[CPU_DATA]][1]
+// CHECK: llvm.store %[[VAL1]], %[[DST1]]
// CHECK-DAG: %[[PROCESSOR_INFO:.+]] = llvm.load %arg2
// CHECK: %[[PROCESSOR_ID:.+]] = llvm.extractvalue %[[PROCESSOR_INFO]][4]
+// With no executable target the feature bitmask is zero; the ukernel receives a
+// cpu_data stack buffer (field 0 constant, fields 1..N from the runtime data).
// CHECK: %[[VAL:.+]] = llvm.call @default_cconv_with_extra_fields
-// CHECK-SAME: (%[[ALLOCA]], %[[Ci32]], %[[Cf64]], %[[DATA]], %[[PROCESSOR_ID]])
+// CHECK-SAME: (%[[ALLOCA]], %[[Ci32]], %[[Cf64]], %[[CPU_DATA]], %[[PROCESSOR_ID]])
// -----
@@ -49,8 +59,8 @@
// site instead of the function entry block, it would re-execute on every
// loop iteration without ever popping the stack back off, overflowing it for
// large iteration counts. So the alloca must be in the entry block and must not
-// reappear in the loop body, the cpu feature patching itself should be in the loop body
-// next to the call.
+// reappear in the loop body, the cpu feature override on the field0 itself should
+// be in the loop body next to the call.
#executable_target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+dotprod", target_triple = "aarch64-none-elf"}>
module {
func.func private @default_cconv_with_extra_fields_in_loop(memref<f32>, i32, f64) -> (f32) attributes {
@@ -72,26 +82,31 @@
return
}
}
-// CHECK: llvm.func @loop_caller(%[[ARG0:.+]]: {{.*}}llvm.ptr{{.*}}, %[[ARG1:.+]]: {{.*}}llvm.ptr{{.*}}, %[[ARG2:.+]]: {{.*}}llvm.ptr{{.*}})
-// Entry block: only the patch buffer's stack slot is reserved here, once.
+// CHECK: llvm.func @loop_caller(%[[ENV:[^:]+]]:
+// Entry block: field 0 is a compile-time constant (`+dotprod` == 4096) and the
+// cpu_data buffer's stack slot is reserved here, once.
// CHECK-NOT: ^{{.+}}:
-// CHECK: %[[PATCHED_DATA:.+]] = llvm.alloca %{{.+}} x i64
+// CHECK: %[[FIELD0:.+]] = llvm.mlir.constant(4096 : i64)
+// CHECK-NOT: ^{{.+}}:
+// CHECK: %[[CPU_DATA:.+]] = llvm.alloca %{{.+}} x i64
// CHECK-NOT: ^{{.+}}:
// CHECK: llvm.br ^[[HEADER:.+]](
// Loop header: just the trip-count check.
// CHECK: ^[[HEADER]]
// CHECK: llvm.cond_br %{{.+}}, ^[[BODY:.+]], ^[[EXIT:.+]]
-// Loop body: no fresh alloca; the compile-time cpu features are patched onto the target
-// environment here in the loop body, storing into the entry-block buffer right before the call.
+// Loop body: no fresh alloca. Field 0 stores the constant cpu features; remaining fields
+// are copied from the runtime processor data.
// CHECK: ^[[BODY]]
// CHECK-NOT: llvm.alloca
-// CHECK: %[[ENV_DATA:.+]] = llvm.getelementptr inbounds %[[ARG0]]
-// CHECK: %[[SRC0:.+]] = llvm.load %[[ENV_DATA]]
-// CHECK: %[[PATCHED0:.+]] = llvm.or %[[SRC0]], %{{.+}}
-// CHECK: llvm.store %[[PATCHED0]], %[[PATCHED_DATA]]
+// CHECK: %[[ENV_DATA:.+]] = llvm.getelementptr inbounds %[[ENV]][4]
+// CHECK: llvm.store %[[FIELD0]], %[[CPU_DATA]]
+// CHECK: %[[SRC1:.+]] = llvm.getelementptr inbounds %[[ENV_DATA]][1]
+// CHECK: %[[VAL1:.+]] = llvm.load %[[SRC1]]
+// CHECK: %[[DST1:.+]] = llvm.getelementptr inbounds %[[CPU_DATA]][1]
+// CHECK: llvm.store %[[VAL1]], %[[DST1]]
// CHECK-NOT: llvm.alloca
// CHECK: llvm.call @default_cconv_with_extra_fields_in_loop
-// CHECK-SAME: %[[PATCHED_DATA]]
+// CHECK-SAME: %[[CPU_DATA]]
// CHECK: llvm.br ^[[HEADER]]
// Loop exit.
// CHECK: ^[[EXIT]]
