kokkos / mdspan

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

Canonical way of handling read-only data in interfaces taking mdspans #331

Open pauleonix opened 3 months ago

pauleonix commented 3 months ago

I have not found any information on this topic. The dot_product example just doesn't take care of const-correctness at all. With span one can just have span<T const> and an implicit cast happens when passing a span<T>. With mdspan things are less straightforward when one wants to stay generic with respect to the AccessorPolicy, i.e.

template <typename T, typename IndexT, std::size_t x_extent,
          std::size_t y_extent, class L, template <typename> class A>
void bar(std::mdspan<T const, std::extents<IndexT, x_extent, y_extent>, L, A<T const>> input);

void foo() {
    auto mat = std::mdspan<float, std::dextents<int, 2>>{};
    bar(mat);
}

doesn't compile as the compiler can't both deduce L and A and cast from an element type of float to an element type of float const. The best solution I could come up with is some kind of explicit cast like

template <typename T, class Extents, class Layout,
          template <typename> class Accessor>
auto read_only_cast(std::mdspan<T, Extents, Layout, Accessor<T>> other) {
    using CT = std::add_const_t<T>;
    return std::mdspan<CT, Extents, Layout, Accessor<CT>>{other};
}

which then allows

void foo() {
    auto mat = std::mdspan<float, std::dextents<int, 2>>{};
    bar(read_only_cast(mat));
}

This naive implementation is probably flawed in its assumption that every Accessor has a single template argument (can be fixed by adding a variadic overload of read_only_cast) and generally I would not expect that the implementation of this kind of facility is left to the user. Is there a simpler way that avoids the need to explicitly specify any template parameters to bar()? I think https://github.com/kokkos/mdspan/blob/a9c54ccd8254cc3d159fdf2adf650dca4e048c97/examples/dot_product/dot_product.cpp#L48-L56 should be updated to reflect any canonical way of specifying that a function takes read-only, input mdspans.

pauleonix commented 3 months ago

After further consideration my read_only_cast is probably even worse than I thought. I came up with an IteratorAccessor templated on the iterator type instead of the element type, i.e. one generally cannot even assume there to be the element type as first template argument to the accessor or at least I have not seen it mandated anywhere. The only idea to fix read_only_cast would be to introduce aConstAccessor that wraps a given accessor and adds const to its element_type and reference.

pauleonix commented 3 months ago

This is what I came up with:

template <class Accessor>
class ConstAccessorAdaptor {
    public:
    using element_type = std::add_const_t<typename Accessor::element_type>;
    using reference = std::add_const_t<typename Accessor::reference>;
    using data_handle_type = typename Accessor::data_handle_type;
    using offset_policy = ConstAccessorAdaptor;

    ConstAccessorAdaptor(Accessor const& to_be_wrapped) : wrapped_{to_be_wrapped} {}

    constexpr reference access(data_handle_type p, std::size_t i) const {
        return wrapped_.access(p, i);
    }

    constexpr data_handle_type offset(data_handle_type p,
                                      std::size_t i) const {
        return wrapped_.offset(p, i);
    }
    private:
    Accessor wrapped_{};
};

template <typename ElementT, class Extents, class Layout, class Accessor>
constexpr auto read_only_cast(cuda::std::mdspan<ElementT, Extents, Layout, Accessor> other) {
    using CElementT = std::add_const_t<ElementT>;
    return cuda::std::mdspan<CElementT, Extents, Layout, ConstAccessorAdaptor<Accessor>>{other};
}
mhoemmen commented 3 months ago

Hi! We've discussed this before, but just haven't had a chance to document it in a nice way. The issue is that accessors in general don't promise that a "const version of the accessor" even exists. This doesn't necessarily even make sense for some custom accessors.

Your solution ConstAccessorAdaptor is reasonable -- it's analogous to the approach of std::ranges::views::as_const -- but it's not fully general. For example, reference may be a proxy reference type, so add_const_t<reference> might not make sense or give you the effect of a "read-only reference." (Consider, e.g., atomic_ref<value_type>.) You'll actually need a wrapper proxy reference type that forbids assignment.

