rapidsai / raft

RAFT contains fundamental widely-used algorithms and primitives for machine learning and information retrieval. The algorithms are CUDA-accelerated and form building blocks for more easily writing high performance applications.
https://docs.rapids.ai/api/raft/stable/
Apache License 2.0
747 stars 189 forks source link

[FEA] Allow construction of raft::mdarray in consteval context #1481

Open dejunlin opened 1 year ago

dejunlin commented 1 year ago

Is your feature request related to a problem? Please describe. Currently raft::mdarray's ContainerPolicy models std::vector, which in principle has constexpr ctor in C++20 but with compiler limitations (needs gcc 12 and clang 16). For the application I'm thinking about, it should be viable to come up with a ContainerPolicy type modeling std::array instead of std::vector. These are some of the issues I can see going down this route:

  1. raft::mdarray cannot be consteval constructed because its ctor depends on raft::resources, which in turn can't be constructed in a consteval context.
  2. raft::mdarray ctor assumes a two-argument call to ContainerPolicy::create(handle, size) where size might not be required for certain.
  3. related to point 2 is the lack of consistency check between the proposed std::array's compile-time size and those in the raft::extents because currently the latter determines the dynamic size as mentioned in point 2.
  4. raft::mdarray can't be list-initializer or constructed from a list of elements, which one would need in order to initialize a consteval mdarray.

Describe the solution you'd like Something like:

using T = int;
using E = raft::extents<uint64_t, 3, 2>;
using L = raft::layout_right;
using A = raft::host_accessor<raft::array_policy<T, 3*2>>; 
// or using A = raft::device_accessor<raft::array_policy<T, 3*2>>; since we have <cuda/std/array>
using Array = raft::mdarray<T, E, L, A>;
constexpr Array a({1,2,3,4,5,6}); // or the equivalent

Describe alternatives you've considered In principle, std::vector should work given enough compiler support but the device mdarray uses a different container policy than std::vector, which I assumes is not able to be consteval constructed. It would be nice to have a unified interface to all this if possible.

Additional context N/A

mhoemmen commented 1 year ago

Would you consider redesigning based on P1684R4, which has a std::array special case and dispatches directly to the container for construction (instead of using a ContainerPolicy, which was dropped in R1)?

dejunlin commented 1 year ago

Hi @mhoemmen , thanks for the suggestion! Are you talking about this specifically so that this can be std::array? Is there some check for the consistency between extents and std::array::size?

Would it make sense to make a generic container policy that maps the element and extents types to either dynamic allocation or static allocation? This policy itself then specialize also to different memory type, e.g., host vs device memory

mhoemmen commented 1 year ago

@dejunlin wrote:

Are you talking about this specifically so that this can be std::array?

Yes, though I think of this in terms of the proposal itself, not necessarily in terms of the implementation.

Is there some check for the consistency between extents and std::array::size?

Yes. The proposal includes phrases like this: "if container_type is a specialization of array, map.required_span_size() <= size(container_type()) is true."

Would it make sense to make a generic container policy that maps the element and extents types to either dynamic allocation or static allocation?

The question isn't really "dynamic allocation" or "static allocation." For example, static_vector (see P0843) uses static allocation, but its size (up to the capacity) can be specified as a run-time size_t value.

Instead, the question is, "if you pass its constructor a size_t, does the resulting container have a size no less than that?" That means more than just "you can construct it with a size_t." In particular, std::array is an aggregate. This means that you could actually construct it with a size_t, but that could just become one of the elements; it wouldn't set the array's size. Custom containers might be constructible from a size_t, but that's not enough; we would need to know the postcondition on the constructor, that the container has at least that size.

The Standard doesn't currently have a concept or a set of requirements for this. A "contiguous container" can be any container whose iterators meet certain requirements (see https://eel.is/c++draft/container.reqmts#68 ). That doesn't say anything about how to construct the container. Conversely, there's no "std::array-like" concept in the Standard, that we could use to make mdarray behave the same for cuda::std::array as for std::array.

P1684 coauthors have been avoiding getting into the business of defining new container concepts, because it's a tarpit of LEWG and LWG discussion. (For feedback from the last LEWG review, see https://github.com/ORNL/cpp-proposals-pub/issues/330 .) That's why P1684 singles out std::array as special, and assumes that any other container_type can be constructed from a size_t and that the resulting container has size at least that. LEWG has asked for more wording about that in the next revision of the paper, which we haven't had time to write yet.

