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.
hunhoffe commentedon Oct 28, 2026
Closing this out — the generic epilogue abstraction landed, just via a different path than the ChanneledUnaryOperator reuse suggested here.
flm.GEMM now carries an Epilogue enum (iron/operators/flm/gemm/op.py) with a configurable epilogue_modes tuple compiled into the build and an epilogue_mask for runtime mode selection, so bias/activation handling is no longer hardcoded per-operator. gemv.py no longer carries any epilogue/gelu code of its own.
If there's still a case this doesn't cover, happy to reopen with the specifics.