GPU Compressors

[ilan-gold]

See https://github.com/NVIDIA/CUDALibrarySamples/tree/master/nvCOMP and https://github.com/rapidsai/kvikio/tree/branch-25.02/python/kvikio and https://github.com/NVIDIA/nvcomp?tab=readme-ov-file

Python bindings already exist so I think we would need rust bindings as well (although maybe we could call the python through rust with pyo3? Maybe a bad performance?). In any case, it would be great to somehow add support for this (somehow emphasized, almost certainly a lot of work) given the speed boost of not having to do the CPU bounce, and zarr-python v3 natively supprting GPU buffers if I'm not mistaken. So this should be doable, but probably a lot of work and maybe it's own package

[LDeakin]

Python bindings already exist so I think we would need rust bindings as well

Unfortunately, no kvikio Rust bindings exist. I tried a few months ago and bailed. Happy to share that repo if you want, but I didn't get too far. Relevant: apache/opendal#5090

So this should be doable, but probably a lot of work and maybe it's own package

Probably!

zarrs does support bypassing the codec chain and storing encoded chunks directly. Maybe that could be useful for experimentation? I support GPU decompression in one of my apps that uses zarrs (not yet open source). Maybe this could serve as some inspiration, although it is specialised to my use case:

1. Encode array on the host with the zarrs gdeflate codec, this is NVIDIA GDeflate
   • This codec is not in the Zarr spec
   • The chunk shape matches the GPU sparse image block shape
2. Retrieve encoded chunks with Array::retrieve_encoded_chunk[s]
   • I've been meaning to add array::ArrayShardedReadableExt::retrieve_encoded_inner_chunk()
3. Decode in Vulkan with the VK_NV_memory_decompression extension into sparse image blocks

[ilan-gold]

I tried a few months ago and bailed

Nice, but probably doesn't bode well for my being able to do it then haha. But sure, let's get more eyes on it, maybe Phil wants to take a look at some point

[ilan-gold]

This codec is not in the Zarr spec

I had no idea codecs were specificed in the spec. I see they are there, but I was under the impression compression was basically a free-for-all, you could use whatever. But yeah, that pipeline looks pretty clean, basically what I thought zarr was supposed to be doing anyway with any codec (aside from the vulkan part). I'll ask in the zulip about this

[ilan-gold]

Commented in the channel, but I think we could in theory start with a 1-1 codec implementation i.e., one in the spec that exists also as a NVIDIA implementation