[LLVMCPU] Use native bf16 converts when the target supports them (#24759)
## Summary
The bf16 arithmetic is getting promoted to f32 in codegen, which is
correct for every CPU (there's no CPU with native bf16 arithmetic yet,
RISC-V `Zvfbfmin` provides conversions only). The conversions around
that promotion, however, are always expanded in software. That prevents
targets that do have bf16 <-> f32 conversion instructions from selecting
them, costing ~6 integer ops per narrow and 2 per widen on every bf16
load and store.
Introduce another LLVMCPULoweringPipeline option: it's set to true on
RISC-V targets with both `Zfbfmin` and `Zvfbfmin` extensions. Both are
required, since with only the vector extension scalar bf16 residue
lowers to `__truncsfbf2` libcalls. bf16 loads/stores and the promotion's
extf/truncf then select as `vfwcvtbf16/vfncvtbf16` (vector) and
`fcvt.s.bf16/fcvt.bf16.s` (scalar). This option is passed to
`ConvertUnsupportedFloatToIntBuffers` and `arith-expand` passes.
## Testing
Measured on a SpaceMiT K3 (X100 cores) (yay, i have access to one now).
Tested on kernels: 1M-element bf16 kernels, native vs current expansion,
outputs bit-identical in every case:
| kernel | native | expanded | speedup |
| :-- | --: | --: | --: |
| elementwise add | 0.800 ms | 1.45 ms | 1.81× |
| f32 to bf16 cast | 0.650 ms | 1.01 ms | 1.55× |
| sigmoid-weighted elementwise chain | 7.83 ms | 8.31 ms | 1.06× |
| bf16 to f32 cast | 0.494 ms | 0.538 ms | 1.09× |
Also tested on whole models: a bf16 `whisper-tiny-en` encoder-decoder
runs in **3.580 s vs 3.936 s (1.10×)** with bit-identical output against
PyTorch reference. A small CLIP-style bf16 model compiles to a 7.5%
smaller vmfb.
---------
Signed-off-by: Zmicier Prybysh <zprybysh@baylibre.com>
diff --git a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
index fda3035..9a8ec2e 100644
--- a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
+++ b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
@@ -259,10 +259,14 @@
void buildTranslationPassPipeline(IREE::HAL::ExecutableTargetAttr targetAttr,
OpPassManager &passManager) final {
- bool enableAArch64SME = isAArch64(targetAttr.getConfiguration()) &&
- hasSMEFeature(targetAttr.getConfiguration());
- buildLLVMCPUCodegenPassPipeline(passManager.nest<ModuleOp>(),
- codegenOptions_, enableAArch64SME);
+ DictionaryAttr config = targetAttr.getConfiguration();
+ LLVMCPUPipelineOptions pipelineOpts;
+ pipelineOpts.cpuOpts = codegenOptions_;
+ pipelineOpts.enableAArch64SME = isAArch64(config) && hasSMEFeature(config);
+ pipelineOpts.enableNativeBf16Converts = isRISCV(config) &&
+ hasZfbfminFeature(config) &&
+ hasZvfbfminFeature(config);
+ buildLLVMCPUCodegenPassPipeline(passManager.nest<ModuleOp>(), pipelineOpts);
buildCodegenTranslationPostProcessingPassPipeline(passManager);
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
index edb73bd..d445c8a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
@@ -168,6 +168,9 @@
isAArch64(targetConfig) && hasSMEFeature(targetConfig);
pipelineOpts.enableAArch64I8mm =
isAArch64(targetConfig) && hasI8mmFeature(targetConfig);
+ pipelineOpts.enableNativeBf16Converts = isRISCV(targetConfig) &&
+ hasZfbfminFeature(targetConfig) &&
+ hasZvfbfminFeature(targetConfig);
pipelineOpts.enablePeeling = isOptEnabled(funcOp, getEnableLoopPeelingStr());
LoweringConfigAttrInterface loweringConfig = getRootLoweringConfig(funcOp);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
index 91fa67f..9984e0c 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -483,8 +483,8 @@
}
static void addLowerToLLVMPasses(OpPassManager &modulePassManager,
- bool enableAArch64SME,
- const CPUCodegenOptions &cpuOpts) {
+ const LLVMCPUPipelineOptions &pipelineOpts) {
+ const CPUCodegenOptions &cpuOpts = pipelineOpts.cpuOpts;
// TODO: Remove the following pass and plumb support for #hal.descriptor_type
// memory space through the stack.
FunctionLikeNest(modulePassManager)
@@ -502,12 +502,13 @@
createLLVMCPUEmitVectorizationRemarksPass)
.addPass(createConvertLinalgToLoopsPass)
.addPass(createConvertBf16ArithToF32Pass)
- .addPass([]() {
- // Convert bf16 buffers to i16. LLVM IR supports fp8 types
- // natively, so we don't need to convert them here.
+ .addPass([&]() {
+ // Convert bf16 buffers to i16, unless the target has native bf16
+ // converts. LLVM IR supports fp8 types natively, so we don't need
+ // to convert them here.
return createConvertUnsupportedFloatToIntBuffersPass(
ConvertUnsupportedFloatToIntBuffersPassOptions{
- /*includeBf16=*/true,
+ /*includeBf16=*/!pipelineOpts.enableNativeBf16Converts,
/*includeF8E5M2=*/false,
/*includeF8E4M3FN=*/false,
/*includeF8E5M2FNUZ=*/false,
@@ -532,7 +533,7 @@
.addPredicatedPass(cpuOpts.useFastMinMaxOps,
createReplaceSlowMinMaxOpsPass);
- if (enableAArch64SME) {
+ if (pipelineOpts.enableAArch64SME) {
modulePassManager.addPass(mlir::arm_sme::createVectorLegalizationPass());
FunctionLikeNest(modulePassManager)
.addPredicatedPass(
@@ -563,7 +564,7 @@
}
VectorTransferLoweringPassOptions transferLoweringOptions;
- if (!enableAArch64SME) {
+ if (!pipelineOpts.enableAArch64SME) {
// The ArmSME dialect has its own (more specific) lowerings for scalable
// vectors that occur later in the pipeline, so only enable the general
// lowerings if SME is not available.
@@ -598,7 +599,9 @@
.addPass(createIREEAffineExpandIndexOpsPass)
.addPass([&]() {
arith::ArithExpandOpsPassOptions options;
- options.includeBf16 = true;
+ // Keep bf16 extf/truncf intact when the target lowers them to native
+ // conversion instructions.
+ options.includeBf16 = !pipelineOpts.enableNativeBf16Converts;
options.includeF8E8M0 = true;
return arith::createArithExpandOpsPass(options);
})
@@ -609,7 +612,7 @@
.addPredicatedPass(cpuOpts.instrumentMemoryAccesses,
createInstrumentMemoryAccessesPass);
- if (enableAArch64SME) {
+ if (pipelineOpts.enableAArch64SME) {
FunctionLikeNest(modulePassManager).addPass([&] {
return createConvertArmSMEToLLVMPass();
});
@@ -682,14 +685,13 @@
}
void buildLLVMCPUCodegenPassPipeline(OpPassManager &modulePassManager,
- const CPUCodegenOptions &cpuOpts,
- bool enableAArch64SME,
+ const LLVMCPUPipelineOptions &pipelineOpts,
bool includeLLVMLowering) {
modulePassManager.addPass(createLowerExecutableUsingTransformDialectPass());
FunctionLikeNest(modulePassManager)
.addPass([&]() {
return createLLVMCPULowerExecutableTargetPass(
- LLVMCPULowerExecutableTargetPassOptions{cpuOpts});
+ LLVMCPULowerExecutableTargetPassOptions{pipelineOpts.cpuOpts});
})
.addPass(createVerifyWorkgroupDistributionPass);
if (clPatchFuncOps) {
@@ -702,7 +704,7 @@
modulePassManager.addPass(IREE::Util::createDropCompilerHintsPass());
if (includeLLVMLowering) {
- addLowerToLLVMPasses(modulePassManager, enableAArch64SME, cpuOpts);
+ addLowerToLLVMPasses(modulePassManager, pipelineOpts);
}
LLVM_DEBUG({
llvm::dbgs() << "LLVMCPU codegen pass pipeline:\n";
@@ -847,6 +849,11 @@
Option<bool> enableArmSME{
*this, "enable-arm-sme",
llvm::cl::desc("Enable the ArmSME lowering pipeline.")};
+ Option<bool> enableNativeBf16Converts{
+ *this, "enable-native-bf16-converts",
+ llvm::cl::desc("Assume the target has hardware bf16 <-> f32 converts "
+ "and keep bf16 in the IR instead of emulating it."),
+ llvm::cl::init(false)};
Option<bool> includeLLVMLowering{
*this, "include-llvm-lowering",
llvm::cl::desc("Include the lowering to LLVM dialect."),
@@ -859,10 +866,13 @@
"Runs the LLVMCPU lowering pipeline",
[](OpPassManager &modulePassManager,
LLVMCPULoweringPipelineOptions const &options) {
- CPUCodegenOptions cpuOpts =
+ LLVMCPUPipelineOptions pipelineOpts;
+ pipelineOpts.cpuOpts =
getCPUCodegenOptionsForTextualPipeline(options);
- buildLLVMCPUCodegenPassPipeline(modulePassManager, cpuOpts,
- options.enableArmSME,
+ pipelineOpts.enableAArch64SME = options.enableArmSME;
+ pipelineOpts.enableNativeBf16Converts =
+ options.enableNativeBf16Converts;
+ buildLLVMCPUCodegenPassPipeline(modulePassManager, pipelineOpts,
options.includeLLVMLowering);
});
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h
index 5e71246..e08c1a9 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h
@@ -85,6 +85,7 @@
bool enableVectorMasking = false;
bool enableAArch64SME = false;
bool enableAArch64I8mm = false;
+ bool enableNativeBf16Converts = false;
bool lowerToAVX2 = false;
};
@@ -146,8 +147,7 @@
/// to LLVM dialect via the structured ops path. The `modulePassManager`
/// should operate on the module within the IREE::HAL::ExecutableOp.
void buildLLVMCPUCodegenPassPipeline(OpPassManager &modulePassManager,
- const CPUCodegenOptions &codegenOptions,
- bool enableAArch64SME = false,
+ const LLVMCPUPipelineOptions &pipelineOpts,
bool includeLLVMLowering = true);
//----------------------------------------------------------------------------//
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.cpp
index 7d1c3b2..d76f7c7 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.cpp
@@ -89,6 +89,14 @@
hasFeature(targetConfig, "+zve64d");
}
+bool hasZfbfminFeature(DictionaryAttr targetConfig) {
+ return hasFeature(targetConfig, "+zfbfmin");
+}
+
+bool hasZvfbfminFeature(DictionaryAttr targetConfig) {
+ return hasFeature(targetConfig, "+zvfbfmin");
+}
+
bool hasSMEFeature(DictionaryAttr targetConfig) {
return hasFeature(targetConfig, "+sme");
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.h b/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.h
index 2840d0d..7cbdd63 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.h
@@ -49,6 +49,12 @@
/// cpu features.
bool hasAnyVFeature(DictionaryAttr targetConfig);
+/// Returns true if the 'targetAttr' contains '+zfbfmin' in its cpu features.
+bool hasZfbfminFeature(DictionaryAttr targetConfig);
+
+/// Returns true if the 'targetAttr' contains '+zvfbfmin' in its cpu features.
+bool hasZvfbfminFeature(DictionaryAttr targetConfig);
+
/// Returns true if the 'targetAttr' contains '+sme' in its cpu features.
bool hasSMEFeature(DictionaryAttr targetConfig);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
index 7bbe656..837115d 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
@@ -40,6 +40,7 @@
"lowering_strategy_from_tuning_spec.mlir",
"peel.mlir",
"pipeline_arm_sme_streaming_mode_tests.mlir",
+ "pipeline_bf16_native_converts_tests.mlir",
"pipeline_disable_distribution_tests.mlir",
"pipeline_full_smoketests.mlir",
"pipeline_gather_tests.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
index 25afeb8..8d75ee9 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
@@ -35,6 +35,7 @@
"lowering_strategy_from_tuning_spec.mlir"
"peel.mlir"
"pipeline_arm_sme_streaming_mode_tests.mlir"
+ "pipeline_bf16_native_converts_tests.mlir"
"pipeline_disable_distribution_tests.mlir"
"pipeline_full_smoketests.mlir"
"pipeline_gather_tests.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_bf16_native_converts_tests.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_bf16_native_converts_tests.mlir
new file mode 100644
index 0000000..25aeb5a
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_bf16_native_converts_tests.mlir
@@ -0,0 +1,96 @@
+// RUN: iree-opt --iree-codegen-llvmcpu-configuration-pipeline --iree-codegen-llvmcpu-lowering-pipeline='enable-native-bf16-converts=true' --split-input-file %s | FileCheck %s --check-prefixes=COMMON,NATIVE
+// RUN: iree-opt --iree-codegen-llvmcpu-configuration-pipeline --iree-codegen-llvmcpu-lowering-pipeline --split-input-file %s | FileCheck %s --check-prefixes=COMMON,EMULATED
+
+// Verifies the enable-native-bf16-converts pipeline option. The target backend
+// derives this option from the target's cpu features (Zfbfmin + Zvfbfmin on
+// RISC-V). Arithmetic is promoted to f32 either way, there is no non-widening
+// bf16 arithmetic (the only bf16-input arithmetic instruction, Zvfbfwma's
+// `vfwmaccbf16`, widens to f32). What the option controls is the conversions
+// around the promotion. When it's set, bf16 storage is kept and the promotion's
+// extf/truncf survive to the LLVM dialect as fpext/fptrunc, which select as
+// native conversion instructions, otherwise bf16 storage becomes i16 and the
+// conversions are expanded into shift/round-bias integer sequences.
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+
+#executable_target_embedded_elf_riscv_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-riscv_64", {cpu_features = "+m,+a,+f,+d,+c,+v,+zfbfmin,+zvfbfmin", data_layout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128", native_vector_size = 16 : index, target_triple = "riscv64-unknown-unknown-eabi-elf"}>
+builtin.module {
+ func.func @bf16_add() attributes {hal.executable.target = #executable_target_embedded_elf_riscv_64_} {
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xbf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xbf16>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(2) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xbf16>>
+ %lhs = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0], sizes = [1024], strides = [1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xbf16>> -> tensor<1024xbf16>
+ %rhs = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0], sizes = [1024], strides = [1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xbf16>> -> tensor<1024xbf16>
+ %init = tensor.empty() : tensor<1024xbf16>
+ %add = linalg.add ins(%lhs, %rhs : tensor<1024xbf16>, tensor<1024xbf16>) outs(%init : tensor<1024xbf16>) -> tensor<1024xbf16>
+ iree_tensor_ext.dispatch.tensor.store %add, %2, offsets = [0], sizes = [1024], strides = [1] : tensor<1024xbf16> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xbf16>>
+ return
+ }
+}
+// Native bf16 add: bf16 loads, fpext/fptrunc around the add, then bf16 store.
+// Emulated: the add is performed in f32 on values reconstructed from i16
+// storage, and the result is rounded back with the shift/bias sequence.
+// COMMON-LABEL: llvm.func @bf16_add
+// NATIVE: %[[LHS:.+]] = llvm.load {{.*}} -> bf16
+// NATIVE: %[[LHSF:.+]] = llvm.fpext %[[LHS]] : bf16 to f32
+// NATIVE: %[[RHS:.+]] = llvm.load {{.*}} -> bf16
+// NATIVE: %[[RHSF:.+]] = llvm.fpext %[[RHS]] : bf16 to f32
+// EMULATED: %[[LHS:.+]] = llvm.load {{.*}} -> i16
+// EMULATED: %[[LHSZ:.+]] = llvm.zext %[[LHS]] : i16 to i32
+// EMULATED: %[[LHSS:.+]] = llvm.shl %[[LHSZ]], {{.*}} : i32
+// EMULATED: %[[LHSF:.+]] = llvm.bitcast %[[LHSS]] : i32 to f32
+// EMULATED: %[[RHS:.+]] = llvm.load {{.*}} -> i16
+// EMULATED: %[[RHSZ:.+]] = llvm.zext %[[RHS]] : i16 to i32
+// EMULATED: %[[RHSS:.+]] = llvm.shl %[[RHSZ]], {{.*}} : i32
+// EMULATED: %[[RHSF:.+]] = llvm.bitcast %[[RHSS]] : i32 to f32
+// COMMON: %[[SUM:.+]] = llvm.fadd %[[LHSF]], %[[RHSF]] {{.*}} : f32
+// NATIVE: %[[RES:.+]] = llvm.fptrunc %[[SUM]] : f32 to bf16
+// NATIVE: llvm.store %[[RES]], {{.*}} : bf16
+// NATIVE-NOT: llvm.lshr
+// NATIVE-NOT: llvm.shl
+// EMULATED: llvm.fcmp "une" %[[SUM]], %[[SUM]]
+// EMULATED: llvm.lshr
+// EMULATED: llvm.store {{.*}} i16
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+
+#executable_target_embedded_elf_riscv_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-riscv_64", {cpu_features = "+m,+a,+f,+d,+c,+v,+zfbfmin,+zvfbfmin", data_layout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128", native_vector_size = 16 : index, target_triple = "riscv64-unknown-unknown-eabi-elf"}>
+builtin.module {
+ func.func @bf16_truncf() attributes {hal.executable.target = #executable_target_embedded_elf_riscv_64_} {
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xf32>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xbf16>>
+ %in = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0], sizes = [1024], strides = [1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xf32>> -> tensor<1024xf32>
+ %init = tensor.empty() : tensor<1024xbf16>
+ %trunc = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]}
+ ins(%in : tensor<1024xf32>) outs(%init : tensor<1024xbf16>) {
+ ^bb0(%a: f32, %out: bf16):
+ %t = arith.truncf %a : f32 to bf16
+ linalg.yield %t : bf16
+ } -> tensor<1024xbf16>
+ iree_tensor_ext.dispatch.tensor.store %trunc, %1, offsets = [0], sizes = [1024], strides = [1] : tensor<1024xbf16> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xbf16>>
+ return
+ }
+}
+// Both paths load the f32 input the same way. Native: the cast stays a single
+// fptrunc (selectable as one narrowing convert) and stores bf16. Emulated: the
+// round-to-nearest-even shift/bias sequence, storing i16.
+// COMMON-LABEL: llvm.func @bf16_truncf
+// COMMON: %[[IN:.+]] = llvm.load {{.*}} -> vector<{{[0-9]+}}xf32>
+// NATIVE: %[[RES:.+]] = llvm.fptrunc %[[IN]] : vector<{{[0-9]+}}xf32> to vector<{{[0-9]+}}xbf16>
+// NATIVE: llvm.store %[[RES]], {{.*}} : vector<{{[0-9]+}}xbf16>
+// NATIVE-NOT: llvm.lshr
+// EMULATED-NOT: llvm.fptrunc {{.*}} to vector<{{[0-9]+}}xbf16>
+// EMULATED: llvm.fcmp "une" %[[IN]], %[[IN]]
+// EMULATED: llvm.lshr
+// EMULATED: llvm.store {{.*}} vector<{{[0-9]+}}xi16>
