kokkos / mdspan

Reference implementation of mdspan targeting C++23
Other
398 stars 65 forks source link

Custom accessor accessing more than one elements #319

Closed christosa-rfpro closed 5 months ago

christosa-rfpro commented 5 months ago

Hi,

Consider an array of 4*N doubles arranged like this:

x0, x1, ..., xN, y0, y1, ..., yN, z0, z1, ..., zN, w0, w1, ..., wN

It might help to visualise the array in two dimensions:

x0, x1, ..., xN,
y0, y1, ..., yN,
z0, z1, ..., zN,
w0, w1, ..., wN

I want to use an mdspan to view that array as a collection of

struct v4d { double x, y, z, w; };

One way to achieve this is with the following accessor:

struct const_v4d_accessor
{
  using offset_policy = const_v4d_accessor;
  using element_type = v4d;
  using reference = element_type;
  using data_handle_type = const double*;

  int stride; // the accessor must "know" how many elements to skip to go from `xi` to `yi`, etc.

  constexpr reference access(data_handle_type p, std::size_t i) const noexcept
  {
      return v4d{p[i], p[i + stride], p[i + 2*stride], p[i + 3*stride]};
  }

  constexpr data_handle_type offset(data_handle_type p, std::size_t i) const noexcept
  {
      return p + i;
  }
};

For example:

const int N = /* the number of vectors (not known at compile-time) */;
std::vector<double> vector_components = get_vector_components(N);

assert(vector_components.size() == (4*N));

// Iterate through vector_components as if it were a v4d array
using ext_t = std::extents<int, std::dynamic_extent>;
auto map = std::layout_right::mapping{ext_t{N}};
auto vectors = std::mdspan{vector_components.data(), map, const_v4d_accessor{N}};

for (int i = 0; i < vectors.extent(0); ++i) {
  const v4d v = vectors[i];
  do_something_with(v);
}

I wonder if having const_v4d_accessor::stride abuses mdspan's interface because that information is typically part of the layout. Do you think this approach makes any sense? Is there a better alternative?

Thanks, Christos

mhoemmen commented 5 months ago

Greetings! I think this is a correct approach. The accessor needs the stride because the stride explains how to turn a (data handle, 1-D index) pair into a "reference" to the element. I agree that it's good to separate "information that belongs in the layout mapping" from "information that belongs in the accessor"; mdspan's design encourages that. However, in this case, the stride is truly part of how the accessor accesses elements.

The discussion in https://github.com/kokkos/mdspan/issues/117 solves a similar problem. There, the user has an array allocated as char (bytes), but representing some larger type. Here, you have an array allocated as double, but representing some larger type. Your struct v4d should have the same minimum alignment requirement as double, so as long as your storage array of double is correctly aligned for double, you shouldn't have to worry about misaligned access.

The typical mdspan idiom in this case would be to allocate a 2-D mdspan of double, like mdspan<double, extents<size_t, dynamic_extent, 4>, layout_left>, and then access (x, y, z, w) using submdspan(x, k, full_extent). There's nothing wrong with your approach, but a lot of users don't bother with creating custom accessors and proxy reference types. It's also easier to experiment with the performance effects of different layouts that way.

christosa-rfpro commented 5 months ago

Very interesting points, thank you!

mhoemmen commented 5 months ago

Thanks for commenting! Please feel free to reopen this issue or open a new one if you have any questions. Thanks!