[proposal] Support of functional interfaces

[andrii0lomakin]

Many tasks of heterogenous computing are variations of the same well-known patterns that differ only in functions called per element.
Like, reduce, for example, when SoftMax and RMS normalization, in a nutshell, differ only in functions processed.

It would be beneficial to support functional interfaces as parameters of kernels that unwind to real calls (with appropriate restrictions, of course).

That will minimize development time and increase the maintainability of kernels developed in TornadoVM.

[jjfumero]

Hi @andrii0lomakin , is this a proposal or an issue?

You can build libraries on top of TornadoVM that solve specific functionality for domains of applications such as LLMs, Linear Algebra, Graphics, Physics, etc.

[andrii0lomakin]

This is a feature request. At least, I have chosen it like this :-) . At the moment, I need to repeat the same boilerplate code repeatedly, and it would benefit me to pass functional interfaces.

[jjfumero]

This is a feature request.

Ok, I open this for discussion for all community members and TornadoVM maintainers.

At the moment, I need to repeat the same boilerplate code repeatedly, and it would benefit me to pass functional interfaces.

What do you mean by passing functional interfaces? At which level? Tasks within TaskGraph receive functional interfaces already. Do you have any example in mind that you can share?

As I mentioned, one can build libraries on top. For example, LLM and transformer library that contain softmax, normalization, reductions and the matmul. Is this what you are referring to?

[andrii0lomakin]

Simplest example.

I have an RMS norm layer that, in a nutshell, consists of reduce kernels that are called in layers as described here: https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf.

Initially, I reduce the squired values and then just call the plain sum kernel in subsequent layers.
I can create a single kernel that accepts Function as a parameter and inline it during compilation, but instead, I need to repeat the same code.

The same is true for SoftMax, which uses a reduce function as its denominator.

I suppose there are plenty of other such use cases.

P.S. I am ready to implement this feature myself, but I need your feedback, of course.

[jjfumero]

Reductions are supported in TornadoVM in two ways:

1. Using the combination of @Parallel and @Reduce for loop-parallelism. The TornadoVM JIT Compiler is able to generate efficient reductions for simple test-cases (min, max, sum, etc).
   Details: https://www.researchgate.net/publication/327871451_Using_Compiler_Snippets_to_Exploit_Parallelism_on_Heterogeneous_Hardware_A_Java_Reduction_Case_Study

Note that the final reduction happens on the device. The TornadoVM JIT compiler generates two kernels from a single reduction kernel (one to run in parallel, and the second to perform the final reductions from the remaining work-groups).

2. Using the Kernel API and the KernelContext to access local memory and barriers. In this case, the developer needs to implement the whole logic. This might be also the path to follow for more complex reductions.

[andrii0lomakin]

@jjfumero seems like I was not clear enough, I know how to implement reduce and have knowledge of barriers and etc.

I mean that if I had the ability to pass the function that I need to call when I perform reduce, I would not need to repeat the same boilerplate code again and again.

For example in my code I have two kernels, one does:

    localSnippet[context.localIdx] = inputTensor.get(currentInputOffset);

and another one

 float value = inputTensor.get(currentInputOffset);
 localSnippet[context.localIdx] = value * value;

in general I can pass just a function in both kernels, but instead I need to repeat the same code again and again.

P.S. In general when I raise the issue I have impression that I always receive information how to implement basics of functionality instead of discussion of concrete issue in depth, probably I need change something in my conversational style to avoid this :-)

[andrii0lomakin]

As for TorandoVM annotations, I find them a good entry point for developers who want to learn about heterogenous computing, but because all those "complications" are done for nothing but performance, at least at this concrete stage they are not quite suitable for production usage, at least in commercial applications or in librariries that are created to support commercial tools.