CUBLAS wrappers with switchable host/device pointer mode

[achirkin]

Add a few overloads for raft-CUBLAS gemv, gemm, axpy functions to support switching between host and device pointer mode. This allows passing some of the parameters (constants alpha, beta) as device pointers, which sometimes improves performance.

By default, CUBLAS context is created in the host pointer mode. To keep this presumption, the device pointer mode is enabled only for the time of a particular CUBLAS call.

This feature is required for rapidsai/cuml#4446.

[achirkin]

rerun tests

[tfeher]

Thanks Artem for this PR! The changes look good to me.

[tfeher]

One thing still worth to discuss with @cjnolet: how many overloads for gemm and gemv do we want to keep in the API? The advantage of the wrappers introduced by Artem is that they follow the standard BLAS argument order and they also expose the ld parameter.

If we are still consolidating the RAFT API then one could consider a follow up PR to remove some of the overloads (together with cleanup in downstream repositories):

• For gemv, I am not convinced that we need any of the additional overloads.
• For gemm, one could argue that the additional versions are useful, because they set several default parameters. I am still not sure if we need all of the overloads.

[cjnolet]

@tfeher,

how many overloads for gemm and gemv do we want to keep in the API?

Indeed, that was the first thing I noticed when I began reviewing this PR. I don't think we need them, and even more specifically, once the mdspan PR is merged, we should be able to acept device_scalar and host_scalar in place of the new pointers (or have them implicitly convertable).

That said, if these changes are urgent, I'm okay creating a Github issue to follow up when the mdsan is ready and getting these merged.

[achirkin]

That said, if these changes are urgent, I'm okay creating a Github issue to follow up when the mdsan is ready and getting these merged.

Thanks, I'd prefer to get this merged now to proceed with rapidsai/cuml#4446, which is almost two month old now.

[tfeher]

@cjnolet I had a discussion with Artem about this:

• We need an additional overload that accepts device pointers (or device_scalar in the future).
• Indeed, the preferred way would be to have an overload that accepts a device scalar or something equivalent.
• But to make it actually useful in Rewrite CD solver using more BLAS cuml#4446, we would need a way to get a device scalar from a device_vector, which we do not have yet.
• Since RMM device_scalar owns the scalar, constructing that would mean memcpy, which we want to avoid.
• Alternatively one could think about accepting a device_span<math_t, 1> once [REVIEW] Span implementation. #399 is merged.

Considering these points, we would prefer to go forward with the current PR and rapidsai/cuml#4446. Once span / mdspan is merged, we shall consolidate the API of gemm / gemmv to use mdspan, and refactor the relevant calls in the CD solver.

[cjnolet]

Since RMM device_scalar owns the scalar, constructing that would mean memcpy, which we want to avoid.

I'd prefer not to use RMM in the public API. The new mdspan will provide a device_scalar_view which can be created from an rmm device_uvector without copying. I would prefer to use that when ready.

Otherwise, I don't see any reason to hold up this PR.

[cjnolet]

I've gone ahead and added our plan to an existing issue, which I created to clean up the ctk wrappers in cuml in general: #475

[cjnolet]

@achirkin, just missing doxygen but otherwise looks great.

[cjnolet]

LGTM. Thanks again @achirkin!

[achirkin]

rerun tests

[cjnolet]

rerun tests

[cjnolet]

@gpucibot merge