Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Make nanobind adapter backward compatible #1790

Merged
merged 3 commits into from
Jun 25, 2024
Merged

Make nanobind adapter backward compatible #1790

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
516a389
Make nanobind adapter backward compatible
havogt Jun 19, 2024
d2701d4
remove accidentally committed file
havogt Jun 19, 2024
cc8746b
address review comments
havogt Jun 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 22 additions & 12 deletions include/gridtools/storage/adapter/nanobind_adapter.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,25 +24,33 @@

namespace gridtools {
namespace nanobind_sid_adapter_impl_ {
#if NB_VERSION_MAJOR >= 2
using array_size_t = nanobind::ssize_t;
#else
using array_size_t = std::size_t;
#endif

// Use `-1` for dynamic stride, use an integral value for static stride.
template <nanobind::ssize_t... Values>
using stride_spec = std::integer_sequence<nanobind::ssize_t, Values...>;
inline constexpr array_size_t dynamic_size = -1;

// Use `dynamic_size` for dynamic stride, use an integral value for static
// stride.
template <array_size_t... Values>
using stride_spec = std::integer_sequence<array_size_t, Values...>;

template <class IndexSequence>
struct dynamic_strides_helper;

template <std::size_t... Indices>
struct dynamic_strides_helper<std::index_sequence<Indices...>> {
using type = stride_spec<(void(Indices), -1)...>;
using type = stride_spec<(void(Indices), dynamic_size)...>;
};

template <std::size_t N>
using fully_dynamic_strides = typename dynamic_strides_helper<std::make_index_sequence<N>>::type;

template <nanobind::ssize_t SpecValue>
auto select_static_stride_value(std::size_t dyn_value) {
if constexpr (SpecValue == -1) {
template <array_size_t SpecValue>
constexpr auto select_static_stride_value(std::size_t dyn_value) {
if constexpr (SpecValue == dynamic_size) {
return dyn_value;
} else {
if (SpecValue != dyn_value) {
Expand All @@ -52,20 +60,20 @@ namespace gridtools {
}
}

template <nanobind::ssize_t... SpecValues, std::size_t... IndexValues>
auto select_static_strides_helper(
template <array_size_t... SpecValues, std::size_t... IndexValues>
constexpr auto select_static_strides_helper(
stride_spec<SpecValues...>, const std::size_t *dyn_values, std::index_sequence<IndexValues...>) {

return gridtools::tuple{select_static_stride_value<SpecValues>(dyn_values[IndexValues])...};
}

template <nanobind::ssize_t... SpecValues>
auto select_static_strides(stride_spec<SpecValues...> spec, const std::size_t *dyn_values) {
template <array_size_t... SpecValues>
constexpr auto select_static_strides(stride_spec<SpecValues...> spec, const std::size_t *dyn_values) {
return select_static_strides_helper(spec, dyn_values, std::make_index_sequence<sizeof...(SpecValues)>{});
}

template <class T,
nanobind::ssize_t... Sizes,
array_size_t... Sizes,
class... Args,
class Strides = fully_dynamic_strides<sizeof...(Sizes)>,
class StridesKind = sid::unknown_kind>
Expand Down Expand Up @@ -93,7 +101,9 @@ namespace gridtools {

namespace nanobind {
using nanobind_sid_adapter_impl_::as_sid;
using nanobind_sid_adapter_impl_::dynamic_size;
using nanobind_sid_adapter_impl_::fully_dynamic_strides;
using nanobind_sid_adapter_impl_::stride_spec;

} // namespace nanobind
} // namespace gridtools
Loading