Add a fused multiply-add function

[carlosgmartin]

Feature request: Add a fused multiply-add function. It should use a hardware implementation if available.

This function computes x * y + z as if to infinite precision, rounding only once to fit the destination type.

Examples:

• https://github.com/smasher164/fma
• https://github.com/nschloe/pyfma
• Fused multiply-add: proposal to add math.fma() python/cpython#73468
• https://en.cppreference.com/w/cpp/numeric/math/fma
• https://www.gnu.org/software/c-intro-and-ref/manual/html_node/Fused-Multiply_002dAdd.html
• https://accurate-algorithms.readthedocs.io/en/latest/ch03fma.html

[jakevdp]

For what it's worth, XLA already fuses operations like this; here's the compiled HLO output on CPU:

import jax

def f(x, y, z):
  return x * y + z

print(jax.jit(f).lower(1, 2, 3).compile().as_text())

HloModule jit_f, is_scheduled=true, entry_computation_layout={(s32[], s32[], s32[])->s32[]}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}, frontend_attributes={fingerprint_before_lhs="94992b2905ce51de85765e674face20f"}

%fused_add (param_0.1: s32[], param_1.1: s32[], param_2: s32[]) -> s32[] {
  %param_1.1 = s32[] parameter(1)
  %param_2 = s32[] parameter(2)
  %multiply.1.1 = s32[] multiply(s32[] %param_1.1, s32[] %param_2), metadata={op_name="jit(f)/jit(main)/mul" source_file="<ipython-input-5-1b552238959a>" source_line=4}
  %param_0.1 = s32[] parameter(0)
  ROOT %add.1.1 = s32[] add(s32[] %multiply.1.1, s32[] %param_0.1), metadata={op_name="jit(f)/jit(main)/add" source_file="<ipython-input-5-1b552238959a>" source_line=4}
}

ENTRY %main.6 (Arg_0.1.0: s32[], Arg_1.2.0: s32[], Arg_2.3.0: s32[]) -> s32[] {
  %Arg_2.3.0 = s32[] parameter(2), metadata={op_name="z"}
  %Arg_1.2.0 = s32[] parameter(1), metadata={op_name="y"}
  %Arg_0.1.0 = s32[] parameter(0), metadata={op_name="x"}
  ROOT %loop_add_fusion = s32[] fusion(s32[] %Arg_2.3.0, s32[] %Arg_0.1.0, s32[] %Arg_1.2.0), kind=kLoop, calls=%fused_add, metadata={op_name="jit(f)/jit(main)/add" source_file="<ipython-input-5-1b552238959a>" source_line=4}
}

Is there something more you need than this kind of fusion in the compiler? Perhaps you're hoping that the compiler can lower this to some hardware-specific kernel?

[hawkinsp]

I'm actually supportive of adding fma because it has different rounding behavior.

However, the question is, what does one do on hardware without an fma? Error? Lower to the non-fma?

[carlosgmartin]

@jakevdp I don't think that guarantees the desired precision. Does it?

@hawkinsp Can it lower to a software implementation of fma? (That would preserve semantics.)

[hawkinsp]

@carlosgmartin Last I checked, emulating fma accurately on hardware without one is quite hard. But maybe there's a way I don't know about.

[carlosgmartin]

@hawkinsp In that case, how about we go with an error for now and fill in the software implementation later? (Some of the links in the OP might be useful for that.)

Another possibility is to fall back to standard multiply-add but also emit a warning (so the user is aware of the loss of precision).

[jakevdp]

I'm still confused about the request here: are you hoping for a JAX-level operator for the user to choose fused mult/add when they want it, or are you hoping for the compiler to use such a kernel automatically?

[carlosgmartin]

@jakevdp For this particular issue, I'd just like access to an fma function that guarantees the exact precision property.

[pearu]

Exposing fma as a jax/lax primitive function that uses hardware implementation when available would be very useful for implementing efficient and accurate algorithms in general. For instance, fma is a must-have primitive for polynomial evaluation that many math function algorithms rely on.

Currently, the accuracy of computations using jax is very sensitive to the choice of backend platform (CPU/GPU/TPU) and dtype combinations with respect to the availability of fma and enabling flush-to-zero (FTZ) mode. For example, pearu/functional_algorithms#114 (comment) contains accuracy tests of evaluating x * y + z on CPU/GPU using different FTZ modes and it clearly shows when XLA uses hardware fma and how enabling FTZ (default for CPU) affects computations when fma is disabled: x * y + z in float16 is always inaccurate (more than necessary) because float16-fma is not typically available in HW; x * y + z in float32/float64 is always accurate because XLA uses fma instructions available in the used HWs. In general, most AMD/Intel processor models not older than, say, 12 years, have fma instructions available for float32/float64. However, there exists HW when this is not true, e.g. on Apple GPUs float64 is emulated and on older CPUs. On the other hand, it is not always advantageous to blindly use fma instructions as it may lead to unexpected results (#26147, openxla/xla#22116).

Finally, here is a lax-based implementation of FMA emulation using floating-point operations only and the algorithm 7 from https://hal.science/hal-04624238 (generated using functional_algorithms Python tool):

import numpy
from jax._src import lax
from jax._src import numpy as jnp
from jax._src.api import jit
from jax._src.typing import Array, ArrayLike
from jax._src.numpy.util import promote_args_inexact

