[FEA] Improvements on bitset class

[lowener]

Related to #1866.
This PR adds useful operations on bitsets: count(), any(), ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[lowener]

/ok to test

[divyegala]

Per suggestions from my last review, pointer access methods should be called data() when they are returning a pointer. data_handle() is specific to mdspan because it may or may not return a pointer

[divyegala]

Not sure what this type is, is this from rmm? Why not use raft::device_vector or rmm::device_uvector?

[cjnolet]

This is buried in the container policy code for RAFT and I'd prefer using RAFT types, even if they are just direct wrappers around thrust or rmm types (it provides us a safe facade to maintain api compatibility even if the underlying impls should somehow change or need to be modified).

[divyegala]

raft::device_vector is the way to go here!

[lowener]

Not really, I chose raft::device_uvector because it is resizable. raft::device_vector isn't, and this resizable feature is very helpful for incremental addition to IVF indexes (and soon CAGRA)

[wphicks]

For my own education, won't raft::device_vector initialize the underlying memory? Would it be worth exposing raft::device_uvector more publicly specifically for cases like this?

[wphicks]

@divyegala You're quite right! Had rmm::device_vector in my mind instead of raft::device_vector. raft::device_vector uses rmm::device_uvector for its container policy, so no initialization.

[wphicks]

Adding my two cents on the resize question, I'd be quite squeamish about using the resources from the constructor. Not only do we not know if that resources object has gone out of scope, but it also makes it much harder to track and ensure stream safety if we are not explicitly passing the stream we want to use when a resize might occur. For example:

auto res1 = device_resources{};  // Some underlying stream
auto arr = device_mdarray{res1, ...};
res1.sync_stream();
auto res2 = device_resources();
// modify underlying data of arr using the stream from res2
res2.sync_stream();
// Is it now safe to copy back data to the host on res2? Maybe. If a resize was triggered in between, we'd need another res1.sync_stream() first

Alternatively, if we sync the stream from res1 in the resize method, we might still have an issue if the data are actively being modified on another stream during the resize.

[divyegala]

@wphicks thanks for typing that out, you are right and that just clarifies me for again why we didn't end up supporting this in the first place anyway.

raft::device_uvector is supposed to be an internal/detail type and I do not belive it should be used directly. If resize() is needed then we should switch the type to rmm::device_uvector. Thoughts on this @cjnolet @wphicks ?

[cjnolet]

Actually, the original suggestion to use a raft type instead of directly using the rmm type in public API code was from me. To make the point more clear, this is about consistency and safeguarding our implementations from changes that we cannot control. This is in a very similar vein to our diligent wrapping of the CUDA math libraries APIs so that we can centralize those calls and change any underlying details should the rug get pulled out from under us.

Upon closer inspection to Mickael's changes, however, I very much agree that we should not be storing the raft::resources instance as object members and Will's example above is one of the very reasons we want to avoid this. We wouldn't otherwise be doing it in the device_container_policy if it weren't for the fact that we needed the deferred allocation. To Divye'a point, the device_uvector had been kept an implementation detail until recently. I would prefer that we fix that problem at some point soon, rather than to continue using this pattern.

[lowener]

I removed raft::resources instance from the bitset class.

[divyegala]

Approved, pending decision from reviewers on whether to use raft::device_uvector or rmm::device_uvector

[lowener]

/ok to test

[cjnolet]

@wphicks are you happy with the current implementation here? I very much agree that we shouldn't be explicitly storing raft::device_resources off anywhere.

I think we should revisit the current implementation of device_container_policy at some point and try to find a better way to implement the deferred allocation strategy without storing off the resources.

[wphicks]

@cjnolet Happy and agree 100% on device_container_policy

[cjnolet]

/merge