Skip to content

[ARM64] Vectorized math approximations, elementwise fusion and buffer deallocation for activations - #6

Open
Aditya Trivedi (adit4443ya) wants to merge 3 commits into
qualcomm:mainfrom
adit4443ya:ac3w
Open

Aditya Trivedi (adit4443ya) wants to merge 3 commits into
qualcomm:mainfrom
adit4443ya:ac3w

Conversation

@adit4443ya

@adit4443ya Aditya Trivedi (adit4443ya) commented Oct 20, 2026 •

Copy link
Copy Markdown

GELU (erf and tanh forms) did not compile, exp-based activations ran a scalar expf per lane, and every
kernel leaked its temporary buffers. This fixes all three and adds Triton and torch-mlir tests for the
common LLM activations. On device (SVE) activations run 4.9-20.2x faster, and both GELUs now compile.

Changes

1. Math approximations and elementwise fusion

  • New pass triton-arm64-math-approximation, run in triton-linalg-to-llvm before convert-math-to-llvm.
    It expands math.erf/exp/tanh into MLIR's polynomial approximations (f16/bf16 computed in f32); the set
    is configurable with ops=.
    • math.erf had no LLVM lowering: gelu_erf failed in mlir-translate with "Dialect `math' not found".
    • math.exp on scalable vectors became llvm.exp.nxv4f32, which the AArch64 backend scalarizes into
      one bl expf per lane.
  • The pass also always expands math.fpowi. torch-mlir writes the x³ in GELU-tanh as fpowi, and its
    vector form cannot lower (llvm.intr.powi takes a scalar exponent). A constant power becomes
    multiplies. This fails on main too.
  • Linalg elementwise fusion now runs first in optimize-triton-shared. triton-shared emits one
    linalg.generic per op; on gelu_tanh's 15 ops, TritonARM64BubbleUpExtractSlice hit the greedy rewrite
    limit and the pipeline failed without a diagnostic. After fusion the kernel is one loop with no
    temporaries.

2. Buffer deallocation

  • Bufferization never freed its temporaries, so every program instance leaked one malloc per temporary.
    I run buildBufferDeallocationPipeline right after TritonARM64VectorLowering, while loops are still
    scf (the SME passes lower them to cf next, and the dealloc pipeline needs structured loops).
  • I inline the @dealloc_helper it generates for buffers yielded from loops, because the driver expects
    one function per kernel module (get_kernel_name).

3. Tests

  • python/examples/test_activations.py: relu, leaky_relu, sigmoid, silu, tanh, gelu_erf, gelu_tanh vs
    PyTorch, including partial blocks and saturation (±inf, |x| > 88).
  • python/torch-mlir/test_activations_torch.py: the same activations from nn modules.
    execute_and_compare now returns the match result, so tests can fail, and takes an optional
    reference. torch-mlir imports GELU-tanh with √(2/π) = 0.797724 instead of 0.7978846 (up to 5e-5 off
    PyTorch), so that case is checked against the imported formula.
  • lit: arm64-math-approximation.mlir, triton-linalg-to-llvm-math-approximation.mlir,
    optimize-triton-shared-elementwise-fusion.mlir, optimize-triton-shared-buffer-deallocation.mlir.

What changes in the generated code (sve, qemu)

Kernel main this PR
gelu_tanh -optimize-triton-shared fails (15 linalg.generic) 1 vectorized loop, 0 allocs
gelu_erf mlir-translate fails on math.erf compiles, erf as 8 vector FMAs
sigmoid asm 3 bl malloc, 1 bl expf in a per-lane loop no calls, 5 SVE fmla
layernorm 1151x8192, optimized IR 24 memref.alloc, 0 dealloc 13 alloc, 15 dealloc
layernorm asm 24 bl malloc, 0 bl free 10 bl malloc, 12 bl free
layernorm peak memory 756 MB 44 MB

Performance

SVE, ns/element (median of 3 runs x 100 iterations), main -> this PR:

Activation 64K 1M 16M
relu 1.35 -> 0.22 (6.1x) 1.10 -> 0.22 (4.9x) 1.06 -> 0.20 (5.2x)
leaky_relu 4.80 -> 0.24 (20.2x) 4.14 -> 0.24 (17.3x) 4.13 -> 0.20 (20.2x)
sigmoid 8.62 -> 1.56 (5.5x) 6.33 -> 0.81 (7.8x) 6.32 -> 0.59 (10.7x)
silu 9.83 -> 1.63 (6.0x) 7.56 -> 0.63 (12.1x) OOM -> 0.63
tanh 11.13 -> 1.77 (6.3x) 9.56 -> 0.87 (11.1x) OOM -> 0.67
gelu_erf doesn't compile -> 1.82 doesn't compile -> 0.79 doesn't compile -> 0.79
gelu_tanh doesn't compile -> 0.94 doesn't compile -> 1.15 doesn't compile -> 0.95

relu and leaky_relu call no exp; their speedup comes from fusion and from no longer allocating a
temporary per block.

- New pass triton-arm64-math-approximation: expands math.erf/exp/tanh into
  MLIR polynomial approximations (math.erf had no LLVM lowering; math.exp
  on scalable vectors became a scalar expf loop) and math.fpowi, whose
  vector form llvm.intr.powi cannot lower.
- Run linalg elementwise fusion at the start of optimize-triton-shared, so
  long chains like GELU stay within the greedy rewrite limit.
- lit tests for both.
Bufferization never freed its temporaries, so every program instance
leaked one malloc per temporary. Run the buffer deallocation pipeline
after vector lowering and inline the generated @dealloc_helper.
- examples/test_activations.py: relu, leaky_relu, sigmoid, silu, tanh,
  gelu (erf, tanh) vs PyTorch, incl. partial blocks and saturation.
- torch-mlir/test_activations_torch.py: the same activations from PyTorch
  modules. execute_and_compare now returns the match result and takes an
  optional reference (torch-mlir imports GELU tanh with sqrt(2/pi) =
  0.797724, so that case is checked against the imported formula).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant