I guess I haven't noticed how confusing this was before. If i was a user looking at a function called reduce_cols_by_key, I'd be expecting that I'd pass my own reduction function in. However, if I were invoking such a function without a reduction function (where another overload were provided that accepts a lambda), I'd expect it to perform a sum by default. It doesn't look like we provide such an option, though. What makes things even more confusing to me is that the reduce_rows_by_key function performs a weighted sum.

I'm not saying we need to rename or remove these functions, but it would be useful if we had a github issue (and a todo here in this file that references it) to allow generic reduction functions. That would make this much more useful and also allow it to be extended for generalized row-norm computations).

Still waiting on todo / github issue here

Same as above- can you put a todo in this file and reference a github issue to add an overload w/ generic reduction functions? I believe the weighted reduction here should be using a fused multiply-add intrinsic but a generalized reduction would make this primitive more widely applicable to different problems.

It looks like the goal here is to use 32-bit indices if possible, when inverting the layout mapping to use a 1-D loop index. This can be done, but there are two correctness issues with your approach.

1. The right quantity to test here is out.required_span_size(), not out.size(). The layout mapping maps the input multidimensional index to the half-open interval of offsets [0, out.required_span_size()).

2. The layout_{left, right, stride}::mapping constructors generally have as a precondition that the required span size of the input extents (and strides, if applicable) be representable as a value of type index_type.

Here is an approach that would address these issues.

template<class T>
constexpr bool is_32_bit_integral_v = std::is_integral_v<T> && sizeof(T) == std::uint32_t;
template<class T>
constexpr bool is_greater_than_32_bit_integral_v = std::is_integral_v<T> && sizeof(T) > std::uint32_t;

if constexpr (is_32_bit_integral_v<typename OutType::index_type>) {
  // ... always call 32-bit version ...
} else if constexpr (is_greater_than_32_bit_integral_v<typename OutType::index_type>) {
  // ... test the value of `required_span_size()`; dispatch to 32-bit or index_type (64 or more bits) as needed ...
} else {
  // ... always use index_type, which is 16 bits or less here ...
}

You'll also want to check the index_type and required_span_size() of the other mdspan. The above approach has the advantage that it only compiles an inner kernel for index types that you actually use.

In point 2, what happens in extreme cases? Consider index_type=uint32_t with extents {2^32, 2}. In this case, will required_span_size() by representable by index_type or will it cause an overflow?

@divyegala required_span_size() is not representable by index_type in this case. For layout_left and layout_right, required_span_size() and size() are the same mathematically. The only difference is the return type (index_type resp. size_t). For layout_stride, though, required_span_size() can be greater than the size(). For other layouts (e.g., the "matrix of a single value" layout that maps all multidimensional indices to the offset zero), required_span_size() can be less than size().

Note that while it's UB for users to violate preconditions, implementations aren't required to check preconditions. The reference implementation of layout_left does not currently check preconditions, as you can see here, for instance. This means two things.

1. If someone gives you a layout_{left,right,stride}::mapping instance (e.g., in an mdspan), then you can assume that the precondition is satisfied.

2. If you are constructing a layout_{left,right,stride}::mapping instance (e.g., by constructing an mdspan with a pointer and extents), then you are responsible for ensuring that the precondition is satisfied.

@divyegala wrote:

In this case, will required_span_size() by representable by index_type or will it cause an overflow?

Those are two separate questions, actually! : - )

1. required_span_size() is not representable by index_type in this case.
2. Giving this extents object to layout_{left,right,stride}::mapping's constructor violates the constructor's precondition. It could overflow, or it could open a portal to the Awesome Dimension and let loose a swarm of nasal demons who search out precondition violators and boop them gently on the nose.

@mhoemmen thanks for the explanations! How do we really represent such edge cases and safely obtain the product of the extents? Sounds like size() is the safe way to obtain the product without violating any pre-conditions since it's representable by size_t?

@divyegala Gladly! : - )

How do we really represent such edge cases and safely obtain the product of the extents?

By the time the user has created a layout mapping, it's already too late. What I mean by that is that if required_span_size() doesn't fit index_type, then the user will likely get the wrong answer when they try to index into the mdspan.

In what follows in my comment, I'll distinguish between "the Preconditions in the spec" and "what the reference implementation does." The reference implementation currently does not check this precondition in the layout mapping. This means that it's possible for users to construct extents for which the mapping's required_span_size() can overflow.

We can prevent this by wrapping mdspan creation to check the extents object for potential overflow, before it goes into a layout mapping's constructor. It's not UB to construct, e.g., dextents<uint16_t, 2>( 2^{15} , 2^{15} ). We just need to intercept that naughty extents value before it goes into a layout mapping's constructor. Otherwise, the layout mapping has the freedom to do whatever it likes, including calling abort().

Our mdarray implementation's conversion to mdspan can also check, but again, we're probably better off making the wrapper explicit and not part of the mdarray proposal. WG21 likes Preconditions and wants violating them to be UB. If we want some specified behavior (e.g., throwing a particular exception, or calling terminate() after printing a helpful error message), then we'll have to implement that ourselves.

Looks like out.required_span_size() does not work. How do I access this from the layout?

It's unclear to me from reading this doc 1) whether or not the dots output is needed when inplace==true and 2) if data == dots when inplace==true. I felt the same while looking over the API for reduce. This might just require some additional documentation in these functions, though I wonder if it's an indication that an update to the API might be needed (I'm just really not sure what that would be as of yet).

It's unclear to me as well, and I mentioned it to another of @mhoemmen 's comments. Reading the impl did not make it any clearer for me.

Maybe we should do a separate overload for the inplace version like we do in raft::matrix. That makes it much more straightforward to the user that it's in place (and which array(s) end up being overwriten).

One thing I realized as I was going through the raft::matrix APIs is that they have an overload for the inplace version and I think that looks more clean because the arguments kind of speak for themselves (e.g no output argument and docs say inout).

Should I do the separate overload in this PR or wait for the move to raft::matrix?

I also really don't think inplace means what we think it means here. Looking at the kernel, it seems like all inplace=true does is run the reduction_op one more time over the output pointer.

Maybe we should do a separate overload for the inplace version like we do in raft::matrix. That makes it much more straightforward to the user that it's in place (and which array(s) end up being overwriten).