kokkos / mdspan

Reference implementation of mdspan targeting C++23
Other
413 stars 69 forks source link

submdspan by using indexes? #363

Closed danieltowner closed 1 week ago

danieltowner commented 2 weeks ago

Can we allow fewer indexes in an mdspan subscript operator to act as a short-hand for a call to submdspan?

My use case is where I have something like rank-3 mdspan, which effectively is representing an array of a 2d array. If I want to extract a single 2d array from that, I can do:

auto original = mdspan(data, std::extents{1024, 8, 8});

submdspan(original, n, std::full_extent, std::full_extent);

but that is a little verbose. This would be more concise:

original[n]

Effectively, if too few indexes are provided to mdspan::operator[] this is equivalent to handing off to submdspan instead, and substituting std::full_extent for the missing indexes. Maybe submdspan could default to std::full_extent when fewer slice arguments are supplied than the rank?

I thought this might have come up before but I haven't been able to find any discussion of it, so apologies if I have overlooked something.

mhoemmen commented 1 week ago

Thanks for the suggestions! As I understand them, they include two different points.

  1. Overload operator[] to do slicing (returning an mdspan that views a subset of elements) as well as element access
  2. Have slicing "auto-broadcast" by treating missing trailing slice specifiers as full_extent

We considered both of these, but decided against them.

Regarding (1), we wanted to separate "function that returns a reference to an element" from "function that returns a slice." This makes it easier for both compilers and software developers to see where element accesses take place. Custom accessors' reference types can be proxy references with nontrivial side effects, and taking slices in a tight inner loop can be expensive in some cases, so it's important that developers be able to distinguish access from slicing.

Regarding (2), while it's true that original[n] would be more concise, we prefer that readers can see the ranks of the input and result directly from the syntax. submdspan(original, n, full_extent, full_extent) clearly takes a rank-3 mdspan and returns a rank-2 mdspan. Another reason we decided against (2) is that either original[n] or submdspan(original, n) syntax might give an incorrect impression that mdspan has a preferred 1-D indexing scheme. That is, it might suggest that users could ask for "the n-th thing in an mdspan" and get a contiguous chunk of it (because the "n-th thing" would naturally be completely separate from the other "things"). This would express a preference for a particular layout, which would go against the layout-generic design of mdspan. For example, putting the implicit full_extent on the right would suggest that original is row-major, because that would make original[n] contiguous.

I hope this helps explain our reasoning! Please feel free to ask more questions.

danieltowner commented 1 week ago

Thanks for the explanation. That makes sense.