Python bindings: adds numpy_view() to hl.Buffer.

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<bool, bool> 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::buffer, const std::string &, bool>(), py::arg("buffer"), py::arg("name") = "", py::arg("reverse_axes") = true)
             .def(py::init_alias<>())
             .def(py::init_alias<const Buffer<> &>())
-            .def(py::init([](Type type, const std::vector<int> &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<int> &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<int> &sizes, const std::vector<int> &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<int> &sizes, const std::vector<int> &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<int> &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<int> &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<Buffer<> (*)(Type, const std::string &)>(Buffer<>::make_scalar), py::arg("type"), py::arg("name") = "")
             .def_static("make_interleaved", static_cast<Buffer<> (*)(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 << "<halide.Buffer of type " << halide_type_to_string(b.type()) << " shape:" << get_buffer_shape(b) << ">";
+                    o << "<halide.Buffer of type " << halide_type_to_string(b.type()) << " shape: " << get_buffer_shape(b) << ">";
                 } else {
                     o << "<undefined halide.Buffer>";
                 }

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()