Basic extension catalog for kernels, plus demo on early extensibility.

[pdamme]

Before I start implementing, I'd like to share an overview of my plans with everyone interested. Feel free to comment.

Goals

• introduce initial version of basic kernel extension catalog, while not changing too many other things, for now (i.e., refactoring to enable extensibility afterwards)
• assumption: types of all intermediates are known at compile-time (the current implementation assumes this too; we will relax that later)
• functionality first, performance second, for now

Planned steps

1. introduce new class for the kernel extension catalog with very basic preliminary interface for adding new kernels and querying the registered kernels
   • representation of registered kernels
     • key: (name of a) DaphneIR op (since we always want a kernel for a specific op)
     • value: complex object
       • list of input types
       • list of output types
       • processing backend (or "API", following the terminology in kernels.json) (CPP, CUDA, FPGAOPENCL, ...)
   • interface
     • method to add a new kernel with its information
     • method to retrieve the information for a DaphneIR op
     • will be extended later
2. populate kernel extension catalog at system start-up based on the information in kernels.json
   • parse that JSON file and enter information in the kernel registry
   • no means for adding kernels exposed to users yet, just using the kernels whose pre-compilation is currently hard-coded
   • requires kernels.json in release artifact (maybe move that file)
   • anyway, using that file is just a temporary solution (see follow-up issues)
3. make DAPHNE compiler use the kernel extension catalog for lowering DaphneIR domain-specific ops to kernel calls
   • recap of current design
     • MarkCUDAOpsPass and MarkFPGAOPENCLOpsPass mark ops to be executed on the respective processing backend if certain conditions are met
     • AdaptTypesToKernelsPass harmonizes input and output types for certain ops, assuming that kernels are usually available for homogeneous type combinations
     • RewriteToCallKernelOpPass rewrites domain-specific DaphneIR ops to CallKernelOp; the name of the kernel function is created by naming convention taking into account the op name, input and output types, and hints on the processing backend; but does not know if that kernel really exists
     • LowerToLLVMPass lowers VectorizedPipelineOp and DistributedPipelineOp to CallKernelOp
   • planned design
     • MarkCUDAOpsPass and MarkFPGAOPENCLOpsPass remain as they are
     • AdaptTypesToKernelsPass remains as it is, but might not be used, for now (but could be useful again later)
     • SelectKernelsPass (new): queries the new kernel extension catalog by (op name, processing backend) and (1) decides which kernel function to call and adds a hint (attribute) to the op, and (2) inserts casts of inputs and outputs where necessary
       • raises an error if no kernel is registered for this operation
       • if a kernel for the required type combination is available, use that
       • otherwise, use any alternative kernel with the lowest cost; whereby the cost is (a) infinity if the kernel cannot be applied without loss of precision, and (b) the number of required casts to insert otherwise
       • this pass is separate from RewriteToCallKernelOpsPass, since the casts it introduces could offer potential for further optimization
     • RewriteToCallKernelOpPass will merely rewrite ops to CallKernelCallOp, based on decisions that have been made before, reflected by attributes of the op
     • LowerToLLVMPass stays as it is (those two ops shall have very generic kernels anyway)

Possible follow-up issues (not to be addressed by this PR)

• steps towards extensibility features
  • make kernel extension catalog accessible to user
  • allow plugging in additional shared libs with more kernels
  • multiple dispatch and support for kernels on super classes of DAPHNE data types
  • add performance profiles (of whatever form) to the kernel registry for cost-based kernel/device selection
• detailed improvements
  • more sophisticated selection of alternative kernels by better modeling of the costs (e.g., take data sizes into account (logical, physical), try not to deflate data more than necessary, rather deflate small than large data objects, …)
  • canonicalization for CastOp (remove cast from x to x; simplify cast from x to y to z to cast from x to z, where possible)

Closes #455.

[pdamme]

@EricMier: I think this could be interesting to you. Once we have the initial extension catalog for kernels, it could become the home of any performance profiles for kernel/device selection. The kernel extension catalog could be accessible throughout the DAPHNE compiler, such that any pass could query its information, e.g., a separate pass selecting a suitable device (as an alternative/improvement to MarkCUDAOpsPass etc.) Feel free to let me know your thoughts and if this sounds in-line with your goals.

