diff --git a/python_bindings/src/halide/halide_/PyBuffer.cpp b/python_bindings/src/halide/halide_/PyBuffer.cpp index 90606677ca36..a59b4b1f043b 100644 --- a/python_bindings/src/halide/halide_/PyBuffer.cpp +++ b/python_bindings/src/halide/halide_/PyBuffer.cpp @@ -311,7 +311,7 @@ class PyBuffer : public Buffer<> { ~PyBuffer() override = default; }; -py::buffer_info to_buffer_info(Buffer<> &b, bool reverse_axes = true) { +py::buffer_info to_buffer_info(const Buffer<> &b, bool reverse_axes) { if (b.data() == nullptr) { throw py::value_error("Cannot convert a Buffer<> with null host ptr to a Python buffer."); } @@ -335,6 +335,48 @@ py::buffer_info to_buffer_info(Buffer<> &b, bool reverse_axes = true) { ); } +// Returns a pair [c_contig, f_contig], where: +// - If c_contig is true, the buffer is stored in the default order on the Halide side. +// - The first dim has stride 1 (innermost first). +// - This is true if `b` was constructed without passing anything to `storage_order`, +// or equivalently, if `storage_order` was [0, 1, 2, ...]. +// - If f_contig is true, the buffer is stored in reversed order on the Halide side. +// - The last dim has stride 1 (innermost last). +// - This is true if `b` was constructed with `storage_order` of [d-1, d-2, ..., 0]. +// - It is possible for a Buffer to be both C and F contiguous (e.g., a scalar or a +// 1D vector), or for a Buffer to be neither (e.g., storage_order=[1, 0, 2] for a 3D +// buffer). +// ELEPHANT: maybe I should just call it [densest_first, densest_last]. But that +// doesn't imply "contiguous". contiguous_densest_first? +std::pair is_any_contiguous(const Buffer<> &b) { + if (b.dimensions() == 0) { + return {true, true}; + } + + const int d = b.dimensions(); + + int c_stride = 1; // stride in elements, not bytes + int f_stride = 1; + bool c_contig = true; + bool f_contig = true; + + for (int i = 0; i < d; ++i) { + const int c_idx = i; + const int f_idx = d - 1 - i; + if (b.raw_buffer()->dim[c_idx].stride != c_stride) { + c_contig = false; + } + c_stride *= b.raw_buffer()->dim[c_idx].extent; + + if (b.raw_buffer()->dim[f_idx].stride != f_stride) { + f_contig = false; + } + f_stride *= b.raw_buffer()->dim[f_idx].extent; + } + + return {c_contig, f_contig}; +} + } // namespace void define_buffer(py::module &m) { @@ -352,32 +394,40 @@ void define_buffer(py::module &m) { // Note that this allows us to convert a Buffer<> to any buffer-like object in Python; // most notably, we can convert to an ndarray by calling numpy.array() - .def_buffer([](Buffer<> &b) -> py::buffer_info { - return to_buffer_info(b, /*reverse_axes*/ true); + + // ELEPHANT: this always reverses axes, which might be surprising? + // We need to update the docs though. + // how about reverse axes only when "C", and does not when "F", otherweise, fail? + .def_buffer([](Buffer<> &self) -> py::buffer_info { + return to_buffer_info(self, /*reverse_axes*/ true); }) - // This allows us to use any buffer-like python entity to create a Buffer<> + .def("numpy_view", [](Buffer<> &self) -> py::array { + const auto [c_contig, f_contig] = is_any_contiguous(self); + if (!c_contig && !f_contig) { + throw py::value_error("Buffer is not contiguous in either C or F order; cannot create numpy view."); + } + const bool reverse_axes = c_contig && !f_contig; + // base = py::cast(self) ensures that `self` outlives the returned value. + return py::array(to_buffer_info(self, reverse_axes), /*base*/ py::cast(self)); }, "Returns a NumPy array that is a view of this Buffer. If the Buffer is C-contiguous (innermost first), reverses the axes to produce a C-contiguous array (innermost last). If the Buffer is F-contiguous (innermost last), does not reverse the axes, producing an F-contiguous array. If the Buffer is not contiguous in either order, raises an error.") + + .def("numpy_view", [](Buffer<> &self, bool reverse_axes) -> py::array { + // base = py::cast(self) ensures that `self` outlives the returned value. + return py::array(to_buffer_info(self, reverse_axes), /*base*/ py::cast(self)); }, py::arg("reverse_axes"), "Returns a NumPy array that is a view of this Buffer. The caller decides whether to reverse axis ordering.") + + // This allows us to use any buffer-like Python entity to create a Buffer<> // (most notably, an ndarray) .def(py::init_alias(), py::arg("buffer"), py::arg("name") = "", py::arg("reverse_axes") = true) .def(py::init_alias<>()) .def(py::init_alias &>()) - .def(py::init([](Type type, const std::vector &sizes, const std::string &name) -> Buffer<> { - return Buffer<>(type, sizes, name); - }), - py::arg("type"), py::arg("sizes"), py::arg("name") = "") + .def(py::init([](Type type, const std::vector &sizes, const std::string &name) -> Buffer<> { return Buffer<>(type, sizes, name); }), py::arg("type"), py::arg("sizes"), py::arg("name") = "") - .def(py::init([](Type type, const std::vector &sizes, const std::vector &storage_order, const std::string &name) -> Buffer<> { - return Buffer<>(type, sizes, storage_order, name); - }), - py::arg("type"), py::arg("sizes"), py::arg("storage_order"), py::arg("name") = "") + // The default storage order is [0, 1, 2, ...], meaning store the first axis densest. + .def(py::init([](Type type, const std::vector &sizes, const std::vector &storage_order, const std::string &name) -> Buffer<> { return Buffer<>(type, sizes, storage_order, name); }), py::arg("type"), py::arg("sizes"), py::arg("storage_order"), py::arg("name") = "") // Note that this exists solely to allow you to create a Buffer with a null host ptr; // this is necessary for some bounds-query operations (e.g. Func::infer_input_bounds). - .def_static( - "make_bounds_query", [](Type type, const std::vector &sizes, const std::string &name) -> Buffer<> { - return Buffer<>(type, nullptr, sizes, name); - }, - py::arg("type"), py::arg("sizes"), py::arg("name") = "") + .def_static("make_bounds_query", [](Type type, const std::vector &sizes, const std::string &name) -> Buffer<> { return Buffer<>(type, nullptr, sizes, name); }, py::arg("type"), py::arg("sizes"), py::arg("name") = "") .def_static("make_scalar", static_cast (*)(Type, const std::string &)>(Buffer<>::make_scalar), py::arg("type"), py::arg("name") = "") .def_static("make_interleaved", static_cast (*)(Type, int, int, int, const std::string &)>(Buffer<>::make_interleaved), py::arg("type"), py::arg("width"), py::arg("height"), py::arg("channels"), py::arg("name") = "") @@ -656,7 +706,7 @@ void define_buffer(py::module &m) { .def("__repr__", [](const Buffer<> &b) -> std::string { std::ostringstream o; if (b.defined()) { - o << ""; + o << ""; } else { o << ""; } diff --git a/python_bindings/test/correctness/buffer.py b/python_bindings/test/correctness/buffer.py index 73e00b62234d..cedfcafad52e 100644 --- a/python_bindings/test/correctness/buffer.py +++ b/python_bindings/test/correctness/buffer.py @@ -4,7 +4,7 @@ import sys -def test_ndarray_to_buffer(reverse_axes=True): +def test_ndarray_to_buffer(reverse_axes): a0 = np.ones((200, 300), dtype=np.int32) # Buffer always shares data (when possible) by default, @@ -49,7 +49,7 @@ def test_ndarray_to_buffer(reverse_axes=True): assert a0[56, 34] == 12 -def test_buffer_to_ndarray(reverse_axes=True): +def test_buffer_to_ndarray(reverse_axes): buf0 = hl.Buffer(hl.Int(16), [4, 6]) assert buf0.type() == hl.Int(16) buf0.fill(0) @@ -93,8 +93,9 @@ def test_buffer_to_ndarray(reverse_axes=True): assert array_shared[1, 2] == 3 assert array_copied[1, 2] == 42 - # Ensure that Buffers that have nonzero mins get converted correctly, - # since the Python Buffer Protocol doesn't have the 'min' concept + # Ensure that Buffers that have nonzero mins get converted correctly. + # Since the Python Buffer Protocol doesn't have the 'min' concept, the + # numpy views will have the cropped extents but no min coordinate. cropped_buf0 = buf0.copy() cropped_buf0.crop(dimension=0, min=1, extent=2) cropped_buf = cropped_buf0.reverse_axes() if not reverse_axes else cropped_buf0 @@ -131,6 +132,131 @@ def test_buffer_to_ndarray(reverse_axes=True): assert cropped_array_copied[0, 2] == 3 +def test_numpy_view_auto_reverse_axes(): + W = 16 + H = 12 + C = 3 + + x = 6 + y = 3 + c = 1 + + # Construct a hl.Buffer with the default storage order ([0, 1, 2]). + halide_buffer_default_order = hl.Buffer(hl.Int(16), [W, H, C]) + assert halide_buffer_default_order.type() == hl.Int(16) + halide_buffer_default_order.fill(0) + halide_buffer_default_order[x, y, c] = 42 + assert halide_buffer_default_order[x, y, c] == 42 + + # Its numpy_view() has: + # - axes reversed + # - The C_CONTIGUOUS flag set. + numpy_view_of_halide_buffer_default_order = halide_buffer_default_order.numpy_view() + assert numpy_view_of_halide_buffer_default_order.shape == (C, H, W) + assert numpy_view_of_halide_buffer_default_order.strides == (384, 32, 2) # (192, 16, 2) * sizeof(int16) + assert numpy_view_of_halide_buffer_default_order.dtype == np.int16 + assert numpy_view_of_halide_buffer_default_order.flags['C_CONTIGUOUS'] + assert numpy_view_of_halide_buffer_default_order[c, y, x] == 42 + + # Modifying the buffer should affect the view. + assert halide_buffer_default_order[x + 1, y + 1, c + 1] == 0 + halide_buffer_default_order[x + 1, y + 1, c + 1] = 99 + assert numpy_view_of_halide_buffer_default_order[c + 1, y + 1, x + 1] == 99 + + # Modifying the view should affect the buffer. + assert numpy_view_of_halide_buffer_default_order[c - 1, y - 1, x - 1] == 0 + numpy_view_of_halide_buffer_default_order[c - 1, y - 1, x - 1] = 100 + assert halide_buffer_default_order[x - 1, y - 1, c - 1] == 100 + + # Construct a hl.Buffer with the reversed storage order ([0, 1, 2]). + # Its shape (from Halide's perspective) is still (w, h, c), but its strides + # are reversed (c is densest). + halide_buffer_reversed_storage_order = hl.Buffer(hl.Int(16), [W, H, C], storage_order=[2, 1, 0]) + assert halide_buffer_reversed_storage_order.type() == hl.Int(16) + halide_buffer_reversed_storage_order.fill(0) + halide_buffer_reversed_storage_order[x, y, c] = 42 + assert halide_buffer_reversed_storage_order[x, y, c] == 42 + + # Its numpy_view() has: + # - axes preserved (i.e., its shape is also (w, h, c)). + # - The C_CONTIGUOUS flag set. This is not obvious: one would initially think that + # it might be F_CONTIGUOUS instead. But no, numpy considers an array to be + # C_CONTIGUOUS if the last axis is the densest (stride = itemsize). This is true + # because we asked hl.Buffer to allocate its last axis (c) to be the densest + # and did not reverse axes. + numpy_view_of_halide_buffer_reversed_storage_order = halide_buffer_reversed_storage_order.numpy_view() + assert numpy_view_of_halide_buffer_reversed_storage_order.shape == (W, H, C) + assert numpy_view_of_halide_buffer_reversed_storage_order.strides == (72, 6, 2) # (36, 3, 1) * sizeof(int16) + assert numpy_view_of_halide_buffer_reversed_storage_order.dtype == np.int16 + assert numpy_view_of_halide_buffer_reversed_storage_order.flags['C_CONTIGUOUS'] # This is not obvious. + assert numpy_view_of_halide_buffer_reversed_storage_order[x, y, c] == 42 + + # Modifying the buffer should affect the view. + assert halide_buffer_reversed_storage_order[x + 1, y + 1, c + 1] == 0 + halide_buffer_reversed_storage_order[x + 1, y + 1, c + 1] = 99 + assert numpy_view_of_halide_buffer_reversed_storage_order[x + 1, y + 1, c + 1] == 99 + + # Modifying the view should affect the buffer. + assert numpy_view_of_halide_buffer_reversed_storage_order[x - 1, y - 1, c - 1] == 0 + numpy_view_of_halide_buffer_reversed_storage_order[x - 1, y - 1, c - 1] = 100 + assert halide_buffer_reversed_storage_order[x - 1, y - 1, c - 1] == 100 + + # Construct a hl.Buffer with neither the first nor last as densest. + buf_neither_order = hl.Buffer(hl.Int(16), [W, H, C], storage_order=[1, 0, 2]) + assert buf_neither_order.type() == hl.Int(16) + buf_neither_order.fill(0) + buf_neither_order[x, y, c] = 42 + assert buf_neither_order[x, y, c] == 42 + + # numpy_view() without an explicit reverse_axes flag will fail. + try: + _ = buf_neither_order.numpy_view() + except ValueError as e: + assert "Buffer is not contiguous in either C or F order; cannot create numpy view" in str(e) + else: + assert False, "Did not see expected exception!" + +def test_numpy_view_explicit_reverse_axes(): + W = 16 + H = 12 + C = 3 + + x = 6 + y = 3 + c = 1 + + # Construct a hl.Buffer with neither the first nor last as densest. + buf_neither_order = hl.Buffer(hl.Int(16), [W, H, C], storage_order=[1, 0, 2]) + assert buf_neither_order.type() == hl.Int(16) + buf_neither_order.fill(0) + buf_neither_order[x, y, c] = 42 + assert buf_neither_order[x, y, c] == 42 + + # numpy_view(reverse_axes=False) should succeed and preserve axis order. + numpy_view_no_reverse = buf_neither_order.numpy_view(reverse_axes=False) + assert numpy_view_no_reverse.shape == (W, H, C) + assert numpy_view_no_reverse[x, y, c] == 42 + assert numpy_view_no_reverse.strides == (24, 2, 384) # (H, 1, H * W) * sizeof(int16) + + # numpy_view(reverse_axes=True) should succeed and reverse axis order. + numpy_view_reverse = buf_neither_order.numpy_view(reverse_axes=True) + assert numpy_view_reverse.shape == (C, H, W) + assert numpy_view_reverse[c, y, x] == 42 + assert numpy_view_reverse.strides == (384, 2, 24) # (H * W, 1, H) * sizeof(int16) + + # Modifying the buffer should affect both views. + assert buf_neither_order[x + 1, y + 1, c + 1] == 0 + buf_neither_order[x + 1, y + 1, c + 1] = 99 + assert numpy_view_no_reverse[x + 1, y + 1, c + 1] == 99 + assert numpy_view_reverse[c + 1, y + 1, x + 1] == 99 + + # Modifying one view should affect the other view and the buffer. + assert numpy_view_no_reverse[x - 1, y - 1, c - 1] == 0 + numpy_view_no_reverse[x - 1, y - 1, c - 1] = 100 + assert numpy_view_reverse[c - 1, y - 1, x - 1] == 100 + assert buf_neither_order[x - 1, y - 1, c - 1] == 100 + + def _assert_fn(e): assert e @@ -169,7 +295,7 @@ def test_bufferinfo_sharing(): a0 = np.ones((20000, 30000), dtype=np.int32) b0 = hl.Buffer(a0) del a0 - for i in range(200): + for _ in range(200): b1 = hl.Buffer(b0) b0 = b1 b1 = None @@ -546,6 +672,8 @@ def make_orig_buf(): test_ndarray_to_buffer(reverse_axes=False) test_buffer_to_ndarray(reverse_axes=True) test_buffer_to_ndarray(reverse_axes=False) + test_numpy_view_auto_reverse_axes() + test_numpy_view_explicit_reverse_axes() test_for_each_element() test_fill_all_equal() test_bufferinfo_sharing()