diff --git a/docs/advanced/pycpp/numpy.rst b/docs/advanced/pycpp/numpy.rst index 57a52f6a8..fbcebc421 100644 --- a/docs/advanced/pycpp/numpy.rst +++ b/docs/advanced/pycpp/numpy.rst @@ -239,27 +239,13 @@ by the compiler. The result is returned as a NumPy array of type The scalar argument ``z`` is transparently replicated 4 times. The input arrays ``x`` and ``y`` are automatically converted into the right types (they are of type ``numpy.dtype.int64`` but need to be ``numpy.dtype.int32`` and -``numpy.dtype.float32``, respectively) +``numpy.dtype.float32``, respectively). -Sometimes we might want to explicitly exclude an argument from the vectorization -because it makes little sense to wrap it in a NumPy array. For instance, -suppose the function signature was +.. note:: -.. code-block:: cpp - - double my_func(int x, float y, my_custom_type *z); - -This can be done with a stateful Lambda closure: - -.. code-block:: cpp - - // Vectorize a lambda function with a capture object (e.g. to exclude some arguments from the vectorization) - m.def("vectorized_func", - [](py::array_t x, py::array_t y, my_custom_type *z) { - auto stateful_closure = [z](int x, float y) { return my_func(x, y, z); }; - return py::vectorize(stateful_closure)(x, y); - } - ); + Only arithmetic, complex, and POD types passed by value or by ``const &`` + reference are vectorized; all other arguments are passed through as-is. + Functions taking rvalue reference arguments cannot be vectorized. In cases where the computation is too complicated to be reduced to ``vectorize``, it will be necessary to create and access the buffer contents diff --git a/include/pybind11/cast.h b/include/pybind11/cast.h index 180f48d2a..0a8ef8318 100644 --- a/include/pybind11/cast.h +++ b/include/pybind11/cast.h @@ -15,6 +15,7 @@ #include "descr.h" #include #include +#include NAMESPACE_BEGIN(pybind11) NAMESPACE_BEGIN(detail) diff --git a/include/pybind11/common.h b/include/pybind11/common.h index 8d04c29bc..387300922 100644 --- a/include/pybind11/common.h +++ b/include/pybind11/common.h @@ -367,6 +367,12 @@ template struct make_index_sequence_impl <0, S...> { typedef index_ template using make_index_sequence = typename make_index_sequence_impl::type; #endif +/// Make an index sequence of the indices of true arguments +template struct select_indices_impl { using type = ISeq; }; +template struct select_indices_impl, I, B, Bs...> + : select_indices_impl, index_sequence>, I + 1, Bs...> {}; +template using select_indices = typename select_indices_impl, 0, Bs...>::type; + /// Backports of std::bool_constant and std::negation to accomodate older compilers template using bool_constant = std::integral_constant; template struct negation : bool_constant { }; diff --git a/include/pybind11/numpy.h b/include/pybind11/numpy.h index 2395b1f40..1e12ede19 100644 --- a/include/pybind11/numpy.h +++ b/include/pybind11/numpy.h @@ -1289,8 +1289,8 @@ public: return *this; } - template const T& data() const { - return *reinterpret_cast(m_common_iterator[K].data()); + template T* data() const { + return reinterpret_cast(m_common_iterator[K].data()); } private: @@ -1340,7 +1340,7 @@ enum class broadcast_trivial { non_trivial, c_trivial, f_trivial }; // buffer; returns `non_trivial` otherwise. template broadcast_trivial broadcast(const std::array &buffers, ssize_t &ndim, std::vector &shape) { - ndim = std::accumulate(buffers.begin(), buffers.end(), ssize_t(0), [](ssize_t res, const buffer_info& buf) { + ndim = std::accumulate(buffers.begin(), buffers.end(), ssize_t(0), [](ssize_t res, const buffer_info &buf) { return std::max(res, buf.ndim); }); @@ -1411,21 +1411,63 @@ broadcast_trivial broadcast(const std::array &buffers, ssize_t & broadcast_trivial::non_trivial; } +template +struct vectorize_arg { + static_assert(!std::is_rvalue_reference::value, "Functions with rvalue reference arguments cannot be vectorized"); + // The wrapped function gets called with this type: + using call_type = remove_reference_t; + // Is this a vectorized argument? + static constexpr bool vectorize = + satisfies_any_of::value && + satisfies_none_of::value && + (!std::is_reference::value || + (std::is_lvalue_reference::value && std::is_const::value)); + // Accept this type: an array for vectorized types, otherwise the type as-is: + using type = conditional_t, array::forcecast>, T>; +}; + template struct vectorize_helper { - typename std::remove_reference::type f; +private: static constexpr size_t N = sizeof...(Args); + static constexpr size_t NVectorized = constexpr_sum(vectorize_arg::vectorize...); + static_assert(NVectorized >= 1, + "pybind11::vectorize(...) requires a function with at least one vectorizable argument"); +public: template - explicit vectorize_helper(T&&f) : f(std::forward(f)) { } + explicit vectorize_helper(T &&f) : f(std::forward(f)) { } - object operator()(array_t... args) { - return run(args..., make_index_sequence()); + object operator()(typename vectorize_arg::type... args) { + return run(args..., + make_index_sequence(), + select_indices::vectorize...>(), + make_index_sequence()); } - template object run(array_t&... args, index_sequence index) { - /* Request buffers from all parameters */ - std::array buffers {{ args.request()... }}; +private: + remove_reference_t f; + + template using param_n_t = typename pack_element::call_type...>::type; + + // Runs a vectorized function given arguments tuple and three index sequences: + // - Index is the full set of 0 ... (N-1) argument indices; + // - VIndex is the subset of argument indices with vectorized parameters, letting us access + // vectorized arguments (anything not in this sequence is passed through) + // - BIndex is a incremental sequence (beginning at 0) of the same size as VIndex, so that + // we can store vectorized buffer_infos in an array (argument VIndex has its buffer at + // index BIndex in the array). + template object run( + typename vectorize_arg::type &...args, + index_sequence i_seq, index_sequence vi_seq, index_sequence bi_seq) { + + // Pointers to values the function was called with; the vectorized ones set here will start + // out as array_t pointers, but they will be changed them to T pointers before we make + // call the wrapped function. Non-vectorized pointers are left as-is. + std::array params{{ &args... }}; + + // The array of `buffer_info`s of vectorized arguments: + std::array buffers{{ reinterpret_cast(params[VIndex])->request()... }}; /* Determine dimensions parameters of output array */ ssize_t nd = 0; @@ -1433,61 +1475,79 @@ struct vectorize_helper { auto trivial = broadcast(buffers, nd, shape); size_t ndim = (size_t) nd; - ssize_t size = 1; - std::vector strides(ndim); - if (ndim > 0) { - if (trivial == broadcast_trivial::f_trivial) { - strides[0] = sizeof(Return); - for (size_t i = 1; i < ndim; ++i) { - strides[i] = strides[i - 1] * shape[i - 1]; - size *= shape[i - 1]; - } - size *= shape[ndim - 1]; - } - else { - strides[ndim-1] = sizeof(Return); - for (size_t i = ndim - 1; i > 0; --i) { - strides[i - 1] = strides[i] * shape[i]; - size *= shape[i]; - } - size *= shape[0]; - } + size_t size = std::accumulate(shape.begin(), shape.end(), (size_t) 1, std::multiplies()); + + // If all arguments are 0-dimension arrays (i.e. single values) return a plain value (i.e. + // not wrapped in an array). + if (size == 1 && ndim == 0) { + PYBIND11_EXPAND_SIDE_EFFECTS(params[VIndex] = buffers[BIndex].ptr); + return cast(f(*reinterpret_cast *>(params[Index])...)); } - if (size == 1) - return cast(f(*reinterpret_cast(buffers[Index].ptr)...)); + array_t result; + if (trivial == broadcast_trivial::f_trivial) result = array_t(shape); + else result = array_t(shape); - array_t result(shape, strides); - auto buf = result.request(); - auto output = (Return *) buf.ptr; + if (size == 0) return result; /* Call the function */ - if (trivial == broadcast_trivial::non_trivial) { - apply_broadcast(buffers, buf, index); - } else { - for (ssize_t i = 0; i < size; ++i) - output[i] = f((reinterpret_cast(buffers[Index].ptr)[buffers[Index].size == 1 ? 0 : i])...); - } + if (trivial == broadcast_trivial::non_trivial) + apply_broadcast(buffers, params, result, i_seq, vi_seq, bi_seq); + else + apply_trivial(buffers, params, result.mutable_data(), size, i_seq, vi_seq, bi_seq); return result; } - template - void apply_broadcast(const std::array &buffers, - buffer_info &output, index_sequence) { - using input_iterator = multi_array_iterator; - using output_iterator = array_iterator; + template + void apply_trivial(std::array &buffers, + std::array ¶ms, + Return *out, + size_t size, + index_sequence, index_sequence, index_sequence) { - input_iterator input_iter(buffers, output.shape); - output_iterator output_end = array_end(output); + // Initialize an array of mutable byte references and sizes with references set to the + // appropriate pointer in `params`; as we iterate, we'll increment each pointer by its size + // (except for singletons, which get an increment of 0). + std::array, NVectorized> vecparams{{ + std::pair( + reinterpret_cast(params[VIndex] = buffers[BIndex].ptr), + buffers[BIndex].size == 1 ? 0 : sizeof(param_n_t) + )... + }}; - for (output_iterator iter = array_begin(output); - iter != output_end; ++iter, ++input_iter) { - *iter = f((input_iter.template data())...); + for (size_t i = 0; i < size; ++i) { + out[i] = f(*reinterpret_cast *>(params[Index])...); + for (auto &x : vecparams) x.first += x.second; + } + } + + template + void apply_broadcast(std::array &buffers, + std::array ¶ms, + array_t &output_array, + index_sequence, index_sequence, index_sequence) { + + buffer_info output = output_array.request(); + multi_array_iterator input_iter(buffers, output.shape); + + for (array_iterator iter = array_begin(output), end = array_end(output); + iter != end; + ++iter, ++input_iter) { + PYBIND11_EXPAND_SIDE_EFFECTS(( + params[VIndex] = input_iter.template data() + )); + *iter = f(*reinterpret_cast *>(std::get(params))...); } } }; +template +vectorize_helper +vectorize_extractor(const Func &f, Return (*) (Args ...)) { + return detail::vectorize_helper(f); +} + template struct handle_type_name> { static PYBIND11_DESCR name() { return _("numpy.ndarray[") + npy_format_descriptor::name() + _("]"); @@ -1496,22 +1556,32 @@ template struct handle_type_name> { NAMESPACE_END(detail) -template -detail::vectorize_helper -vectorize(const Func &f, Return (*) (Args ...)) { - return detail::vectorize_helper(f); -} - +// Vanilla pointer vectorizer: template -detail::vectorize_helper +detail::vectorize_helper vectorize(Return (*f) (Args ...)) { - return vectorize(f, f); + return detail::vectorize_helper(f); } -template ::type::operator())>::type> +// lambda vectorizer: +template ::operator())>::type> auto vectorize(Func &&f) -> decltype( - vectorize(std::forward(f), (FuncType *) nullptr)) { - return vectorize(std::forward(f), (FuncType *) nullptr); + detail::vectorize_extractor(std::forward(f), (FuncType *) nullptr)) { + return detail::vectorize_extractor(std::forward(f), (FuncType *) nullptr); +} + +// Vectorize a class method (non-const): +template ())), Return, Class *, Args...>> +Helper vectorize(Return (Class::*f)(Args...)) { + return Helper(std::mem_fn(f)); +} + +// Vectorize a class method (non-const): +template ())), Return, const Class *, Args...>> +Helper vectorize(Return (Class::*f)(Args...) const) { + return Helper(std::mem_fn(f)); } NAMESPACE_END(pybind11) diff --git a/tests/test_numpy_vectorize.cpp b/tests/test_numpy_vectorize.cpp index 34197469c..b1f820813 100644 --- a/tests/test_numpy_vectorize.cpp +++ b/tests/test_numpy_vectorize.cpp @@ -20,6 +20,17 @@ std::complex my_func3(std::complex c) { return c * std::complex(2.f); } +struct VectorizeTestClass { + VectorizeTestClass(int v) : value{v} {}; + float method(int x, float y) { return y + (float) (x + value); } + int value = 0; +}; + +struct NonPODClass { + NonPODClass(int v) : value{v} {} + int value; +}; + test_initializer numpy_vectorize([](py::module &m) { // Vectorize all arguments of a function (though non-vector arguments are also allowed) m.def("vectorized_func", py::vectorize(my_func)); @@ -40,6 +51,22 @@ test_initializer numpy_vectorize([](py::module &m) { m.def("selective_func", [](py::array_t, py::array::c_style>) { return "Complex float branch taken."; }); + // Passthrough test: references and non-pod types should be automatically passed through (in the + // function definition below, only `b`, `d`, and `g` are vectorized): + py::class_(m, "NonPODClass").def(py::init()); + m.def("vec_passthrough", py::vectorize( + [](double *a, double b, py::array_t c, const int &d, int &e, NonPODClass f, const double g) { + return *a + b + c.at(0) + d + e + f.value + g; + } + )); + + py::class_ vtc(m, "VectorizeTestClass"); + vtc .def(py::init()) + .def_readwrite("value", &VectorizeTestClass::value); + + // Automatic vectorizing of methods + vtc.def("method", py::vectorize(&VectorizeTestClass::method)); + // Internal optimization test for whether the input is trivially broadcastable: py::enum_(m, "trivial") .value("f_trivial", py::detail::broadcast_trivial::f_trivial) diff --git a/tests/test_numpy_vectorize.py b/tests/test_numpy_vectorize.py index 7ae777227..362b03626 100644 --- a/tests/test_numpy_vectorize.py +++ b/tests/test_numpy_vectorize.py @@ -159,3 +159,45 @@ def test_trivial_broadcasting(): assert vectorized_func(1, y2, 1).flags.f_contiguous assert vectorized_func(z1[1::4, 1::4], y2, 1).flags.f_contiguous assert vectorized_func(y1[1::4, 1::4], z2, 1).flags.c_contiguous + + +def test_passthrough_arguments(doc): + from pybind11_tests import vec_passthrough, NonPODClass + + assert doc(vec_passthrough) == ( + "vec_passthrough(" + "arg0: float, arg1: numpy.ndarray[float64], arg2: numpy.ndarray[float64], " + "arg3: numpy.ndarray[int32], arg4: int, arg5: m.NonPODClass, arg6: numpy.ndarray[float64]" + ") -> object") + + b = np.array([[10, 20, 30]], dtype='float64') + c = np.array([100, 200]) # NOT a vectorized argument + d = np.array([[1000], [2000], [3000]], dtype='int') + g = np.array([[1000000, 2000000, 3000000]], dtype='int') # requires casting + assert np.all( + vec_passthrough(1, b, c, d, 10000, NonPODClass(100000), g) == + np.array([[1111111, 2111121, 3111131], + [1112111, 2112121, 3112131], + [1113111, 2113121, 3113131]])) + + +def test_method_vectorization(): + from pybind11_tests import VectorizeTestClass + + o = VectorizeTestClass(3) + x = np.array([1, 2], dtype='int') + y = np.array([[10], [20]], dtype='float32') + assert np.all(o.method(x, y) == [[14, 15], [24, 25]]) + + +def test_array_collapse(): + from pybind11_tests import vectorized_func + + assert not isinstance(vectorized_func(1, 2, 3), np.ndarray) + assert not isinstance(vectorized_func(np.array(1), 2, 3), np.ndarray) + z = vectorized_func([1], 2, 3) + assert isinstance(z, np.ndarray) + assert z.shape == (1, ) + z = vectorized_func(1, [[[2]]], 3) + assert isinstance(z, np.ndarray) + assert z.shape == (1, 1, 1)