[Bounty $2,500] Optimise/improve accuracy for pow(x, y) fp32 (non-integer exponent)

Tenstorrent Bounties

Issue ID: I_kwDOI9Wqc88AAAABIZbW5Q

Summary

The ttnn.pow(x, y) fp32 implementation (the exp_21f-based path added in PR #25025, which closed #23529) was a big improvement over the old kernel, but it still lost up to 27 ULP for non-integer exponents, while integer exponents stayed bit-exact (0 ULP). The op is the composite y * log2(x) followed by 2^(...), so the fp32 log2 and exp2 rounding errors compound. The CogVideo case that motivated #23529 (10000 ** 1.7984) was exactly one of those weak spots.

This PR is a follow-up accuracy improvement on top of the accepted #25025 fix, in the same class as the transcendental accuracy bounties exp2 #46006, expm1 #45046, atan2 #41029. It brings all tested non-integer exponents to within 3 ULP — including the long-mantissa 1.7984 CogVideo case (19 → 2 ULP) — while keeping integer exponents bit-exact.

Current state and fix (measured on Wormhole B0)

Dense sweep, 1000 log-spaced bases x ∈ [0.5, 50000], fp32 input, fp32_dest_acc_en=true, math_approx_mode=false. Golden = torch.pow (fp32). Both columns are hardware-measured on Wormhole B0 with a confirmed clean JIT recompile: “before” from the unmodified #25025 kernel, “after” with this fix applied.

exponent y Max ULP (before) Max ULP (after) after: points > 3 ULP after: points > 10 ULP
2.0 (integer) 0 0 0 / 1000 0 / 1000
3.0 (integer) 0 0 0 / 1000 0 / 1000
0.5 4 1 0 / 1000 0 / 1000
1.5 19 2 0 / 1000 0 / 1000
1.7984 (CogVideo) 19 2 0 / 1000 0 / 1000
2.5 27 3 0 / 1000 0 / 1000

2.5 has a single point at exactly 3.0 ULP (p99 = 2.0); every other non-integer exponent is ≤ 2 ULP. bf16 is unaffected (≤ 2 ULP, unchanged).

Performance (DEVICE KERNEL DURATION [ns], single 32×32 fp32 tile, 53 calls, measured via tracy): after ≈ 7990 ns steady-state (avg 8024, range 7918–8299). The fix adds a branch-free two-sum, a Veltkamp split of one operand, and one extra Newton-Raphson step — a handful of SFPU ops, no extra passes over the tile. The unmodified #25025 kernel was not separately profiled; a strict before/after should stash the four headers and re-run the same driver.

Accuracy sweep for pow(x, 2.5) fp32, before vs after the fix. Before: max 27 ULP, error grows with x (compounded log2 + exp2 rounding). After: max 3 ULP, mean 0.36, 0 points > 3 ULP.

Before (#25025 implementation):

After (this fix):

Root cause: the fp32 path (_sfpu_binary_power_f32_ / _sfpu_unary_power_61f_updated_ in ckernel_sfpu_binary_pow.h / ckernel_sfpu_unary_power.h) computes x^y = 2^(y · log2(x)). Integer y hits an exact iterative multiply path (0 ULP), but non-integer y goes through the log2 + exp2 composite. The killer was that z = y·log2(x) was collapsed into a single fp32 before 2^z: for large |x| the large integer part of log2(x) squeezes out the fractional mantissa bits, so up to ~20 bits of the fractional exponent were lost → 20–27 ULP. A secondary floor came from the reciprocal 1/(m+1) in the log2 atanh series, which had a ~2 ULP residual after two Newton-Raphson steps.

Fix (implemented, this PR)

Three changes, applied consistently to the unary and binary fp32 paths on both Wormhole B0 and Blackhole:

  1. Carry z = y·log2(x) as an unevaluated double-float (z_hi, z_lo) instead of one fp32 — z_hi the large integer-valued term, z_lo the fractional term. A register-frugal helper _sfpu_pow2_f32_accurate_hilo_(z_hi, z_lo) uses Knuth TwoSum (branch-free, exact) to form s = round(z_hi+z_lo) and its exact residual e, reduces s by k = round(s) (s − k exact by Sterbenz), adds back e, then feeds the reduced small argument into the existing Cody-Waite 2^f polynomial and scales by 2^k via setexp. The tensor–tensor path feeds the same reduction into _sfpu_exp_fp32_accurate_.

  2. Veltkamp split of pow (2^12+1 splitter) so the pow · exp_f32 product is exact. exp_f32 is the base’s integer exponent (a small integer); splitting pow into a 12-bit hi/lo pair makes both pow_hi·exp_f32 and pow_lo·exp_f32 exact (12-bit half × ≤7-bit integer fits a 24-bit significand), so the low partial product rides in z_lo and no bit of the large integer term is lost. This is what fixed the long-mantissa 1.7984 case. Only pow is split (not a full Dekker two-product on log2(base)) to keep the simultaneously-live vector count within the SFPU register-allocator budget — a full two-product on the 16×-inlined kernel triggers a maximum number of generated reload insns ICE on Wormhole.

  3. Third Newton-Raphson iteration on 1/(m+1) in the log2 series, driving the reciprocal to full fp32 precision and removing the residual floor that kept 1.5/2.5 above threshold.

Success criteria

  • fp32 pow(x, y), non-integer y: Max ULP ≤ 3 on x ∈ [0.5, 50000] for all tested exponents (including long-mantissa 1.7984). :white_check_mark:
  • Integer exponents remain 0 ULP (verified: 2.0, 3.0 → 0). :white_check_mark:
  • bf16 path unchanged (already ≤ 2 ULP). :white_check_mark:
  • Performance not meaningfully regressed (~8.0 µs / 32×32 fp32 tile on WH, measured via tracy). :white_check_mark: (baseline #25025 kernel pending a strict re-profile)
  • New ULP unit tests for non-integer exponents (0.5, 1.5, 2.5, 1.7984). :white_check_mark:
  • Applied to both Wormhole B0 and Blackhole kernels. :white_check_mark:

Status of this PR against the criteria

  • All non-integer exponents within 3 ULP: 2.5 27 → 3, 1.5 19 → 2, 0.5 4 → 1; integers stay 0.
  • 1.7984 (CogVideo) 19 → 2 ULP — solved via the Veltkamp split of pow (partial two-product) that stays inside the register budget.

Why this matters

  • pow with non-integer exponents is used across production models (DeepSeek, CogVideo, Stable Diffusion). The CogVideo 10000 ** 1.7984 case from #23529 (closed by #25025) still showed a large error; this PR resolves it.
  • Same class and demand as exp2 / expm1 / atan2 accuracy bounties.
  • The fix is self-contained in the two pow SFPU LLK headers (× 2 architectures), on top of the current #25025 implementation.

(All numbers above are hardware-measured on Wormhole B0)