Relaxed BFloat16 Dot Product instruction

[Maratyszcza]

Introduction

BFloat16 is a 16-bit floating-point format that represents the IEEE FP32 numbers truncated to the high 16 bits. BFloat16 numbers have the same exponential range as the IEEE FP32 numbers, but way fewer mantissa bits, and is primarily used in neural network computations which are very tolerable to reduced precision. Despite being a non-IEEE format, BFloat16 computations are supported in the recent and upcoming processors from Intel (Cooper Lake, Alder Lake, Sapphire Rapids), AMD (Zen 4), and ARM (Cortex-A510, Cortex-A710, Cortex-X2), and can be efficiently emulated on older hardware by zero-padding BFloat16 numbers with zeroes to convert to IEEE FP32.

Both x86 and ARM BFloat16 extensions provide instructions to compute an FP32 dot product of two two-element BFloat16 elements with addition to the FP32 accumulator. However, there're some differences in details:

x86 introduced support for BFloat16 computations with the AVX512-BF16 extension, and the BFloat16 dot production functionality is exposed via the VDPBF16PS (Dot Product of BF16 Pairs Accumulated into Packed Single Precision) instruction. The instruction VDPBF16PS dest, src1, src2 computes

dest.fp32[i] = FMA(cast<fp32>(src1.bf16[2*i]), cast<fp32>(src2.bf16[2*i]), dest.fp32[i])
dest.fp32[i] = FMA(cast<fp32>(src1.bf16[2*i+1]), cast<fp32>(src2.bf16[2*i+1]), dest.fp32[i])

for each 32-bit lane i in the accumulator register dest. Additionally, denormalized numbers in inputs are treated as zeroes and denormalized numbers in outputs are replaced with zeroes.

ARM added support for BFloat16 computations with the ARMv8.2-A BFDOT instruction (mandatory with ARMv8.6-A), which implements the following computation:

temp.fp32[i] = cast<fp32>(src1.bf16[2*i]) * cast<fp32>(src2.bf16[2*i])
temp2.fp32[i] = cast<fp32>(src1.bf16[2*i+1]) * cast<fp32>(src2.bf16[2*i+1])
temp.fp32[i] = temp.fp32[i] + temp2.fp32[i]
dest.fp32[i] = dest.fp32[i] + temp.fp32[i]

where all multiplications and additions are unfused and use non-IEEE Round-to-Odd rounding mode. Additionally, denormalized numbers in inputs are treated as zeroes, denormalized numbers in outputs are replaced with zeroes.

ARMv9.2-A introduced additional "Extended BFloat16" (EBF16) mode, which modifies the behavior of BFDOT as follows:

temp.fp32[i] = cast<fp32>(src1.bf16[2*i] * src2.bf16[2*i] + src1.bf16[2*i+1] * src2.bf16[2*i+1])
dest.fp32[i] = dest.fp32[i] + temp.fp32[i]

where temp.fp32[i] is calculated as a fused sum-of-products operations with only a single (standard, Round-to-Nearest-Even) rounding at the end.

Furthermore, ARMv8.6-A introduced BFMLALB/BFMLALT instructions which perform widening fused multiply-add of even/odd BFloat16 elements to an IEEE FP32 accumulator. A pair of these instructions can be similarly used to implement BFloat16 Dot Product primitive.

What are the instructions being proposed?

I propose a 2-element dot product with accumulation instruction with BFloat16 input elements and FP32 accumulator & output elements. The instruction has relaxed semantics to allow lowering to native VDPBF16PS (x86) and BFDOT (ARM) instructions, as well as FMA instructions where available.

I suggest f32x4.relaxed_dot_bf16x8_add_f32x4 as a tentative name for the instruction. In a sense it is a floating-point equivalent of the i32x4.dot_i16x8_add_s instruction proposed in WebAssembly/simd#127.

What are the semantics of these instructions?

y = f32x4.relaxed_dot_bf16x8_add_f32(a, b, c) computes

y.fp32[i] = y.fp32[i] + cast<fp32>(a.bf16[2*i]) * cast<fp32>(b.bf16[2*i]) + cast<fp32>(a.bf16[2*i+1]) * cast<fp32>(b.bf16[2*i+1])

The relaxed nature of the instruction manifests in several allowable options:

Evaluation ordering options

We permit two evaluation orders for the computation:

• Compute cast<fp32>(a.bf16[2*i]) * cast<fp32>(b.bf16[2*i]) + cast<fp32>(a.bf16[2*i+1]) * cast<fp32>(b.bf16[2*i+1]) in the first step, then compute the sum with y.fp32[i] in the second step.
• Compute y.fp32[i] += cast<fp32>(a.bf16[2*i]) * cast<fp32>(b.bf16[2*i]) in the first step, then compute y.fp32[i] += cast<fp32>(a.bf16[2*i+1]) * cast<fp32>(b.bf16[2*i+1]) in the second step.