[pdamme]

@corepointer: JFYI, since you've also worked on the lowering from DaphneIR ops to kernel calls, especially MarkCUDAOpsPass and MarkFPGAOPENCLOpsPass.

[pdamme]

This PR is ready to be merged now, from my point of view. It contains the basic extension catalog for kernels (see above, with a few minor deviations from the initial ideas) plus a few more commits that show its end-to-end use.

Of course, feedback is welcome, but I'm not asking for a code review. However, as this PR changes a few important things, I would like to give everyone some time to have a look before I merge it. Unless anyone speaks up, I will merge this PR next Monday (Apr 29). The commits in this PR are meant to be "rebased & merged", not "squashed & merged".

These are the most important changes (see the commit messages for more details):

• Introduction of an initial extension catalog for kernels (a73dfbe)
  The DAPHNE compiler finally knows which kernels exist in pre-compiled form, since the available kernels are registered in a data structure, the kernel extension catalog. This catalog is populated from a JSON file at DAPHNE start-up, the JSON file is automatically generated along with the kernel instantiations and resides in lib/ (such that it should already be covered by our scripts for deployment and release packaging). When generating calls to pre-compiled kernels, the DAPHNE compiler looks up the kernel function names in the catalog, instead of using a naming convention.
• Automatic determination of the required kernel libs (894f563)
  The DAPHNE compiler additionally knows in which shared lib each kernel is. Based on that, the DAPHNE compiler knows which shared libs we really need to link with for the DaphneDSL script at hand. The information on the shared libs is part of the kernel catalog JSON files. Thus, the way the paths to these libs were managed before could be simplified.
• Initial support for kernel hints in DaphneDSL, a test case of a simple kernel extension, and initial docs on kernel extensions (9eb00f1, 62d8dc6, 131ed82)
  These commits mainly add new material, so it's less likely they conflict with anyone else's interests. I decided to add them to this PR, since they showcase the end-to-end use and added value of the extension catalog. It is essentially a cleaned-up and extended version of the demo on extensibility I've shown at the DAPHNE General Assembly meeting in December 2023.

This PR is just the first step towards extensibility in DAPHNE and there are many more aspects and details we will work on in the future.

[philipportner]

I'm exited to play around with this in the next couple days. After only reading the new docs and looking at the examples, this already looks really awesome, great work @pdamme !

[pdamme]

Thanks. I'm sure it still has several limitations and can be improved from an efficiency and usability point of view, but it is a first step that we can build upon.

[philipportner]

I wrote a couple other extension and that, following the example in doc/Extensions.md worked quite well.
The only thing I would like to see in doc/Extensions.md is how you get the correct mnemonic of a Daphne op.
For example, when doing the same for the aggMax built-in, the correct mnemonic in the kernel_extensions.json would be maxAll. The only place where (besides at some points in the IR) I would be able to see that would be in DaphneOps.td where the op is defined. Maybe an explanation or link to figure out what the correct mnemonic is would be helpful.

Was also confused by the wrong results in our internal kernel, but I just saw you have already created an issue for that.

Otherwise, I think this looks great so far!

[corepointer]

Awesome work @pdamme 👍
I read through all five commits with adequate care (I hope) to understand what's going on but I did not compile and run it yet. So I don't have any particular question upfront without getting my hands dirty. But I'm looking forward to getting into the details and to see catalog.json 👀

[pdamme]

Thanks for the feedback, @philipportner. Happy to hear that you succeeded writing some little extensions. Indeed, the documentation is not comprehensive yet. I didn't want to spend to much time on that as we may still change a few things. However, I will add a short note (for developers) on where to get essential information to the docs (op mnemonics in DaphneOps.td and kernel interfaces in the generated kernels.cpp).

Thanks also to @corepointer. Merging it in is not super urgent, so simply let me know if you would like to have a closer look and try it out before I merge it. I would be happy to merge it by the end of this week, though (unless any significant concerns arise).

[pdamme]

I've rebased this PR and tried to adapt the changes it contains to the new error handling from #706 as good as I could (hope I didn't miss anything). I've also added a note on where DAPHNE developers can find references of the op mnemonics and kernel interfaces to the docs on writing a custom kernel extension.