[Issue]: fx.exp2 lowers to an OCML libcall instead of native v_exp_f32, and `fastmath` is silently dropped on that path

[yiijin]

Summary

Comparing a raw rocdl.exp2 escape hatch against the equivalent fx.exp2 call, two issues showed up:

1. fx.exp2 lowers to llvm.call @__ocml_exp2_f32, not the native v_exp_f32. The inlined OCML
   body wraps every call site in denormal/overflow range handling, costing +5 VALU and +3 VGPR per
   call site, for numerically identical results on the input range the kernel actually uses.
2. fastmath="fast" cannot be used to opt out. It is accepted by the Python API and reaches MLIR
   as math.exp2 %x fastmath<fast> : f32, but is dropped in the math-to-rocdl lowering — the
   emitted llvm.call carries no fast-math flags and the final ISA is byte-identical to the
   non-fastmath version. There is currently no way to reach the native instruction through the
   public flydsl.expr API.

Repro

Elementwise C = exp2(A) with 4 exp2 call sites per thread. Input is -rand() * 20, i.e.
non-positive exponents only — the range GDN decay terms and softmax scores live in, where OCML's
range handling is provably unreachable. The three variants differ only in this expression:

def _exp2(x, v):
    if v == "rocdl":                                                     # A: current workaround
        return fx.Float32(fx.rocdl.exp2(fx.Float32.ir_type, x.ir_value()))
    if v == "fx":                                                        # B: public API
        return fx.exp2(x)
    return fx.exp2(x, fastmath="fast")                                   # C: public API + fastmath

Save the script below as min_exp2.py, then dump the ISA for each variant:

for v in "rocdl A" "fx B" "fx_fast C"; do set -- $v
  FLYDSL_DUMP_IR=1 FLYDSL_RUNTIME_ENABLE_CACHE=0 FLYDSL_DUMP_DIR=/tmp/isa_$2 python3 min_exp2.py $1
done

All three print max_abs_err=0.000e+00 against torch.exp2.

min_exp2.py

import sys

import torch

import flydsl.compiler as flyc
import flydsl.expr as fx

VARIANT = sys.argv[1] if len(sys.argv) > 1 else "rocdl"
N = 4


def _exp2(x, v):
    if v == "rocdl":
        return fx.Float32(fx.rocdl.exp2(fx.Float32.ir_type, x.ir_value()))
    if v == "fx":
        return fx.exp2(x)
    return fx.exp2(x, fastmath="fast")


@flyc.kernel
def exp2_kernel(A: fx.Tensor, C: fx.Tensor, variant: fx.Constexpr):
    i = fx.block_idx.x * 256 + fx.thread_idx.x
    for j in fx.range_constexpr(N):
        C[i, j] = _exp2(A[i, j], variant)