Fusion options

• Either both steps of the computation are fused, or both are unfused

Rounding options

• Operations that comprise the computation are rounded with either Round-to-Nearest-Even or Round-to-Odd rounding modes.

How will these instructions be implemented?

x86/x86-64 processors with AVX512-BF16 instruction set

• f32x4.relaxed_dot_bf16x8_add_f32x4
  • c = f32x4.relaxed_dot_bf16x8_add_f32x4(a, b, c) is lowered to VDPBF16PS xmm_c, xmm_a, xmm_b
  • y = f32x4.relaxed_dot_bf16x8_add_f32x4(a, b, c) is lowered to VMOVDQA xmm_y, xmm_c + VDPBF16PS xmm_c, xmm_a, xmm_b

ARM64 processors with BF16 extension (using BFDOT instructions)

• f32x4.relaxed_dot_bf16x8_add_f32x4
  • y = f32x4.relaxed_dot_bf16x8_add_f32x4(a, b, c) is lowered to MOV Vy.16B, Vc.16B + BFDOT Vy.4S, Va.8H, Vb.8H
  • c = f32x4.relaxed_dot_bf16x8_add_f32x4(a, b, c) is lowered to BFDOT Vc.4S, Va.8H, Vb.8H

ARM64 processors with BF16 extension (using BFMLALB/BFMLALT instructions)

• f32x4.relaxed_dot_bf16x8_add_f32x4
  • y = f32x4.relaxed_dot_bf16x8_add_f32x4(a, b, c) is lowered to MOV Vy.16B, Vc.16B + BFMLALB Vy.4S, Va.4H, Vb.4H + BFMLALT Vy.4S, Va.8H, Vb.8H
  • c = f32x4.relaxed_dot_bf16x8_add_f32x4(a, b, c) is lowered to BFMLALB Vc.4S, Va.4H, Vb.4H + BFMLALT Vc.4S, Va.8H, Vb.8H

Reference lowering through the WAsm Relaxed SIMD instruction set

• f32x4.relaxed_dot_bf16x8_add_f32x4
  • y = f32x4.relaxed_dot_bf16x8_add_f32x4(a, b, c) is lowered to:
    • const v128_t a_lo = wasm_i32x4_shl(a, 16)
    • const v128_t b_lo = wasm_i32x4_shl(b, 16)
    • const v128_t a_hi = wasm_v128_and(a, wasm_i32x4_const_splat(0xFFFF0000))
    • const v128_t b_hi = wasm_v128_and(b, wasm_i32x4_const_splat(0xFFFF0000))
    • y = f32x4.relaxed_fma(a_lo, b_lo, c)
    • y = f32x4.relaxed_fma(a_hi, b_hi, y)

How does behavior differ across processors? What new fingerprinting surfaces will be exposed?

The use of AVX512-BF16 VDPBF16PS instruction can be detected through testing for behavior on denormal inputs and outputs.
The use of ARM BFDOT instruction can be detected through testing for behavior on denormal inputs and outputs, or testing for Round-to-Odd rounding. However, an ARM implementation using a pair of BFMLALB/BFMLALT instructions would be indistinguishable from a generic lowering using FMA instructions (but probably slower than BFDOT lowering).

What use cases are there?

• Matrix-Matrix Multiplication in Intel oneDNN library

[ngzhian]

@yurydelendik @dtig
please add some engine implementor feedback on this instruction.
specifically, this instruction requires AVX512-BF16 extension on the x86 side, and ARMv8.6-A with BF16 extension on the ARM side. Please comment on possible support for these instructions.

[yurydelendik]

(Typo to fix: in reference lowering code y = f32x4.relaxed_fma(a_lo, b_lo, y) -> y = f32x4.relaxed_fma(a_lo, b_lo, c))

[abrown]

I wanted to also comment on what we discussed about bf16 conversions to and from fp32 (correct me if anything is incorrect): @Maratyszcza is proposing that the native conversion instructions (e.g., VCVTNE2PS2BF16) not be included in relaxed SIMD because the conversion can be done by zero-extending the bf16 to fill the bottom 16 bits of the fp32 (or truncating, in the opposite-direction conversion), since this will have the same semantics as round-to-zero. @Maratyszcza made the point that these conversions, event though lossy, are rare in the domain he works in (?).

[Maratyszcza]

