[REVIEW] Develop new mdspan-ified multi_variable_gaussian interface

[mhoemmen]

Develop a new multi_variable_gaussian interface that uses mdspan. The new interface uses free functions, rather than a class.

TODO:

• Do not expose workspaces to the public API (see e.g., Implement raft::stats API with mdspan #802 (comment)); use RMM instead.
• Add a public API not in in the detail namespace, and test it.
• File issues for refactoring and improving the implementation.

[cjnolet]

Do not expose workspaces to the public API (see e.g., Implement stats API with mdspan #802 (comment)); use RMM instead.

If you have time to support it, adding an optional rmm::mr::device_memory_resource (or even better, add our own interface that can wrap rmm's device and host memory resources) to the handle could be super helpful to get the workspace out of the public API. We can also use it on the ivf flat search API.

[mhoemmen]

@cjnolet wrote:

If you have time to support it, adding an optional rmm::mr::device_memory_resource (or even better, add our own interface that can wrap rmm's device and host memory resources) to the handle could be super helpful to get the workspace out of the public API. We can also use it on the ivf flat search API.

I'll be happy to add an optional rmm::mr::device_memory_resource parameter. For adding it to the handle instead, do you have a sense for what that interface should look like? I was thinking of a free function like this.

device_memory_resource get_default_memory_resource(const handle_t&);

The const reference is necessary because of all the interfaces that take const handle&, but storing the memory resource in the handle would call for the same mutable pattern as before.

[cjnolet]

The const reference is necessary because of all the interfaces that take const handle&, but storing the memory resource in the handle would call for the same mutable pattern as before.

I think that could work. To make this compatible w/ some of the ideas we've been tossing around for the more library agnostic raft handle, what do you think about sticking this in a file called rmm_resource_handler.hpp? Just thinking out loud here, but it would be nice if the handle ultimately just became a key/value store and primitives could request what they need from the handle. Eventually, we could have includes like cublas_resource_handler.hpp, which would lazy-load a cublas handle and store it on the raft handle when it's needed. Of course, this disregards the immutability of the const qualifier and might require some care (locks?) to avoid conflicts in concurrent updates. The handle is coupled to a single "user" stream, but what if these resources are updated in a worker stream?

The basic idea here would be that if no primitives are invoked that need cublas, there's no reason for the user to ever have to interact with it, and thus no reason to even bother creating a cublas handle at all.

[mhoemmen]

@cjnolet wrote:

Just thinking out loud here, but it would be nice if the handle ultimately just became a key/value store and primitives could request what they need from the handle.

I've filed this as a feature request: #856

Thanks! : - )

[cjnolet]

Thanks for these wrappers, Mark! This looks great overall. Just minor things and I think it's ready to go.

[cjnolet]

LGTM. Thanks mark! The new API looks really clean.

[cjnolet]

@gpucibot merge

[cjnolet]

@gpucibot merge