This policy itself then specialize also to different memory type, e.g., host vs device memory

That's a separate question. P1684 addresses this in two ways.

  1. mdarray has constructors that take an allocator.
  2. mdarray has constructors that take const container_type& or container_type&&, so you can create the container separately and copy (const container_type&) or move (container_type&&) it into the mdarray.

Way (2) raises the question of allocator propagation on copy or move (see propagate_on_container_copy_assignment and propagate_on_container_move_assignment in allocator_traits).

See also https://github.com/ORNL/cpp-proposals-pub/issues/263 ("mdarray iteration over elements needs a parallel execution policy").

dejunlin commented 1 year ago

@mhoemmen

Instead, the question is, "if you pass its constructor a size_t, does the resulting container have a size no less than that?"

What is the real usage case for allocating a larger capacity than the logical size? Are we talking about functionalities such as std::vector::push_back?

Back to your previous suggestion about delegating to std::array construction, does it require the user to specify the static size of std::array? It seems the requirements of setting both extents<I, ns, ...> and std::array<T, nTotal> are cumbersome. My idea was to abstract that away from the user so that they don't have to deal with the extra-level of complexity by having to set the Container type in the most general usage case.

dejunlin commented 1 year ago

BTW, I realize I am not up to date with the mdarray proposal so I need to do more reading. Do you have slides or presentation on the topic with high-level ideas or plans presented as visuals so that it makes it easier to grasp the concepts? I found it easy to get lost in reading the open-std proposal text, e.g., when dealing with the details of which revision does what or what components are responsible for solving what problems

mhoemmen commented 1 year ago

@dejunlin I'll contact you offline.

cjnolet commented 1 year ago

@dejunlin @mhoemmen just FYI, we do have it on our todo-list to update our implementation of mdarray to match the most recent proposal version. We can certainly increase the priority of updating the mdarray implementation if that would help with this issue.

@dejunlin to help w/ items 1 and 2, we (RAFT) generally suggest using the factory functions to create mdarray/mdspan instances instead of going directly through the constructors. However, we could add a no-arg create() member to the ContainerPolicy in the meantime and a no-arg constructor to the mdarray. Would that be helpful for your immediate need? I also can't see why we shouldn't be able to support items 3 and 4 in the meantime as well.

@divyegala do you have any additional thoughts here?

dejunlin commented 1 year ago

@cjnolet Thanks for the suggestions! I was just talking to Mark to clear my thoughts on a deeper level.

However, we could add a no-arg create() member to the ContainerPolicy in the meantime and a no-arg constructor to the mdarray

I think this should be a forwarding method instead that takes whatever the underlying container needs, as @divyegala suggested on the slack thread.

we (RAFT) generally suggest using the factory functions to create mdarray/mdspan instances instead of going directly through the constructors

I think the ctor is often necessary in a lower-level metaprogramming context. I have encountered some minor issues using factory functions but I don't remember the specifics now.

cjnolet commented 1 year ago

I think this should be a forwarding method instead that takes whatever the underlying container needs, as @divyegala suggested on the slack thread.

@dejunlin just to make the solution explicit here in the Github issue (so that we can verify correctness and completeness), do you mind elaborating on your thoughts here a little bit?

dejunlin commented 1 year ago

just to make the solution explicit here in the Github issue (so that we can verify correctness and completeness), do you mind elaborating on your thoughts here a little bit?

I don't think I have enough details about the current implementation to make specific suggestions, but it seems now the ctor is taking:

mdarray a{handle, layout, policy};

where the arguments might not be necessary for std::array construction so my suggestion was to provide something like:

template < class ... Args >
explicit mdarray(Args&& ... args) : ... {}

with some metaprogramming to check argument types, consistency and to construct the respective members.

mhoemmen commented 1 year ago

@dejunlin Would the idea of this variadic mdarray constructor be to extract any extents or layout mapping arguments, then forward the rest to the container? If so, then you might want to familiarize yourself with the three different conventions for passing an allocator into a container. Please see the "Uses-allocator construction" section of https://en.cppreference.com/w/cpp/memory/uses_allocator .