BFloat16 numbers are usually only used for storage, up-casted to IEEE FP32 for computations. The proposed f32x4.relaxed_dot_bf16x8_add_f32x4 instruction too does internal computations in FP32 and produce FP32 result. Only the final result need to be down-casted to BFloat16 for store to memory, this downcast can be done by truncating the high 16 bits of the FP32 number, which is sufficiently accurate if it happens relatively rarely, e.g. at operator boundaries in NN computations.

[justinmichaud]

I would like to add some implementer feedback from JSC/WebKit.

There is a lot of interesting feedback in this thread, and I have a few questions.

Some of this was mentioned in the meeting, but do we have any data about:

• The performance/speedup of using this instruction, particularly on non-AVX512 intel chips?
• How often we expect this instruction to be useful? As in, how general is it?
• Are there any other mainstream programming languages or libraries that support it?

In general, we are concerned about instructions that might make deterministic fingerprinting much easier, and also about how they might cause compatibility/interoperability issues for users who browse the web on less-common CPU architectures.

An important part of the Relaxed-SIMD proposal for us in general is the fact that it justifies this risk with very compelling performance data. We would not consider implementing BFloat16 unless we have proof that this would speed up an entire class of reasonably popular programs.

Overall, since Relaxed-SIMD is in stage 3, we think that this change should be put in a separate proposal.

[Maratyszcza]

Thanks for the feedback @justinmichaud (and great to see JSC/WebKit being more active in WAsm SIMD)!

The performance/speedup of using this instruction, particularly on non-AVX512 intel chips?

Performance on x86 CPUs without AVX512-BF16 extension would be similar to software emulation of the f32x4.relaxed_dot_bf16x8_add_f32x4 with Relaxed SIMD instructions. It would still be faster than end-to-end FP32 processing if the computation is bandwidth-bound on the machine. Note that when BF16 format was originally introduced, it was handled purely in software; hardware support arrived years later.

On Galaxy S22 with Snapdragon 8 Gen 1 I see throughput of 4 BFDOT/cycle (Cortex-X2 cores) / 2 BFDOT/cycle (Cortex-A710 cores) / 1 BFDOT/cycle (Cortex-A510 cores). This is the same throughput as for single-precision FMLA instructions, but each BFDOT instruction produce twice more results, leading to twice higher peak performance.

How often we expect this instruction to be useful? As in, how general is it?

This instruction is known to be helpful for neural network computations (both for inference and training). Background Effects in Google Meet and Teachable Machine are examples of applications using such functionality.

Are there any other mainstream programming languages or libraries that support it?

The native VDPBF16PS and BFDOT instructions are exposed in C/C++ as intrinsic functions _mm512_dpbf16_ps and vbfdotq_f32 (+variants of these). Among machine learning frameworks, both TensorFlow and PyTorch support BFloat16 data type (with emulation in lieu of hardware support).

we are concerned about instructions that might make deterministic fingerprinting much easier
Note that the implementation using BFMLALB/BFMLALT instructions on ARM is bitwise-compatible with pre-BF16 emulation path.

Overall, since Relaxed-SIMD is in stage 3, we think that this change should be put in a separate proposal.

Unlike native platforms, WebAssembly doesn't have the capability to detect supported WebAssembly ISA extensions at run-time: if a WAsm engine doesn't support any instruction in a WAsm module, the whole will be rejected, even if the instruction is never executed in runtime. Unfortunately, this means that WebAssembly extensions can't be as granular as native ISA extensions: developers have to rebuild the whole webapp for each set of WAsm extensions they target, and having an extension with a single instruction in it is overburdening for the developers.

Please note that it is permissible to add new instructions at stage 3 (WebAssembly SIMD added dozens of instructions while at Stage 3). It is Stage 4 is where the spec is frozen.

[justinmichaud]

Thanks for the fast response!

On Galaxy S22 with Snapdragon 8 Gen 1 I see throughput of 4 BFDOT/cycle (Cortex-X2 cores) / 2 BFDOT/cycle (Cortex-A710 cores) / 1 BFDOT/cycle (Cortex-A510 cores). This is the same throughput as for single-precision FMLA instructions, but each BFDOT instruction produce twice more results, leading to twice higher peak performance.

Thanks! Is there a real benchmark or library that is available to be compiled with this instruction vs just the normal set of relaxed-simd instructions? The two examples that you linked seem interesting, but I don't see how they could be ported today.

I am also wondering if there is an implementation using BFMLALB/BFMLALT that is complete enough to collect real-world performance numbers, since that is the only implementation we would consider on ARM.

We really need to see data in the context of a full benchmark or application here, because on our CPUs instruction throughput doesn't usually indicate general performance.

Note that the implementation using BFMLALB/BFMLALT instructions on ARM is bitwise-compatible with pre-BF16 emulation path.

