Operator fusion is a critical optimization in machine learning compilers that eliminates intermediate tensor materialization, reducing memory bandwidth requirements. However, aggressive fusion increases register pressure and can cause register spills that negate performance gains.
Is there any better way? Or anyone working into it? I would be interested to have a look and contribute.
Let me know if I’m missing anything..
Modern machine learning workloads consist of sequences of tensor operations amenable to fusion—combining multiple operations into a single kernel to eliminate intermediate tensor storage. For example, a transformer attention layer computes:
Attention(Q,K,V) = softmax(QK^T / √d_k) · V
This involves multiple elementwise operations (scaling, softmax components, multiplication) that can be fused to reduce memory traffic. Without fusion, each intermediate result is written to memory and read back; with fusion, intermediate values reside in registers.
However, fusion is not universally beneficial. Aggressive fusion creates large kernel bodies requiring many simultaneous live values, exhausting available registers and causing register spills, memory traffic that negates the original savings.
I could come up with CostFuse, a cost-model driven fusion framework for MLIR’s Linalg dialect that balances fusion benefits against register pressure constraints through two configurable parameters: maximum chain length and maximum body operations.
Your views on this will be highly appreciated.
See: [RFC] Target Description and Cost Model in MLIR
Such infrastructure depends a lot on what kind of target you have in mind. Most MLIR based compilers have their own in various stages of completeness, but unlikely to contribute upstream due to various factors.
Generic upstream cost models are very hard to build and maintain and can take many years to become usable and profitable. LLVM has a few and they were all multi-year efforts.
We (in the tensor compiler group) are slowly building from the ground up, you’re more than welcome to join the effort. Do you have a particular target and/or area you want to contribute first?