@jit
def fma_general(x: ArrayLike, y: ArrayLike, z: ArrayLike) -> Array:
    """Emulate fused add-multiply: x * y + z

        Applicability:
          x * y and x * y + z are finite.

        Accuracy:

          The ULP distance between the result and correct value is <= 1.

          fma(x, y, z) is exact (ULP distance is zero) for most inputs.
          If not, either

          - overflow occured in fma arithmetics (the return value is nan).
            Use fix_overflow to workaround overflow in Dekker product or 2Sum;
          - underflow occured in fma arithmetics. Use possibly_zero_z to
            fix z == 0 cases;
          - or the target uses flush-to-zero (FTZ) mode. Disable FTZ in target.

        Algorithms:
          a9: Algorithm 9 from https://hal.science/hal-04575249
          a7: Algorithm 7 from https://hal.science/hal-04624238
          a8: Algorithm 8 from https://hal.science/hal-04624238
          apmath: Using floating-point expansions.

    FA tracing parameters: fix_overflow=True assume_fma=False algorithm=a7 functional=True scale=True size=None possibly_zero_z=True
    """

    (
        x,
        y,
        z,
    ) = promote_args_inexact("fma_general", x, y, z)
    multiply_x_y = lax.mul(x, y)
    add_372 = lax.add(multiply_x_y, z)
    abs_x = lax.abs(x)
    x_max_lst = [numpy.float64(1.7976931080746007e308), numpy.float32(3.401993e38), numpy.float16(63500.0)]
    dtype_index_x = _np_dtypes.index(x.dtype.type)
    x_max = x_max_lst[dtype_index_x]
    constant_1 = numpy.array(1, dtype=x.dtype.type)
    lt_77 = lax.lt(abs_x, constant_1)
    N_lst = [numpy.float64(134217728.0), numpy.float32(4096.0), numpy.float16(64.0)]
    N = N_lst[dtype_index_x]
    invN_lst = [numpy.float64(7.450580596923828e-09), numpy.float32(0.00024414062), numpy.float16(0.01563)]
    invN = invN_lst[dtype_index_x]
    select_79 = jnp.where(lt_77, x, lax.mul(x, invN))
    multiply_N_select_79 = lax.mul(N, select_79)
    subtract_82 = lax.sub(multiply_N_select_79, lax.sub(multiply_N_select_79, select_79))
    constant_0 = numpy.array(0, dtype=x.dtype.type)
    negative_x_max = lax.neg(x_max)
    select_122 = jnp.where(
        lax.le(abs_x, x_max),
        jnp.where(lt_77, subtract_82, lax.mul(subtract_82, N)),
        jnp.where(lax.lt(x, constant_0), negative_x_max, x_max),
    )
    abs_y = lax.abs(y)
    lt_93 = lax.lt(abs_y, constant_1)
    select_95 = jnp.where(lt_93, y, lax.mul(y, invN))
    multiply_N_select_95 = lax.mul(N, select_95)
    subtract_98 = lax.sub(multiply_N_select_95, lax.sub(multiply_N_select_95, select_95))
    select_126 = jnp.where(
        lax.le(abs_y, x_max),
        jnp.where(lt_93, subtract_98, lax.mul(subtract_98, N)),
        jnp.where(lax.lt(y, constant_0), negative_x_max, x_max),
    )
    multiply_129 = lax.mul(select_122, select_126)
    largest_lst = [numpy.float64(1.7976931348623157e308), numpy.float32(3.4028235e38), numpy.float16(65500.0)]
    largest = largest_lst[dtype_index_x]
    subtract_y_select_126 = lax.sub(y, select_126)
    subtract_x_select_122 = lax.sub(x, select_122)
    subtract_373 = lax.sub(add_372, multiply_x_y)
    add_383 = lax.add(
        jnp.where(
            lax.le(lax.abs(multiply_129), largest),
            lax.add(
                lax.add(
                    lax.add(lax.add(lax.neg(multiply_x_y), multiply_129), lax.mul(select_122, subtract_y_select_126)),
                    lax.mul(subtract_x_select_122, select_126),
                ),
                lax.mul(subtract_x_select_122, subtract_y_select_126),
            ),
            constant_0,
        ),
        jnp.where(
            lax.le(lax.abs(subtract_373), largest),
            lax.add(lax.sub(multiply_x_y, lax.sub(add_372, subtract_373)), lax.sub(z, subtract_373)),
            constant_0,
        ),
    )
    add_384 = lax.add(add_372, add_383)
    subtract_385 = lax.sub(add_384, add_372)
    return lax.add(add_384, jnp.where(lax.le(lax.abs(subtract_385), largest), lax.sub(add_383, subtract_385), constant_0))

Its accuracy tests are available in pearu/functional_algorithms#114 (comment) that could be summarized as follows:

• float16 fma_general provides accuracy <= 1 ULP for all platforms
• float32/float64 fma_general provides accuracy <= 1 ULP for GPU and for CPU with FTZ disabled. Enabling FTZ on CPU ruins the floating-point algorithm for inputs that lead to subnormals within computations that are unexpectedly replaced with zeros.
• When fma is available in HW, one should use fma instructions instead of fma_general as it will be naturally much more efficient. The performance of emulated FMA as above is about 2-6 times smaller than the performance of evaluating x * y + z directly (a subtle note: XLA silently using fma instructions actually transform fma_general to a smaller set of instructions than one would expect from reading fma_general).