From 7b74d2fa419c6fd1973bc2a69de6640d4721553b Mon Sep 17 00:00:00 2001 From: Wenzel Jakob Date: Wed, 7 Jun 2023 15:16:18 +0200 Subject: [PATCH] Simplification pass over the preceding commit --- include/nanobind/eigen/dense.h | 352 +++++++++++++++++++-------------- 1 file changed, 206 insertions(+), 146 deletions(-) diff --git a/include/nanobind/eigen/dense.h b/include/nanobind/eigen/dense.h index 7a4e6d42..4df7f48d 100644 --- a/include/nanobind/eigen/dense.h +++ b/include/nanobind/eigen/dense.h @@ -18,62 +18,90 @@ static_assert(EIGEN_VERSION_AT_LEAST(3, 3, 1), NAMESPACE_BEGIN(NB_NAMESPACE) -/// Types for func. arguments that are compatible with various flavors of arrays +/// Function argument types that are compatible with various array flavors using DStride = Eigen::Stride; template using DRef = Eigen::Ref; template using DMap = Eigen::Map; NAMESPACE_BEGIN(detail) +/// Determine the number of dimensions of the given Eigen type template -constexpr int num_dimensions = bool(T::IsVectorAtCompileTime) ? 1 : 2; +constexpr int ndim_v = bool(T::IsVectorAtCompileTime) ? 1 : 2; -template struct StrideExtr { - using Type = Eigen::Stride<0, 0>; +/// Extract the compile-time strides of the given Eigen type +template struct stride { + using type = Eigen::Stride<0, 0>; }; -template struct StrideExtr> { - using Type = StrideType; +template struct stride> { + using type = StrideType; }; -template using Stride = typename StrideExtr::Type; - -/// Is true for Eigen types that are known at compile-time to hold contiguous memory only, which includes all specializations of Matrix and Array, -/// and specializations of Map and Ref with according stride types and shapes. A (compile-time) stride of 0 means "contiguous" to Eigen. -template constexpr bool requires_contig_memory = - (Stride::InnerStrideAtCompileTime == 0 || Stride::InnerStrideAtCompileTime == 1) && - (num_dimensions == 1 || - Stride::OuterStrideAtCompileTime == 0 || - Stride::OuterStrideAtCompileTime != Eigen::Dynamic && Stride::OuterStrideAtCompileTime == T::InnerSizeAtCompileTime); - -/// Is true for StrideTypes that can describe the contiguous memory layout of the plain Eigen type T. -template constexpr bool can_map_contig_memory = - (StrideType::InnerStrideAtCompileTime == 0 || StrideType::InnerStrideAtCompileTime == 1 || StrideType::InnerStrideAtCompileTime == Eigen::Dynamic) && - (num_dimensions == 1 || - StrideType::OuterStrideAtCompileTime == 0 || - StrideType::OuterStrideAtCompileTime == Eigen::Dynamic || - StrideType::OuterStrideAtCompileTime == T::InnerSizeAtCompileTime); - -/// Alias ndarray for a given Eigen type, to be used by type_caster::from_python, which calls type_caster>::from_python. -/// If the Eigen type is known at compile-time to handle contiguous memory only, then this alias makes type_caster>::from_python -/// either fail or provide an ndarray with contiguous memory, triggering a conversion if necessary and supported by flags. -/// Otherwise, this alias makes type_caster>::from_python either fail or provide an ndarray with arbitrary strides, -/// which need to be checked for compatibility then. There is no way to ask type_caster for specific strides other than c_contig and f_contig. -/// Hence, if an Eigen type requires non-contiguous strides (at compile-time) and type_caster> provides an ndarray with unsuitable strides (at run-time), -/// then type_caster::from_python just fails. Note, however, that this is rather unusual, since the default stride type of Map requires contiguous memory, -/// and the one of Ref requires a contiguous inner stride, while it can handle any outer stride. -template using array_for_eigen_t = ndarray< - typename T::Scalar, +template struct stride> { + using type = StrideType; +}; + +template using stride_t = typename stride::type; + +/** \brief Identify types with a contiguous memory representation. + * + * This includes all specializations of ``Eigen::Matrix``/``Eigen::Array`` and + * certain specializations of ``Eigen::Map`` and ``Eigen::Ref``. Note: Eigen + * interprets a compile-time stride of 0 as contiguous. + */ +template +constexpr bool is_contiguous_v = + (stride_t::InnerStrideAtCompileTime == 0 || + stride_t::InnerStrideAtCompileTime == 1) && + (ndim_v == 1 || stride_t::OuterStrideAtCompileTime == 0 || + (stride_t::OuterStrideAtCompileTime != Eigen::Dynamic && + stride_t::OuterStrideAtCompileTime == T::InnerSizeAtCompileTime)); + +/// Identify types with a static or dynamic layout that support contiguous storage +template +constexpr bool can_map_contiguous_memory_v = + (stride_t::InnerStrideAtCompileTime == 0 || + stride_t::InnerStrideAtCompileTime == 1 || + stride_t::InnerStrideAtCompileTime == Eigen::Dynamic) && + (ndim_v == 1 || stride_t::OuterStrideAtCompileTime == 0 || + stride_t::OuterStrideAtCompileTime == Eigen::Dynamic || + stride_t::OuterStrideAtCompileTime == T::InnerSizeAtCompileTime); + +/* This type alias builds the most suitable 'ndarray' for the given Eigen type. + In particular, it + + - matches the underlying scalar type + - matches the number of dimensions (i.e. whether the type is a vector/matrix) + - matches the shape (if the row/column count is known at compile time) + - matches the in-memory ordering when the Eigen type is contiguous. + + This is helpful because type_caster> will then perform the + necessary conversion steps (if given incompatible input) to enable data + exchange with Eigen. + + A limitation of this approach is that ndarray does not support compile-time + strides besides c_contig and f_contig. If an Eigen type requires + non-contiguous strides (at compile-time) and we are given an ndarray with + unsuitable strides (at run-time), type casting will fail. Note, however, that + this is rather unusual, since the default stride type of Eigen::Map requires + contiguous memory, and the one of Eigen::Ref requires a contiguous inner + stride, while handling any outer stride. +*/ + +template +using array_for_eigen_t = ndarray< + Scalar, numpy, std::conditional_t< - num_dimensions == 1, + ndim_v == 1, shape<(size_t) T::SizeAtCompileTime>, shape<(size_t) T::RowsAtCompileTime, (size_t) T::ColsAtCompileTime>>, std::conditional_t< - requires_contig_memory, + is_contiguous_v, std::conditional_t< - num_dimensions == 1 || T::IsRowMajor, + ndim_v == 1 || T::IsRowMajor, c_contig, f_contig>, any_contig>>; @@ -92,7 +120,9 @@ template constexpr bool is_eigen_xpr_v = is_eigen_v && !is_eigen_plain_v && !is_eigen_sparse_v && !std::is_base_of_v, T>; -template struct type_caster && is_ndarray_scalar_v>> { +template +struct type_caster && + is_ndarray_scalar_v>> { using Scalar = typename T::Scalar; using NDArray = array_for_eigen_t; using NDArrayCaster = make_caster; @@ -100,16 +130,21 @@ template struct type_caster && i NB_TYPE_CASTER(T, NDArrayCaster::Name); bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { - NDArrayCaster caster; + // We're in any case making a copy, so non-writable inputs area also okay + using NDArrayConst = array_for_eigen_t; + make_caster caster; if (!caster.from_python(src, flags, cleanup)) return false; - const NDArray &array = caster.value; - if constexpr (num_dimensions == 1) + + const NDArrayConst &array = caster.value; + if constexpr (ndim_v == 1) value.resize(array.shape(0)); else value.resize(array.shape(0), array.shape(1)); - // array_for_eigen_t ensures that array holds contiguous memory. + + // The layout is contiguous & compatible thanks to array_for_eigen_t memcpy(value.data(), array.data(), array.size() * sizeof(Scalar)); + return true; } @@ -122,10 +157,10 @@ template struct type_caster && i } static handle from_cpp(const T &v, rv_policy policy, cleanup_list *cleanup) noexcept { - size_t shape[num_dimensions]; - int64_t strides[num_dimensions]; + size_t shape[ndim_v]; + int64_t strides[ndim_v]; - if constexpr (num_dimensions == 1) { + if constexpr (ndim_v == 1) { shape[0] = v.size(); strides[0] = v.innerStride(); } else { @@ -170,7 +205,7 @@ template struct type_caster && i policy == rv_policy::move ? rv_policy::reference : policy; object o = steal(NDArrayCaster::from_cpp( - NDArray(ptr, num_dimensions, shape, owner, strides), + NDArray(ptr, ndim_v, shape, owner, strides), array_rv_policy, cleanup)); return o.release(); @@ -178,7 +213,9 @@ template struct type_caster && i }; /// Caster for Eigen expression templates -template struct type_caster && is_ndarray_scalar_v>> { +template +struct type_caster && + is_ndarray_scalar_v>> { using Array = Eigen::Array; using Caster = make_caster; @@ -195,11 +232,23 @@ template struct type_caster && is_ } }; -/// Caster for Eigen::Map +/** \brief Type caster for ``Eigen::Map`` + + The ``Eigen::Map<..>`` type exists to efficiently access memory provided by a + caller. Given that, the nanobind type caster refuses to turn incompatible + inputs into a ``Eigen::Map`` when this would require an implicit + conversion. +*/ + template -struct type_caster, enable_if_t && is_ndarray_scalar_v>> { +struct type_caster, + enable_if_t && + is_ndarray_scalar_v>> { using Map = Eigen::Map; - using NDArray = array_for_eigen_t; + using NDArray = + array_for_eigen_t, + const typename Map::Scalar, + typename Map::Scalar>>; using NDArrayCaster = type_caster; static constexpr bool IsClass = false; static constexpr auto Name = NDArrayCaster::Name; @@ -208,13 +257,7 @@ struct type_caster, enable_if_t, enable_if_t is true, then StrideType is known at compile-time to only cope with contiguous memory. - // Then since caster.from_python has succeeded, caster.value now surely provides contiguous memory, and so its strides surely fit. - if constexpr (!requires_contig_memory) { - // A stride that is dynamic at compile-time copes with any stride at run-time. + // Check for memory layout compatibility of non-contiguous 'Map' types + if constexpr (!is_contiguous_v) { + // Dynamic inner strides support any input, check the fixed case if constexpr (StrideType::InnerStrideAtCompileTime != Eigen::Dynamic) { - // A stride of 0 at compile-time means "contiguous" to Eigen, which is always 1 for the inner stride. - int64_t expected_inner_stride = StrideType::InnerStrideAtCompileTime == 0 ? 1 : StrideType::InnerStrideAtCompileTime; - if (expected_inner_stride != (num_dimensions == 1 || !T::IsRowMajor ? caster.value.stride(0) : caster.value.stride(1))) + // A compile-time stride of 0 implies "contiguous" .. + int64_t is_expected = StrideType::InnerStrideAtCompileTime == 0 + ? 1 /* .. and equals 1 for the inner stride */ + : StrideType::InnerStrideAtCompileTime, + is_actual = caster.value.stride( + (ndim_v != 1 && T::IsRowMajor) ? 1 : 0); + + if (is_expected != is_actual) return false; } - if constexpr (num_dimensions == 2 && StrideType::OuterStrideAtCompileTime != Eigen::Dynamic) { - int64_t expected_outer_stride = - StrideType::OuterStrideAtCompileTime == 0 - ? T::IsRowMajor ? caster.value.shape(1) : caster.value.shape(0) - : StrideType::OuterStrideAtCompileTime; - if (expected_outer_stride != (T::IsRowMajor ? caster.value.stride(0) : caster.value.stride(1))) + + // Analogous check for the outer strides + if constexpr (ndim_v == 2 && StrideType::OuterStrideAtCompileTime != Eigen::Dynamic) { + int64_t os_expected = StrideType::OuterStrideAtCompileTime == 0 + ? caster.value.shape(T::IsRowMajor ? 1 : 0) + : StrideType::OuterStrideAtCompileTime, + os_actual = caster.value.stride(T::IsRowMajor ? 0 : 1); + + if (os_expected != os_actual) return false; } } @@ -247,10 +295,10 @@ struct type_caster, enable_if_t]; - int64_t strides[num_dimensions]; + size_t shape[ndim_v]; + int64_t strides[ndim_v]; - if constexpr (num_dimensions == 1) { + if constexpr (ndim_v == 1) { shape[0] = v.size(); strides[0] = v.innerStride(); } else { @@ -261,40 +309,34 @@ struct type_caster, enable_if_t, shape, handle(), strides), + NDArray((void *) v.data(), ndim_v, shape, handle(), strides), rv_policy::reference, cleanup); } StrideType strides() const { - constexpr int is = StrideType::InnerStrideAtCompileTime, - os = StrideType::OuterStrideAtCompileTime; + constexpr int IS = StrideType::InnerStrideAtCompileTime, + OS = StrideType::OuterStrideAtCompileTime; int64_t inner = caster.value.stride(0), outer; - if constexpr (num_dimensions == 1) + if constexpr (ndim_v == 1) outer = caster.value.shape(0); else outer = caster.value.stride(1); - if constexpr (num_dimensions == 2 && T::IsRowMajor) + if constexpr (ndim_v == 2 && T::IsRowMajor) std::swap(inner, outer); - // Compile-time strides of 0 must be passed as such to constructors of StrideType, to avoid assertions in Eigen. - if constexpr (is == 0) { - // Ensured by stride checks in from_python_: - // assert(inner == 1); + // Eigen may expect a stride of 0 to avoid an assertion failure + if constexpr (IS == 0) inner = 0; - } - if constexpr (os == 0) { - // Ensured by stride checks in from_python_: - // assert(num_dimensions == 1 || outer == (T::IsRowMajor ? int64_t(caster.value.shape(1)) : int64_t(caster.value.shape(0)))); + if constexpr (OS == 0) outer = 0; - } - if constexpr (std::is_same_v>) + if constexpr (std::is_same_v>) return StrideType(inner); - else if constexpr (std::is_same_v>) + else if constexpr (std::is_same_v>) return StrideType(outer); else return StrideType(outer, inner); @@ -302,85 +344,103 @@ struct type_caster, enable_if_t == 1) + if constexpr (ndim_v == 1) return Map(t.data(), t.shape(0), strides()); else return Map(t.data(), t.shape(0), t.shape(1), strides()); } }; +/** \brief Caster for Eigen::Ref + + Compared to the ``Eigen::Map`` type caster above, the reference caster + accepts a wider set of inputs when it is used in *constant reference* mode + (i.e., ``Eigen::Ref``). In this case, it performs stride conversions + (except for unusual non-contiguous strides) as well as conversions of the + underlying scalar type (if implicit conversions are enabled). -/// Caster for Eigen::Ref + For non-constant references, the caster matches that of ``Eigen::Map`` and + requires an input with the expected layout (so that changes can propagate to + the caller). +*/ template -struct type_caster, enable_if_t && is_ndarray_scalar_v>> { +struct type_caster, + enable_if_t && + is_ndarray_scalar_v>> { using Ref = Eigen::Ref; + + /// Potentially convert strides/dtype when casting constant references + static constexpr bool MaybeConvert = + std::is_const_v && + // Restrict to contiguous 'T' (limitation in Eigen, see PR #215) + can_map_contiguous_memory_v; + + /// Eigen::Map caster with fixed strides using Map = Eigen::Map; - using DMap = Eigen::Map; using MapCaster = make_caster; + + // Extended version taking arbitrary strides + using DMap = Eigen::Map; using DMapCaster = make_caster; - using DmapMatches = typename Eigen::internal::traits::template match::type; - static constexpr bool can_map_contig_mem = can_map_contig_memory; + + /** + * The constructor of ``Ref`` uses one of two strategies + * depending on the input. It may either + * + * 1. Create a copy ``Ref::m_object`` (owned by Ref), or + * 2. Reference the existing input (non-owned). + * + * When the value below is ``true``, then it is guaranteed that + * ``Ref()`` owns the underlying data. + */ + static constexpr bool DMapConstructorOwnsData = + !Eigen::internal::traits::template match::type::value; + static constexpr bool IsClass = false; - static constexpr auto Name = const_name>(DMapCaster::Name, MapCaster::Name); + static constexpr auto Name = + const_name(MapCaster::Name, DMapCaster::Name); + template using Cast = Ref; MapCaster caster; - DMapCaster dcaster; - - - /// In short: - /// - type_caster> supports no conversions, independent of flags. - /// - type_caster> - /// + supports stride conversions, independent of flags, except for uncommon strides. - /// + It additionally supports conversions to T::Scalar if flags say so, - /// and if either a cleanup_list is passed, or if Ref is guaranteed to map its own data. - /// - /// type_caster> supports stride conversions independent of flags, because if the intention was to not allow them, - /// then the bound function would most probably expect a Map instead of a Ref. - /// - /// Both Ref and Ref map data. - /// Like for Map, type_caster>::from_python does not support conversions, and for the same reasons. - /// But unlike Ref, instead of mapping external data, Ref may alternatively map data that it owns itself. - /// Ref then maps its member variable m_object, having copy-constructed it from the passed Eigen type. - /// The primary use case of Ref is as function argument that either maps the caller's data, or a suitably converted copy thereof. - /// Hence, unlike with Map and Ref, a Ref that maps a (converted) copy is intended, - /// and thus, type_caster>::from_python may support conversions. - /// It first calls the type_caster for matching strides, not supporting conversions. - /// If that fails, it calls the one for arbitrary strides. Since conversions to T::Scalar create a temporary ndarray, - /// conversions are supported only if flags say so, and if either a cleanup_list is passed (that keeps the temporary alive), - /// or if Ref is guaranteed to map its own data (having copied the temporary), which is ensured only if DmapMatches::value is false. - /// - /// Unfortunately, if src's scalar type needs to be converted, then the latter means that e.g. - /// cast>(src) succeeds, while - /// cast< DRef>(src) fails - - /// even though DRef would be expected to support a superset of the types supported by Ref. - /// - /// Ref::m_object holds contiguous memory, which Ref silently fails to map if this is impossible given StrideType - /// and the passed object's shape. If mapping fails, then Ref is left with mapping nullptr. - /// While this could be considered below, it is not done for efficiency reasons: - /// due to Ref's missing move constructor, its unusual copy constructor, and since C++ does not guarantee named return value optimizations, - /// the Ref would need to be created only for checking it, and created a second time for returning it, - /// which seems too costly for a Ref that owns its data. - /// Instead of checking thoroughly after construction, conversion fails if it is known at compile-time that mapping may fail, - /// even though it may actually succeed in some of these cases at run-time (e.g. StrideType::OuterStrideAtCompileTime==4, - /// and a row-major Matrix with a dynamic number of columns and 4 columns at run-time). - /// Once Ref defines a move constructor https://gitlab.com/libeigen/eigen/-/issues/2668, this restriction may be lifted. + struct Empty { }; + std::conditional_t dcaster; + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { - if constexpr (std::is_const_v) - return caster.from_python(src, flags, cleanup) || - can_map_contig_mem && - dcaster.from_python_(src, (!DmapMatches::value || cleanup) ? flags : flags & ~(uint8_t)cast_flags::convert, cleanup); - else - return caster.from_python(src, flags, cleanup); + // Try a direct cast without implicit conversion first + if (caster.from_python(src, flags, cleanup)) + return true; + + // Potentially convert strides/dtype when casting constant references + if constexpr (MaybeConvert) { + /* Generating an implicit copy requires some object to assume + ownership. During a function call, ``dcaster`` can serve that + role (this case is detected by checking whether ``cleanup`` is + defined). When used in other situatons (e.g. ``nb::cast()``), + the created ``Eigen::Ref<..>`` must take ownership of the copy. + This is only guranteed to work if DMapConstructorOwnsData. + + If neither of these is possible, we disable implicit + conversions. */ + + if (!cleanup && !DMapConstructorOwnsData) + flags &= ~(uint8_t) cast_flags::convert; + + if (dcaster.from_python_(src, flags, cleanup)) + return true; + } + + return false; } operator Ref() { - if constexpr (std::is_const_v) + if constexpr (MaybeConvert) { if (dcaster.caster.value.is_valid()) return Ref(dcaster.operator DMap()); + } + return Ref(caster.operator Map()); } - }; NAMESPACE_END(detail)