[Vulkan] gelu shader returns NaN for inputs below about -13 (unclamped tanh argument overflows fp32)

[msluszniak]

🐛 Describe the bug

The gelu shader in backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml evaluates the tanh approximation with an unclamped argument:

- NAME: gelu
  OPERATOR: 0.5 * X * (1 + tanh(sqrt(2 / 3.141593) * (X + 0.044715 * X * X * X)))

The argument grows cubically in X, so an activation of x = -13.24 produces

sqrt(2/pi) * (-13.24 + 0.044715 * (-13.24)^3) = -93.4

A driver that evaluates tanh(y) as (e^y - e^-y) / (e^y + e^-y) overflows fp32 once |y| > ~88, giving inf/inf = NaN. The threshold in X works out to about -13.0, which matches what I observe exactly.

Note that the tanh op a few lines below in the same file already clamps for this reason:

- NAME: tanh
  OPERATOR: tanh(clamp(X, -15.0, 15.0))

so only gelu is exposed.

Observed behaviour

Running the Whisper-tiny encoder (openai/whisper-tiny, encoder only, fp32, Vulkan delegate) on a Samsung Galaxy S10+ (Mali-G76, Android 12, driver v1.r32p1) returns an all-NaN output tensor. The runner reports success and normal per-iteration timings, so the failure is silent.

I bisected the encoder by exporting progressive prefixes of its graph. Results on device, against a CPU reference:

stage	result
gelu(conv1(x))	correct, cos 1.00000024
gelu(conv2(...))	NaN in 2 of 576000 elements
after transformer block 0	NaN in 576000 of 576000
full encoder	NaN in 576000 of 576000

The two bad elements are at [0, 349, 521] and [0, 349, 836]. Those are precisely the two most negative pre-activations in the tensor: -13.2371 and -13.1652, and -13.2371 is the global minimum. Every other element is finite. The first LayerNorm plus attention then spreads the two NaNs across the entire tensor, which is why the model output is uniformly NaN.

Reproduction

import torch, torch.nn.functional as F
x = torch.tensor([-13.24])
# CPU: -0.0 ; Vulkan on Mali-G76: NaN
F.gelu(x, approximate="tanh")

Any graph containing a gelu whose input drops below about -13 will reproduce it.

Versions

• ExecuTorch main (also present on current origin/main)
• Device: Samsung Galaxy S10+ (SM-G975F), Mali-G76, Android 12, Vulkan driver v1.r32p1

Note that ExecuTorch 1.4.1 is unaffected, because there gelu still carries a clamp that was subsequently dropped from main:

# v1.4.1
OPERATOR: 0.5 * X * (1 + tanh(clamp(sqrt(2 / 3.141593) * (X + 0.044715 * X * X * X), -10.0, 10.0)))

I confirmed this directly: the same .pte and the same input on the same device gives a correct result under a 1.4.1 runtime (cos 0.99999714, 13 of 13 runs) and all-NaN under main (10 of 10 runs).

Fix

Clamping the argument to the same +/-15 the tanh op uses restores correct output on Mali: cos 0.99999702, no NaN, 10 of 10 runs bit-identical. The clamp costs no accuracy, since 1 - tanh(15) = 1.9e-13 is well below fp32 epsilon and the result is bit-identical for every input that gets clamped.

PR follows.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin