I have a computational graph (DAG) consisting of element-wise operations (potentially with broadcasting) and reshape/reduce operations (reshaping/sum/max). I'm trying to understand how vertical kernel fusion is performed and whether it's something that can be done without developing/using external tools like a linear algebra compiler or some custom IR format (like HLO IR for JAX).

The first thing that came to my mind was that I could simply have symbols corresponding to element-wise operations (like * for multiplication and sin( for sine) and once the computational graph is assembled, literally concatenate these symbols corresponding to element-wise operations into one big operation (custom schema to keep track of data pointer names, but still I imagine that's possible). And since I'm only using element-wise operations, the data access is very predictable and linear, so it's perhaps possible to automatically generate performant-ish kernels for both CPU and potentially other processing units.

Is this possible? Am I on the right track here? Or is another approach typically used? If so, where can I read more about it.

thanks :)



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.