Any advices on using std::mdspan with Kokkos submdspan?

[weslleyspereira]

With \<T>LAPACK (https://github.com/tlapack/tlapack), we have been using mdspan from this repository. We wanted to be able to test std::mdspan from C++23 stdlib with Kokkos submdspan class. We currently rely on submdspan. Some questions arise:

• Would you recommend using std::mdspan from C++23 stdlib with Kokkos submdspan? Is that something that you have tested?
• Is there any guidelines on how to better integrate this repo with C++23?

Thanks for the Great job!!

[mhoemmen]

Thanks for the kind words! : - )

@weslleyspereira wrote:

We wanted to be able to test std::mdspan from C++23 stdlib with Kokkos submdspan class. We currently rely on submdspan.

I haven't tried doing this, but I'm not sure if it's possible without writing at least a little bit of code. The main issue is that the reference implementation of submdspan (in https://github.com/kokkos/mdspan ) returns mdspan without namespace qualification, as you can see here. This means that the only way to get it to return std::mdspan and not (say) Kokkos::mdspan would be to define the appropriate macros so that MDSPAN_IMPL_STANDARD_NAMESPACE is std. However, if you do that, you'll be getting conflicting definitions of the reference mdspan implementation with the libc++ implementation. For example, the reference implementation of layout_left includes the hidden friend submdspan_mapping here, but the libc++ implementation might not have that (at least not for C++23).

If I had to do this, I would try writing a little wrapper around submdspan that takes and returns std::mdspan, but dispatches to Kokkos::submdspan underneath. You'll need to call your wrapper always with namespace qualification, so that you don't get the wrong submdspan via argument-dependent lookup.

namespace my {

// Pull strided_slice into your namespace.
// This won't work if libc++ already defines it.
// In that case, you'll need a to_kokkos / to_std conversion pair,
// and you'll need to apply it to the `slices...` parameter pack of submdspan below.
using ::Kokkos::strided_slice;

// Convert std::mdspan to Kokkos::mdspan
template<class ElementType, class Extents, class Layout, class Accessor>
MDSPAN_INLINE_FUNCTION constexpr auto
to_kokkos(const ::std::mdspan<ElementType, Extents, Layout, Accessor>& x)
{
  return ::Kokkos::mdspan{x.data_handle, to_kokkos(x.mapping()), to_kokkos(x.accessor())};
}

template<class ElementType>
MDSPAN_INLINE_FUNCTION constexpr ::Kokkos::default_accessor<ElementType>
to_kokkos(const ::std::default_accessor<ElementType>&)
{
  return {};
}

// repeat for the other layout mappings.  layout_stride::mapping
// will need two arguments: x.extents() and x.strides().
template<class Extents>
MDSPAN_INLINE_FUNCTION constexpr auto
to_kokkos(const ::std::layout_left::template mapping<Extents>& x)
{
  return ::Kokkos::layout_left::template mapping<Extents>{to_kokkos(x.extents())};
}

template<IndexType, ::std::size_t ... Exts>
MDSPAN_INLINE_FUNCTION constexpr auto
to_kokkos(const ::std::extents<IndexType, Exts...>& x)
{
  return [&x]<std::size_t ... Indices> (std::index_sequence<Indices>) {
    return ::Kokkos::extents<IndexType, Exts...>{x.extents(Indices...)...};
  } (std::make_index_sequence<x.rank()>());
}

// ... define to_std overloads for Kokkos things analogously ...

template <class ElementType, class Extents, class LayoutPolicy,
          class AccessorPolicy, class... SliceSpecifiers>
MDSPAN_INLINE_FUNCTION
constexpr auto
submdspan(const mdspan<ElementType, Extents, LayoutPolicy, AccessorPolicy>& src,
  SliceSpecifiers... slices)
{
  return to_std(::Kokkos::submdspan(to_kokkos(src), slices...));
}

} // namespace my

[crtrott]

Couple things: use our mdspan from the Kokkos namespace then it will just work. I will start putting submdspan into stdlib soon. We can certainly look at interoperability more.

[weslleyspereira]

Thanks for the comments!

I agree it is better to use @crtrott suggestion and just use your mdspan + submdspan.

I will start putting submdspan into stdlib soon.

I'll keep an eye out for that. Thanks!

[mhoemmen]

I agree it is better to use @crtrott suggestion and just use your mdspan + submdspan.

I agree; you'll have a Bad Time if the types get mixed up. Christian probably has plans for improving std::mdspan interoperability as it becomes more widely available.