kokkos / mdspan

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

Optional 32-bit indexing for CUDA? #220

Closed VasuAgrawal closed 1 year ago

VasuAgrawal commented 1 year ago

I'm looking at switching some of my CUDA code over to use mdspan instead of hand-rolled accessors. One of the criteria I have, though, is that I'm able to perform most / all of the indexing and offset computations in 32-bit, rather than (usually) 64-bit that is size_t. This requirement is for register efficiency in CUDA cores, where the registers are 32-bit and not 64-bit. I read through this paper about mdspan, and tried looking briefly through the source code for extents, but I didn't see any options to switch from size_t to e.g. uint32_t for extents and indexing. Is this something that is currently supported and I just haven't found it yet, or something that may be supported in the future, or something that's out of scope?

Thanks!

mhoemmen commented 1 year ago

@VasuAgrawal The first template parameter of extents is the index type. You can use int32_t there for 32-bit extents. The result of .extents(r) is index_type, which would be int32_t in this case.

constexpr int m = 4;
constexpr int n = 5;
std::vector<float> v(m * n);
std::mdspan m{v.data(), std::extents<int, m, n>{}};
static_assert(std::is_same_v<decltype(m.extent(0), int>);