gemv: add optional fused GELU epilogue (NPU2)#135
Open
atassis wants to merge 1 commit into
Open
Conversation
Adds an opt-in `epilogue` parameter to GEMV. `epilogue="none"` (default) leaves the output unchanged and is byte-identical to before; `epilogue="gelu"` applies a GELU (tanh approximation) to each output tile inside the producing core, fusing an activation that would otherwise be a separate elementwise pass over the GEMV output. The GELU is applied once per full output tile in core_body (after the matvec inner-loop fills all rows), not per matvec call, whose m_input tile can be smaller than the 16-wide activation vector. To avoid duplicating the GELU math, aie2p/gelu.cc is refactored: the per-16-lane computation is factored into a shared helper used by both the existing out-of-place gelu_tanh_approx_bf16 and a new in-place gelu_tile_bf16. In-place is aliasing-correct (a single restrict pointer, each slot read then written) whereas reusing the out-of-place kernel with input==output would violate its restrict contract. Because gelu.cc is aie2p-only, the fused epilogue is NPU2-only; GEMV raises NotImplementedError if requested on NPU1. When enabled the core links a (matvec, gelu) archive. Test: test_gemv_gelu compares the fused output against a gelu(A @ B) golden across three shapes, with the standard latency/bandwidth/throughput metrics. Verified on NPU2 (15/15 pass).
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
1 participant
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.
Adds an opt-in
epilogueparameter toGEMV.epilogue="none"(default) is byte-identical to today;epilogue="gelu"fuses a GELU (tanh approximation) over each output tile inside the producing core, folding an activation that would otherwise be a separate elementwise pass over the GEMV output.Details
m_outputtile incore_body(after the matvec inner-loop has filled all rows), not per matvec call whosem_inputtile can be smaller than the 16-wide activation vector.aie_kernels/aie2p/gelu.ccis refactored so the per-16-lane computation is a shared helper used by both the existing out-of-placegelu_tanh_approx_bf16and a new in-placegelu_tile_bf16. In-place is aliasing-correct (a singlerestrictpointer, each slot read then written); reusing the out-of-place kernel withinput == outputwould violate itsrestrictcontract, so a dedicated in-place entry point is the correct way to share the code.gelu.cclives inaie2p/, the fused epilogue is NPU2-only;GEMVraisesNotImplementedErrorif it is requested on NPU1. When enabled, the core links a(matvec, gelu)archive; the default path still links the single matvec object.Test
test_gemv_gelucompares the fused output against agelu(A @ B)golden across three shapes, using the same latency/bandwidth/throughput@pytest.mark.metricsastest_gemv. Verified on NPU2: 15/15 pass.This is opt-in and default-preserving; it also composes with the existing
func_prefixfusion support (#123) for use as a fused block.