[RFC] improve interoperability with `mdspan`

[fwyzard]

To make it easier to pass an alpaka::Buf to a kernel and use it with the alpaka::uniformElements loops, I am thinking of looking into:

• defining alpaka::MdSpan<T, ...::dextents<...>, ...> as an alias for std::mdspan (c++23) or std::experimental::mdspan (from Kokkos's reference implementation), with the template arguments corresponding to what is returned by alpaka::experimental::getMdSpan<Buf> edit this exists already :-)
• making alpaka::Buf implicitly convertible to alpaka::MdSpan (with the correct template arguments)
• making alpaka::Vec implicitly convertible to std::array

This way we should be able to write (pseudo)code like

struct Kernel {
  void operator()(Acc3D const& acc, alpaka::MdSpan<float, 3> data, Vec3D size) const {
    for (auto index: alpaka::uniformElements(acc, size)) {
      use(data[index]);
    }
  }
};

void function(Queue& queue) {
  Vec3D size = { 32, 16, 50 };
  auto buf = alpaka::allocAsyncBuf<float>(queue, size);
  auto grid = alpala::WorkDivElements<Acc3D>{...};
  alpaka::exec<Acc3D>(queue, grid, Kernel{}, buf, size);
}

@psychocoderHPC can you remind me what are the limitations of std::mdspan that prompted you to implement your own version for alpaka 3 ?

@SimeonEhrig I remember we discussed this some time ago, but I cannot find the links to your code any more :-(

comments, ideas or objections ?

[psychocoderHPC]

One reason I decided not to use std::mdspan for now is that operator()[] only supports std::array. I don't remember all the details, but I think there were many issues with implicit conversion. I was concerned about the fact that they do not support any tuple-like or array-like type for indexing. Initially, I was wrapping std::mdspan to support using alpaka::Vec, but then I decided that if I needed to wrap it, I could write my own span type.

Another feature I need in alpaka3 is compile-time alignment. This is most likely possible somewhere in the policies of mdspan, but the simpler solution was to use my own type.

For shared memory, I have a special MdSpan that works on top of a C MD-arrays, which results in better PTX code because some address calculations happen directly in the PTX load command. I never achieved this kind of PTX command, even when I used constexpr to calculate addresses.

std::mdspan is a C++23 feature, and with the Kokkos version, I had issues when using it on Godbolt because, in modern compilers, mdspan was not in the experimental namespace.

In short, for making progress, using a C++23 feature in C++20 was too time-consuming, and the benefit was not large enough. This does not mean that we should not use std::mdspan at some point, but currently, there are other things to do.

[fwyzard]

One reason I decided not to use std::mdspan for now is that operator()[] only supports std::array. I don't remember all the details, but I think there were many issues with implicit conversion.

Ah, I see... yes, after checking online, it looks like operator[] will not use implicit conversions - it requires the argument to be the exact type, or at most an inheriting type.

So - how about making alpaka::Vec<T, D> inherit from std::array<T, D::value> instead of containing a C array T[D::value] ?
From a quick check it looks like it could work - but maybe it's a change too intrusive for alpaka 2.x ?

[fwyzard]

After playing around with some code, I realise another issue:

making alpaka::Buf implicitly convertible to alpaka::MdSpan

won't help, because the implicit conversions are not used when the function (in our case, Kernel::operator()) is templated and has overloaded...

Or, at least, I couldn't get it to work 🤷🏻.

[fwyzard]

So here is another idea: when passing an object to a kernel, we could apply some predefined transformation, to be customised via a new trait.

For example:

alpaka::exec<Acc>(queue, grid, kernel, buffer, size);

could automatically convert the buffer argument, and call the kernel function as

kernel(acc, getMdSpan(buffer), size);

The conversion would be registered by defining

struct alpaka::KernelDecay<Buffer> {
    using type = alpaka::experimental::MdSpan<...>;

    type apply(Buffer& buffer) {
        return alpaka::experimental::getMdSpan(buffer);
    }
};

Suggestion for better names, function signatures, approach, etc, are warmly welcome.
The goal is just to simplify using alpaka buffers and views with alpaka kernels.

[SimeonEhrig]

@fwyzard Here is the prototype code, which I wrote. The basic idea is double cast. The buffer is casted to a wrapper type. Via optional template argument, I describe how the interface of the type looks like, e.g. raw pointer, std::span or std::mdspan. If the kernel argument is a generic type, the wrapper type is passed to kernel. If the kernel defines a specific type, like a raw pointer, the wrapper type is automatically casted to it. E.g. our wrapper type has the std::mdpsan interface. it will be casted to a raw pointer and the size of the object is only the size of the memory address.

#include <concepts>
#include <iostream>
#include <span>
#include <typeinfo>
#include <utility>
#include <vector>

constexpr auto getName(auto &&t) {
  if constexpr (requires { t.getName(); }) {
    return t.getName();
  } else {
    return typeid(t).name();
  }
}

template <typename TData>
class SpanAccessor : public std::span<TData, std::dynamic_extent> {
public:
  SpanAccessor(TData *data, std::size_t size)
      : std::span<TData, std::dynamic_extent>{data, size} {}

  TData &operator[](std::size_t const index) { return *(this)[index]; }


  constexpr auto getName() { return "SpanAccessor"; }

  operator std::span<TData, std::dynamic_extent>() const {
    return {this->data(), this->size()};
  }

  operator TData *() const { return this->data(); }
};


template <typename TData> struct PointerAccessor {
  TData *m_data;
  std::size_t m_size;

public:
  PointerAccessor(TData *data, std::size_t size) : m_data(data), m_size(size) {}

  constexpr auto getName() { return "PointerAccessor"; }

  TData &operator[](std::size_t const index) { return m_data[index]; }

  operator std::span<TData, std::dynamic_extent>() const {
    return {m_data, m_size};
  }

  operator TData *() const { return m_data; }
};

template <typename TData> class Buffer {
  std::vector<TData> m_data;

public:
  using Data = TData;
  Buffer(std::vector<TData> data) : m_data(data) {}

  constexpr auto getName() { return "Buffer"; }
  std::vector<TData> *data() { return m_data.data(); }
  auto size() const { return m_data.size(); }

  operator SpanAccessor<TData>() { return {m_data.data(), m_data.size()}; }
  operator PointerAccessor<TData>() { return {m_data.data(), m_data.size()}; }
};

template <typename T>
concept isBuffer = requires {
  typename std::decay_t<T>::Data;
  requires std::is_same_v<std::decay_t<T>,
                          Buffer<typename std::decay_t<T>::Data>>;
};

struct ToSpan {};

struct ToPointer {};

template <typename TAccessor, isBuffer TBuffer>
constexpr auto cast_helper(TBuffer &&buf) {
  using Buf = std::decay_t<TBuffer>;
  using Data = typename Buf::Data;

  if constexpr (std::is_same_v<TAccessor, ToSpan>) {
    return static_cast<SpanAccessor<Data>>(buf);
  } else if constexpr (std::is_same_v<TAccessor, ToPointer>) {
    return static_cast<PointerAccessor<Data>>(buf);
  } else {
    throw std::runtime_error("unsupported type");
  }
}

template <typename Accessor, typename T> auto getAccessorBuffer(T &&t) {
  if constexpr (isBuffer<T>) {
    return cast_helper<Accessor>(std::forward<T>(t));
  } else {
    return std::forward<T>(t);
  }
}

template <typename TAccessor = ToSpan, typename TFunc, typename... TArgs>
auto exec_cast(TFunc &&func, TArgs &&...arg) {
  return func(getAccessorBuffer<TAccessor>(std::forward<TArgs>(arg))...);
}

template <auto Indent = 0> void print_names(auto... args) {
  ((std::cout << std::string(Indent, ' ') << getName(args) << "\n"), ...);
}

template <typename T>
concept SpanLike = requires(T t) {
  t[0];
  t.size();
};

int main() {
  Buffer b{std::vector{1, 2, 3, 4}};

  auto generic = [](auto input) {
    std::cout << "generic: " << getName(input) << "\n";
    return input[0];
  };

  auto req_pointer = [](int *input) {
    std::cout << "req_pointer: " << getName(input) << "\n";
    return input[0];
  };

  auto req_span = [](std::span<int> input) {
    std::cout << "req_span: " << getName(input) << "\n";
    return input[0];
  };

  std::cout << "=== exec_cast ===\n";

  std::cout << "\nexec_cast(req_pointer, b);\n  ";
  exec_cast(req_pointer, b);

  std::cout << "\nexec_cast(req_span, b);\n  ";
  exec_cast(req_span, b);

  std::cout << "\nexec_cast(generic, b);\n  ";
  exec_cast(generic, b);

  std::cout << "\nexec_cast<ToPointer>(generic, b);\n  ";
  exec_cast<ToPointer>(generic, b);

  std::cout << "\nexec_cast multi arguments\n";
  exec_cast([](int a, auto b, std::span<int> c, int *d, auto e,
               auto f) { print_names<2>(a, b, c, d, e, f); },
            1, b, b, b, 1.f, b);

  auto worker_span = [](SpanLike auto input) {
    std::cout << "  Input type: " << getName(input) << "\n";
    std::cout << "  process: ";
    for (auto i = 0; i < input.size(); ++i) {
      std::cout << input[i] << " ";
    }
    std::cout << "\n";
  };

  std::cout << "\nworker_span(): \n";
  exec_cast(worker_span, b);

  auto worker_pointer = [](int *input, std::size_t const size) {
    std::cout << "  Input type: " << getName(input) << "\n";
    std::cout << "  process: ";
    for (auto i = 0; i < size; ++i) {
      std::cout << input[i] << " ";
    }
    std::cout << "\n";
  };

  std::cout << "\nworker_pointer(): \n";
  exec_cast(worker_pointer, b, b.size());

  return 0;
}

[SimeonEhrig]

@psychocoderHPC implemented it already in alpaka3, except he cast to alpaka::mdspan all the time.
alpaka-group/alpaka3#124