Any plan to provide a native interface between numpy/cupy arrays and mdspan ?

[blegouix]

Hello, I asked for an AI to generate an interface between cupy arrays and cuda::std::mdspan:

namespace py = pybind11;

template <class MDSpan>
MDSpan cupy_array_to_mdspan(py::handle obj, const char *arg_name) {
  if (!py::hasattr(obj, "__cuda_array_interface__")) {
    throw py::type_error(std::string(arg_name) +
                         " must expose __cuda_array_interface__");
  }

  py::dict cai =
      py::reinterpret_borrow<py::dict>(obj.attr("__cuda_array_interface__"));

  // shape
  if (!cai.contains("shape")) {
    throw py::type_error(std::string(arg_name) +
                         ".__cuda_array_interface__ missing 'shape'");
  }
  py::tuple shape = cai["shape"].cast<py::tuple>();
  if (shape.size() != 3) {
    throw py::value_error(std::string(arg_name) + " must be a 3D array");
  }

  const int d0 = shape[0].cast<int>();
  const int d1 = shape[1].cast<int>();
  const int d2 = shape[2].cast<int>();

  // dtype
  if (!cai.contains("typestr")) {
    throw py::type_error(std::string(arg_name) +
                         ".__cuda_array_interface__ missing 'typestr'");
  }
  const std::string typestr = cai["typestr"].cast<std::string>();

  // Accept float32 only
  // Typical values are "<f4", "=f4", sometimes "|f4"
  if (!(typestr == "<f4" || typestr == "=f4" || typestr == "|f4")) {
    throw py::type_error(std::string(arg_name) +
                         " must have dtype float32; got typestr='" + typestr +
                         "'");
  }
  
  / data pointer
  if (!cai.contains("data")) {
    throw py::type_error(std::string(arg_name) +
                         ".__cuda_array_interface__ missing 'data'");
  }
  py::tuple data = cai["data"].cast<py::tuple>();
  if (data.size() != 2) {
    throw py::type_error(std::string(arg_name) +
                         ".__cuda_array_interface__['data'] must be a 2-tuple");
  }

  auto ptr_value = data[0].cast<std::uintptr_t>();
  float *ptr = reinterpret_cast<float *>(ptr_value);
  if (ptr == nullptr) {
    throw py::value_error(std::string(arg_name) + " has null device pointer");
  }

  // strides are in BYTES in __cuda_array_interface__
  cuda::std::array<int, 3> stride_elems{};
  if (cai.contains("strides") && !cai["strides"].is_none()) {
    py::tuple strides = cai["strides"].cast<py::tuple>();
    if (strides.size() != 3) {
      throw py::value_error(std::string(arg_name) +
                            " strides must have length 3");
    }

    constexpr std::size_t itemsize = sizeof(float);
    const auto s0b = strides[0].cast<std::ptrdiff_t>();
    const auto s1b = strides[1].cast<std::ptrdiff_t>();
    const auto s2b = strides[2].cast<std::ptrdiff_t>();

    if (s0b % itemsize || s1b % itemsize || s2b % itemsize) {
      throw py::value_error(std::string(arg_name) +
                            " byte strides are not divisible by sizeof(float)");
    }

    stride_elems[0] = static_cast<int>(s0b / itemsize);
    stride_elems[1] = static_cast<int>(s1b / itemsize);
    stride_elems[2] = static_cast<int>(s2b / itemsize);
  } else {
    // C-contiguous default
    stride_elems[2] = 1;
    stride_elems[1] = d2;
    stride_elems[0] = d1 * d2;
  }
  
  / data pointer
  if (!cai.contains("data")) {
    throw py::type_error(std::string(arg_name) +
                         ".__cuda_array_interface__ missing 'data'");
  }
  py::tuple data = cai["data"].cast<py::tuple>();
  if (data.size() != 2) {
    throw py::type_error(std::string(arg_name) +
                         ".__cuda_array_interface__['data'] must be a 2-tuple");
  }

  auto ptr_value = data[0].cast<std::uintptr_t>();
  float *ptr = reinterpret_cast<float *>(ptr_value);
  if (ptr == nullptr) {
    throw py::value_error(std::string(arg_name) + " has null device pointer");
  }

  // strides are in BYTES in __cuda_array_interface__
  cuda::std::array<int, 3> stride_elems{};
  if (cai.contains("strides") && !cai["strides"].is_none()) {
    py::tuple strides = cai["strides"].cast<py::tuple>();
    if (strides.size() != 3) {
      throw py::value_error(std::string(arg_name) +
                            " strides must have length 3");
    }

    constexpr std::size_t itemsize = sizeof(float);
    const auto s0b = strides[0].cast<std::ptrdiff_t>();
    const auto s1b = strides[1].cast<std::ptrdiff_t>();
    const auto s2b = strides[2].cast<std::ptrdiff_t>();

    if (s0b % itemsize || s1b % itemsize || s2b % itemsize) {
      throw py::value_error(std::string(arg_name) +
                            " byte strides are not divisible by sizeof(float)");
    }

    stride_elems[0] = static_cast<int>(s0b / itemsize);
    stride_elems[1] = static_cast<int>(s1b / itemsize);
    stride_elems[2] = static_cast<int>(s2b / itemsize);
  } else {
    // C-contiguous default
    stride_elems[2] = 1;
    stride_elems[1] = d2;
    stride_elems[0] = d1 * d2;
  }

Callable with:

PYBIND11_MODULE(projector_python, m) {
py::class_<Projector>(m, "Projector")
    .def(py::init<const ProjectorParams &>(), py::arg("params"))
    .def(
        "__call__",
        [](Projector &self, py::object projections_obj,
           py::object volume_obj) {
          auto projections =
              cupy_array_to_mdspan<Projector::ProjectionsMDSpan>(
                  projections_obj, "projections");
          auto volume = cupy_array_to_mdspan<Projector::VolumeMDSpan>(
              volume_obj, "volume");

          self(projections, volume);
        },
        py::arg("projections"), py::arg("volume"));
}

It compiles but I did not test it yet extensively so maybe it works, maybe it does not, but since cupy14 recently introduced the .mdspan() conversion function is there any plan to provide a native interface upon it in pybind11 ?

Regards,