@flyc.jit
def run(A: fx.Tensor, C: fx.Tensor, variant: fx.Constexpr, stream: fx.Stream = fx.Stream(None)):
    M, _ = A.shape.unpack()
    exp2_kernel(A, C, variant).launch(grid=(M // 256, 1, 1), block=(256, 1, 1), stream=stream)


M = 1024
A = -torch.rand(M, N, dtype=torch.float32, device="cuda") * 20.0
C = torch.zeros(M, N, dtype=torch.float32, device="cuda")
run(A, C, VARIANT, stream=torch.cuda.Stream())
torch.cuda.synchronize()
print(f"variant={VARIANT} max_abs_err={(C - torch.exp2(A)).abs().max().item():.3e}")

Issue 1 — OCML libcall costs +5 VALU per call site

Opcode counts and .amdhsa_next_free_vgpr from the dumped 21_final_isa.s:

metric	A rocdl.exp2	B fx.exp2	C fx.exp2(fastmath="fast")
total instructions	35	62	62
v_exp_f32	4	4	4
v_ldexp_f32 / v_cmp_gt_f32 / v_cndmask_b32 / v_add_f32	0 / 0 / 0 / 0	4 / 4 / 8 / 4	4 / 4 / 8 / 4
VGPR	6	9	9

The v_exp_f32 that does the actual work is 4 in all three. The delta is 20 range-handling VALU
(5 per call site) plus 4 s_nop and 3 instructions of fixed overhead.

Issue 2 — fastmath is dropped in the math-to-rocdl lowering

cmp /tmp/isa_B/exp2_kernel_0/21_final_isa.s /tmp/isa_C/... reports no difference, so
fastmath="fast" has no effect on codegen at all. Counting the attribute per pass stage shows where
it is lost — it survives to stage 09 and disappears in stage 10, which contains the math-to-rocdl
conversion:

09_canonicalize.mlir                                 fastmath=4  math.exp2=4  ocml=0
10_convert_scf_to_cf_cse_convert_gpu_to_rocdl.mlir   fastmath=0  math.exp2=0  ocml=5

// 09_canonicalize.mlir
%10 = math.exp2 %9 fastmath<fast> : f32

// 10_convert_scf_to_cf_cse_convert_gpu_to_rocdl.mlir -- attribute gone
%18 = llvm.call @__ocml_exp2_f32(%17) : (f32) -> f32

Final LLVM IR (20_llvm_ir.ll):

; A: rocdl.exp2
%19 = call float @llvm.amdgcn.exp2.f32(float %18)

; C: fx.exp2(x, fastmath="fast") -- no fast flag
%19 = call float @__ocml_exp2_f32(float %18)

Environment

• OS: Ubuntu 24.04.4 LTS, kernel 6.8.0-38-generic, x86_64
• CPU: 2× AMD EPYC 9655 96-Core
• GPU: gfx950 (CDNA4), 256 CU, 288 GiB HBM
• ROCm 7.2.4 (HIP 7.2.53211, AMD clang 22.0.0git roc-7.2.4)
• flydsl 0.3.0 (wheel), PyTorch 2.10.0+rocm7.2.4

Context

Found while reviewing aiter's gdn_prepare FlyDSL kernel, which keeps a raw rocdl.exp2 helper
instead of calling fx.exp2; a reviewer asked whether that was still necessary. That kernel (17 exp2
call sites) matches the model above: 1686 → 1774 instructions (+88 = 17 × 5 + 3), VGPR 61 → 64, and
again byte-identical ISA with and without fastmath. There the latency difference is within noise and
occupancy does not change, so the cost is instruction budget and register pressure; in exp2-dense
inner loops (softmax / flash-attention) the +5 per call site lands directly in the hot path.

FlyDSL's own kernels already bypass the public API for this reason: kernels/common/utils.py:22 calls
llvm.call_intrinsic("llvm.amdgcn.exp2.f32") (wrapped by exp2_f32_fast at :32),
kernels/attention/flash_attn_utils.py:3746 uses rocdl.exp2, and
kernels/attention/swa_gfx950.py:334 uses hw_exp2_v16 / hw_exp2_scalar.

Expected behaviour / suggested fix

1. Fix the math-to-rocdl lowering: when math.exp2 carries a fastmath flag including afn (or
   fast), lower it to rocdl.exp2 / llvm.amdgcn.exp2.f32 rather than the OCML call, so
   fx.exp2(x, fastmath="fast") produces the same ISA as rocdl.exp2.
2. At minimum, stop dropping fastmath silently: propagate it onto the emitted llvm.call
   (call fast float @__ocml_exp2_f32) so LLVM can do the substitution itself. A public API parameter
   that is accepted and then has literally no effect is the more surprising part of this report.

Promoting exp2_f32_fast to a public flydsl.expr API (e.g. fx.exp2(..., approx=True)) would also
help: it currently lives under kernels.*, and some distribution channels do not ship it — aiter's
bundled FlyDSL snapshot has no kernels/common/utils.py, so downstream code ends up writing a
try/except ModuleNotFoundError fallback to rocdl.exp2.