Repository navigation
Generic abstraction for prologue/epilogue elementwise unary operator fusion onto GEMV/GEMM #193
Description
We're currently hard-coding which unary operators ("activation functions") may be applied/fused onto a GEMV on to a core:
IRON/iron/operators/gemv/op.py
Lines 58 to 61 in deb6e1e
A nicer "Kernel" abstraction that encodes (i) the function name (extern C), (ii) the argument specification expected (MLIR func.func signature), (iii) how to compile and link it in would be good at some point in the future. If we have to list every single activation function in every GEMV and GEMM kernel as a custom prologue/epilogue, it will get old quick.
We could probably reuse ChanneledUnaryOperator for this. Extract the kernel and its calling convention out of the ChanneledUnaryOperator. A single prologue=ChanneledUnaryOperator(...) would enode all the information needed -- no listing of all functions inside GEMV.
cc @atassis
Activity
atassis commentedon Sep 8, 2026
Thanks for pointing on this, since while I am trying to make things work- I skip deeper design review for now. This should be properly thought, after I finish with my current tasks- I'll definitely think on that.