Three-State Logic for FMA

You would think that whether or not a backend has true native Fused Multiply-Add (FMA) instructions would be pretty cut and dry, right? Until recently, Thermite types indicated if they have true native FMA instructions via a HAS_TRUE_FMA associated const bool. However, that was insufficient, and broken at worst.

# Why Indicate True FMA

It's quite helpful to have algorithms decide on certain paths or optimizations, at compile-time, based on if true FMA is available. For one, it's often as fast as a single multiply while doing both multiply and addition, so that alone accelerates many functions. It also only rounds once, so it can bypass extra effort required for that, such as some compensated arithmetic.

if const { V::HAS_TRUE_FMA } { /* ... */ }

was used all over Thermite's math libraries for these kind of optimizations.

It's also worth knowing Thermite has two flavors of FMA, which are important:

• mul_add and family are always correctly rounded fused multiply-add, regardless of backend. If there is no hardware FMA, it falls back to a bit-identical vectorized polyfill, which can be quite expensive.
• mul_adde and family are the opportunistic forms. They are fused when FMA is a single instruction, but plain a * b + c multiply+add where it isn't. Always fast, allowed to differ between backends.

mul_adde and family are used almost everywhere. Might as well opt for the fastest form available.

# WebAssembly

The WASM relaxed-simd proposal added relaxed_madd and relaxed_nmadd instructions. "Relaxed" meaning they act like mul_adde/nmul_adde, where it will use true FMA if available but otherwise fallback to a*b+c for performance. Nothing is guaranteed about which one you get, the engine is free to do whatever it wants. The only promise is that the choice is fixed for the lifetime of the module instance, which is what makes detecting it once at startup sound.

However, one change I made recently was to opportunistically enable true FMA on WASM if it's detected at startup. Using some WASM tricks, I inject a bit of code that runs before main(), updating a static mut global variable:

#[used]
#[unsafe(link_section = ".init_array")]
static INIT_WASM_RELAXED_FMA: extern "C" fn() = {
    extern "C" fn init() { detect_relaxed_fma(); }
    init
};

The linker collects .init_array entries into __wasm_call_ctors, which runs before user code.

mul_add and family then branch on that flag at runtime to opt into true FMA. Both arms are bit-identical, so that only trades speed.

# What Broke

mul_add was fine, that was always guaranteed to be the correctly rounded fused multiply-add. However, mul_adde behaves differently. The default behavior of mul_adde was to check V::HAS_TRUE_FMA and either call mul_add, or do a*b+c, determined at compile time.

So when mul_adde for WebAssembly specifically was overridden to ignore V::HAS_TRUE_FMA and lower straight to relaxed_madd, leaving the engine to decide, we had cases where math algorithms chose the V::HAS_TRUE_FMA == false path but mul_adde used true FMA internally.

Normally that's fine, but notably "sum/difference of products" breaks entirely.

For example, multiplying a quaternion by its own conjugate should give exactly $(0, 0, 0, |q|^2)$. The fast quaternion product computes each lane as a sum of four products, two of them folded into mul_addes:

let sum12 = Self::mul_adde(x, rhs_x_signed, Self::mul(w, rhs));
let sum34 = Self::mul_adde(z, rhs_z_signed, Self::mul(y, rhs_y_signed));

For $q \cdot \bar{q}$ the vector lanes are made of pairs like $t - t$. If every product is rounded, that's $\mathrm{fl}(t) - \mathrm{fl}(t) = 0$, exact, for free. If mul_adde fuses, one side of the pair is the exact product and the other is rounded, and what survives is the rounding error:

q = (0.1, 0.2, 0.3, 0.927...)
q * conj(q) = [1.5188877683131223e-19, 0.0, -5.551115123125783e-17, 1.0]

The quaternion product has a fast arm and an accurate arm that compensates for this. On backends with no FMA, the accurate path routed back to the fast arm, on the reasoning that all-rounded products already cancel exactly. Correct on SSE2. On WASM, with HAS_TRUE_FMA = false and an engine that fuses, it produced exactly the garbage above, in the path whose documentation promises exact zeros. The 4x4 determinant and inverse minors had the same defect.

In effect, HAS_TRUE_FMA could not be trusted. Not because the value was wrong, but because bool has no way to say "I don't know".

# Tribool

Many years ago I created the rust tribool crate, based off of Boost's Tribool class. It's basically just an enum:

pub enum Tribool {
    False,
    True,
    Indeterminate, // Unknown
}

with various three-state logic methods on it. Well, what would you call the state in which we don't know if FMA is present or not, other than Indeterminate?

The crate fit perfectly, so I set out and replaced V::HAS_TRUE_FMA with

const HAS_NATIVE_FMA: Tribool;

The funny part is after the migration, no algorithm in Thermite mentions Indeterminate, instead preferring one of:

// "Is FMA cheap here? i.e.: dedicated instruction"
if const { matches!(V::HAS_NATIVE_FMA, tribool::True) } { /* ... */ }

// "Could two products round differently from each other?"
// Note the boolean NOT
if const { !matches!(V::HAS_NATIVE_FMA, tribool::False) } { /* ... */ }

As it turns out, the old bool const was close, but not exactly what was needed for informing algorithms about the FMA behavior. Depending on the situation we need to know if FMA is cheap (i.e., a single instruction), or just if we can rely on certain rounding behaviors or not.

# Composites

One more place the old bool was wrong was for composite types, such as Dual or Complex numbers, which reported HAS_TRUE_FMA = false unconditionally. However, if the question is "Is mul_adde faster than a*b+c?", that's more interesting.

Nearly every consumer of HAS_TRUE_FMA cared because of performance reasons, not because of the single rounding. Even then, making use of FMA for Dual/Complex's own mul_adde does improve precision some. Therefore, these now simply forward the inner type's HAS_NATIVE_FMA.

# Conclusion

In the end, most algorithms didn't really care, but when !false != true, making use of Tribool is both useful and cool.