If I were implementing this, I would do something like this.

template<class Accessor>
as_const(const Accessor& acc) {
  return const_accessor_adapter<Accessor>(acc);
}

// Specialize for default_accessor
template<class ElementType>
as_const(const default_accessor<ElementType>& acc) {
  return default_accessor<std::add_const_t<ElementType>>{};
}

template<class ElementType, class Extents, class Layout, class Accessor>
constexpr auto as_const(mdspan<ElementType, Extents, Layout, Accessor> x) {
  auto const_acc = as_const(x.accessor());
  using const_data_handle_type = typename decltype(const_acc)::data_handle_type;
  return mdspan{const_data_handle_type{x.data_handle()}, x.mapping(), const_acc}; // use CTAD
}
pauleonix commented 3 months ago

Yeah, I somewhat assumed that reference wrappers would support reference semantics (including contents not being modifiable when the wrapper is const) like e.g. thrust::device_reference. But atomic_ref is a very good point.

mhoemmen commented 3 months ago

@pauleonix We've been asked the question "How do I make a view-of-const mdspan from a general mdspan?" enough times that we really should write a C++ Standard Library proposal to do it. Would you be interested in contributing?

pauleonix commented 3 months ago

@mhoemmen I would lie if I said that I wasn't interested, but I certainly don't have the headspace right now to work myself into the proposal process from zero.

But if someone would write a draft I would be happy to (proof-) read, discuss and give feedback or however else I can contribute on refining it while hopefully learning more about the process (and mdspan) for future opportunities.

nmm0 commented 2 months ago

@pauleonix We've been asked the question "How do I make a view-of-const mdspan from a general mdspan?" enough times that we really should write a C++ Standard Library proposal to do it. Would you be interested in contributing?

Well this issue just popped up in Kokkos, I made a rather kludgy fix: https://github.com/kokkos/kokkos/pull/7053/files#diff-fdbb0c6e835c8457ab0f560f668ba721fb72de7f4f68fc9ec565c3b729608040 . Just thought I might add to more uses cases. I'd be up for helping with a paper on this. I'm busy with personal matters this month but I could work on one for Wroclaw

mhoemmen commented 2 months ago

@nmm0 Thanks for posting! Accessors in general don't promise to have any template parameters. Even if they have a template parameter, they don't promise that it is the element type. Even if they do have a template parameter that is the element type, then they don't promise that Accessor<add_const_t<Accessor::element_type>> does anything meaningful.

For example, Accessor::reference may be a proxy reference type and not just Accessor::element_type&. Thus, add_const_t<reference> might not make sense, and/or it might not give you the effect of a "read-only reference." (Consider, e.g., atomic_ref<value_type>.) You'll actually need a wrapper proxy reference type that forbids assignment. This is analogous to the approach of std::ranges::views::as_const, which specializes for known types (e.g., span<T> -> span<const T>) and uses a generic iterator wrapper otherwise.

nmm0 commented 2 months ago

@nmm0 Thanks for posting! Accessors in general don't promise to have any template parameters. Even if they have a template parameter, they don't promise that it is the element type. Even if they do have a template parameter that is the element type, then they don't promise that Accessor<add_const_t<Accessor::element_type>> does anything meaningful.

For example, Accessor::reference may be a proxy reference type and not just Accessor::element_type&. Thus, add_const_t<reference> might not make sense, and/or it might not give you the effect of a "read-only reference." (Consider, e.g., atomic_ref<value_type>.) You'll actually need a wrapper proxy reference type that forbids assignment. This is analogous to the approach of std::ranges::views::as_const, which specializes for known types (e.g., span<T> -> span<const T>) and uses a generic iterator wrapper otherwise.

Yeah this was a bit of a hack since we are only expected some specific accessors in our tests. I think specialization might be the best thing here because it can also give meaningful errors for accessors that may not work with const types.