Implement a cache for vk::ShaderModule objects and vk::Pipeline objects

[softcookiepp]

Currently every time a shader is to be used with a different workgroup size and set of buffers, a new kp::Algorithm must be created.
This introduces a lot of unnecessary overhead, as identical shader modules and pipelines are created every time a user wants to use a different set of buffers with the underlying SPIR-V.

Existing machine learning libraries like Pytorch allocate and deallocate buffers in an unpredictable manner, and thus an attempt to integrate Kompute would result in extremely slow performance and possible memory overuse from all the duplicate shader modules and pipelines.

There is a lot that can be done to remedy this problem, but a good first step is to implement separate caches for both vk::ShaderModule and vk::Pipeline obects. This will alleviate a large amount of overhead.

[softcookiepp]

This is what I propose it will look like. Realistically, the proposed kp::ShaderModuleCache could be eliminated and the map of ShaderModule objects it holds could be held by the kp::Manager class, but it would still have to somehow communicate with the kp::Algorithm instances to retrieve cached shader modules. I am unsure of how to do this, and would prefer that the method for hashing SPIR-V is tied to a cache object anyways; this will greatly simplify the rest of the development process.

As for the reason for the ShaderModule class itself, eventually it would increase user friendliness of Kompute in general by implementing more things like SPIR-V reflection and built-in shader source compilation (like OpenCL, OpenGL, and CUDA.)
In general, I am a fan of making iterative changes to designs when needed and maintaining flexibility for new features, rather than trying to get everything perfect on the first try; this reflects that.

[softcookiepp]

It is also worth mentioning that none of these new classes will be exposed externally. Kompute's user-visible API won't be changed by this at all; it is only for making the existing interface more efficient.

[softcookiepp]

Also, I understand the hesitancy to add more classes. However, Vulkan has a lot of complexity that is simply not taken into account with Kompute's existing classes to the extent that it can be. If the codebase is not made flexible in advance, it will be far too easy to code yourself into a corner of unmaintainability. 11:01 of this video kind of describes why this is a problem.

The kp::Algorithm class is inching close to this already. there is still a lot of work to be done in order to make it more manageable, but at least caching the vk::ShaderModule objects will reduce a ton of the overhead that existing design choices have caused.

[axsaucedo]

Thanks for the PR + diagram - the overall motivation sounds like in the right direction, and I agree there is an opportunity here, but we need to be aware that this does not come without tradeoffs, and the design / approach will have different implications on the framework and long term maintenance. Let’s first align on the design and the pros / cons before we start a discussion on whether adding this complexity makes sense or not as this context is necessary - especially as looking at the current design choices in the initial PR makes clear we need to do a few iterations on this discussion. For context and example, Kompute used to have a separate Shader class and it was removed to reduce complexity.

Taking a step back to your overall proposal, there seem to be two objectives, one is to separate the overhead of the shader module creation / management from the algorithm to reduce overhead, and the other is to abstract the overhead of the pipeline creation and management from the respective algorithm and tensors. Both of these seem like they can be approached independently from each other (and noting that the heavier compilation/baking cost tends to happen during pipeline creation rather than shader module creation).

On the ShaderModule side, looking at the current design it seems like an intermediate approach could be to keep a simpler implementation, where we can re-introduce a Shader class which can hold the ShaderModules. At this stage however I’m not yet clear on the value of the Manager class managing the cache itself of the ShaderModule, as it would be an added complexity for what can be arguably not significant overhead, so best would be to start with a limited iteration with a backward compatible API. Namely the Algorithm can receive a shader (str) as it does today, and create an owned underlying Shader class - alternatively it can receive a Shader class which would be managed by the user (as opposed to managed by the Manager). We should also clarify whether this external shader path is only used during initialisation or whether the Algorithm retains a reference (e.g. for rebuild paths).

Once we iterate on this implementation, we can look at the opportunity to cache as part of the pipeline caching, which has a much larger set of considerations and design choices that will be required, but we can iterate once we cross that bridge.

[softcookiepp]

Understood. I would be happy with simply implementing kp::ShaderModule abstraction for now (mostly done) then, and probably rename it to kp::Shader or something to avoid confusion with the vulkan.hpp class. Does that sound like a good objective for now?

[axsaucedo]

Yes let's go forward and we can iterate accordingly.

[softcookiepp]

Perfect! The latest changes to my pull request now reflect this direction.