Repository navigation
[ARM64] Vectorized math approximations, elementwise fusion and buffer deallocation for activations - #6
Open
Aditya Trivedi (adit4443ya) wants to merge 3 commits into
Open
[ARM64] Vectorized math approximations, elementwise fusion and buffer deallocation for activations#6Aditya Trivedi (adit4443ya) wants to merge 3 commits into
Aditya Trivedi (adit4443ya) wants to merge 3 commits into
Conversation
- 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).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
GELU (erf and tanh forms) did not compile, exp-based activations ran a scalar
expfper lane, and everykernel 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
triton-arm64-math-approximation, run intriton-linalg-to-llvmbeforeconvert-math-to-llvm.It expands
math.erf/exp/tanhinto MLIR's polynomial approximations (f16/bf16 computed in f32); the setis configurable with
ops=.math.erfhad no LLVM lowering:gelu_erffailed in mlir-translate with "Dialect `math' not found".math.expon scalable vectors becamellvm.exp.nxv4f32, which the AArch64 backend scalarizes intoone
bl expfper lane.math.fpowi. torch-mlir writes the x³ in GELU-tanh asfpowi, and itsvector form cannot lower (
llvm.intr.powitakes a scalar exponent). A constant power becomesmultiplies. This fails on main too.
optimize-triton-shared. triton-shared emits onelinalg.genericper op; on gelu_tanh's 15 ops,TritonARM64BubbleUpExtractSlicehit the greedy rewritelimit and the pipeline failed without a diagnostic. After fusion the kernel is one loop with no
temporaries.
2. Buffer deallocation
mallocper temporary.I run
buildBufferDeallocationPipelineright afterTritonARM64VectorLowering, while loops are stillscf(the SME passes lower them tocfnext, and the dealloc pipeline needs structured loops).@dealloc_helperit generates for buffers yielded from loops, because the driver expectsone function per kernel module (
get_kernel_name).3. Tests
python/examples/test_activations.py: relu, leaky_relu, sigmoid, silu, tanh, gelu_erf, gelu_tanh vsPyTorch, including partial blocks and saturation (±inf, |x| > 88).
python/torch-mlir/test_activations_torch.py: the same activations fromnnmodules.execute_and_comparenow returns the match result, so tests can fail, and takes an optionalreference. torch-mlir imports GELU-tanh with √(2/π) = 0.797724 instead of 0.7978846 (up to 5e-5 offPyTorch), so that case is checked against the imported formula.
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)
-optimize-triton-sharedfails (15linalg.generic)math.erfbl malloc, 1bl expfin a per-lane loopfmlamemref.alloc, 0deallocalloc, 15deallocbl malloc, 0bl freebl malloc, 12bl freePerformance
SVE, ns/element (median of 3 runs x 100 iterations), main -> this PR:
relu and leaky_relu call no
exp; their speedup comes from fusion and from no longer allocating atemporary per block.