This is good to know. So if I am understanding things correctly, the current state of the world is that this is software emulated on every single chip that is available today except for ARMv8+, and ARMv8+ can only be distinguished from emulation by performance (which is fine)?

Please note that it is permissible to add new instructions at stage 3 (WebAssembly SIMD added dozens of instructions while at Stage 3). It is Stage 4 is where the spec is frozen.

This makes sense, but I would just like to note that this instruction seems like a pretty big departure from the others in the proposal. In addition, it seems like there are almost no desktop CPUs that natively support this instruction yet, so maybe it is too soon to standardize?

[Maratyszcza]

Is there a real benchmark or library that is available to be compiled with this instruction vs just the normal set of relaxed-simd instructions?

Alibaba's MNN and Tencent's NCNN support BF16 computations (via software conversions to/from FP32), but AFAICT neither use the ARM BF16 extension yet. It is also worth noting that native software targeting ARM BF16 is unlikely to use BFDOT or BFMLALB + BFMLALT pair, because BF16 extension also includes a more powerful BFloat16 matrix multiplication instruction BFMMLA. The latter, however, is too exotic and hard to emulate to be exposed in WebAssembly.

My plan is to implement BF16 GEMM (matrix-matrix multiplication) kernels in XNNPACK using generic NEON / BFDOT / BFMLALB + BFMLALT pair to evaluate the potential of the f32x4.relaxed_dot_bf16x8_add_f32x4 instruction. GEMM is responsible for 50-80% of runtime in the convolutional neural network (now-common for Computer Vision tasks) and often 90%+ of runtime in the transformer-type neural networks (now-common in Natural Language Processing and replacing CNNs for Computer Vision), so it is a good proxy for the overall neural network inference performance.

So if I am understanding things correctly, the current state of the world is that this is software emulated on every single chip that is available today except for ARMv8+, and ARMv8+ can only be distinguished from emulation by performance (which is fine)?

This is correct. To be specific, BF16 extension on ARM was introduced as optional in ARMv8.2-A and became mandatory in ARMv8.6-A.

This makes sense, but I would just like to note that this instruction seems like a pretty big departure from the others in the proposal. In addition, it seems like there are almost no desktop CPUs that natively support this instruction yet, so maybe it is too soon to standardize?

I see little risk in adding this instruction, as BF16 extension is mandatory in ARMv8.6-A (and ARMv9), and thus eventually all ARM-based computers will have it. The latest Windows-based laptops (e.g. Lenovo ThinkPad X13s) already support BF16 extension. In the desktop x86 space, AMD Zen 4 is reported to support AVX512-BF16, so Intel will have to match to compete.

However, if we don't add this instruction, WebAssembly SIMD will miss out on the nearly 2X speedup for AI models, which would put privacy-focused technologies like Federated Learning as risk.

[penzn]

I see little risk in adding this instruction, as BF16 extension is mandatory in ARMv8.6-A (and ARMv9), and thus eventually all ARM-based computers will have it. The latest Windows-based laptops (e.g. Lenovo ThinkPad X13s) already support BF16 extension. In the desktop x86 space, AMD Zen 4 is reported to support AVX512-BF16, so Intel will have to match to compete.

To reiterate an important point to make sure it doesn't get lost in a long discussion: we need to hear back from engine maintainers before we conclude those instructions are available to the engines. Implementing support for new ISA extension requires implementing new decode, operands, etc, which comes with maintenance costs. I have brought this up before and it did not look like there was interest (I personally would welcome this changing 😉).

[yurydelendik]

please add some engine implementor feedback on this instruction. specifically, this instruction requires AVX512-BF16 extension on the x86 side, and ARMv8.6-A with BF16 extension on the ARM side. Please comment on possible support for these instructions.

I opened the issue at https://bugzilla.mozilla.org/show_bug.cgi?id=1778751 with intent to implement the reference lowering in SpiderMonkey. We feel that there is a need to support BFloat16 for matrix multiplication / ML use cases. In the future we will adjust the solution for ARM (and AVX) BF16 extensions as they become popular.

[yurydelendik]

There is inconsistency in arguments ordering. For fma(a, b, c) we are doing result = a + b * c, but for relaxed_dot_bf16x8_add_f32(a, b, c) we changed the accumulator placement result = dot(a, b) + c. Shall we keep a as running sum?

(Also reference lowering has typo related to this)

[Maratyszcza]

@yurydelendik Good point! However, it is fma(a, b, c) that is behaves inconsistently:

• Fused Multiply-Add as specified has addition first, then multiplication
• fma(a, b, c) is inconsistent with the same-named C/C++ functions

Lets have discussion in the FMA proposal #27