Optimize Arm f32 GEMM using FMLA by element #679
Merged
Conversation
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
|
There is a The broadcast lane + FMA fusion is convenient as broadcasting a lane is a more generally useful operation. |
robertknight
changed the title
Suppress this new lint until I have time to fix existing failures.
This method is no longer `unsafe`.
On Arm this maps to the `vdupq_laneq_*` instructions. LLVM is conveniently able to fuse `vdupq_laneq` + `vfmaq` into an indexed FMLA operation. For other architectures we currently fall back to store + load.
robertknight
force-pushed
the
arm-gemm-fmla-by-element
branch
from
0d6495b to
e5f241a
Compare
Arm supports an FMLA instruction variant [^1] which multiples one lane of a vector by values in another vector and then accumulates into a destination. This enables reducing the number of loads from A in each iteration of the GEMM microkernel by a factor of 4. Instead of broadcasting one scalar value from A at a time into a register used as an FMA operand, we can load a vector of 4 elements from A and then each FMA specifies which of the 4 elements to use. The way this is implemented is by using `broadcast_lane` + `mul_add` generic operations (`vdupq_laneq_f32` followed by `fmla`) which LLVM will fuse into a single `fmla`. [^1]: https://developer.arm.com/documentation/ddi0596/2021-03/SIMD-FP-Instructions/FMLA--by-element---Floating-point-fused-Multiply-Add-to-accumulator--by-element--?lang=en
robertknight
force-pushed
the
arm-gemm-fmla-by-element
branch
from
31764c1 to
ffeb6fd
Compare
|
Test with a ModernBERT base model on an M3 Pro: Before: ~420ms mean |
robertknight
added a commit
that referenced
this pull request
Apr 27, 2025
Similar to #679, use the by-element variant of UDOT [^1] to reduce the number of loads in the int8 GEMM microkernel. Instead of performing 4 scalar 32-bit loads and broadcasting the result, perform one 128-bit load and use indexed UDOT to implicitly broadcast a 32-bit lane for each of the 4 rows. [^1]: https://developer.arm.com/documentation/100069/0609/A64-SIMD-Vector-Instructions/UDOT--vector--by-element-
robertknight
changed the title
robertknight
added a commit
that referenced
this pull request
Apr 28, 2025
Similar to #679, use the by-element variant of UDOT [^1] to reduce the number of loads in the int8 GEMM microkernel. Instead of performing 4 scalar 32-bit loads and broadcasting the result, perform one 128-bit load and use indexed UDOT to implicitly broadcast a 32-bit lane for each of the 4 rows. [^1]: https://developer.arm.com/documentation/100069/0609/A64-SIMD-Vector-Instructions/UDOT--vector--by-element-
robertknight
added a commit
that referenced
this pull request
Apr 28, 2025
Similar to #679, use the by-element variant of UDOT [^1] to reduce the number of loads in the int8 GEMM microkernel. Instead of performing 4 scalar 32-bit loads and broadcasting the result, perform one 128-bit load and use indexed UDOT to implicitly broadcast a 32-bit lane for each of the 4 rows. [^1]: https://developer.arm.com/documentation/100069/0609/A64-SIMD-Vector-Instructions/UDOT--vector--by-element-
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.
Arm supports an FMLA instruction variant 1 which multiples one lane of a
vector by values in another vector and then accumulates into a destination.
This enables reducing the number of loads from A in each iteration of the GEMM
microkernel by a factor of 4. Instead of broadcasting one scalar value from A at
a time into a register used as an FMA operand, we can load a vector of 4
elements from A and then each FMA specifies which of the 4 elements to use.
The way this is implemented is by using
broadcast_lane+mul_addgenericoperations (
vdupq_laneq_f32followed byfmla) which LLVM will fuse into asingle
fmla.Future work:
Footnotes
https://developer.arm.com/documentation/ddi0596/2021-03/SIMD-FP-Instructions/FMLA--by-element---Floating-point-fused-Multiply-Add-to-accumulator--by-element--?lang=en ↩