Makes hl.Buffer <--> Python Buffer Protocol axis convention consistent throughout (#8953)

Halide interacts with Python in three place:
- Python bindings
- PyCallablle (`compile_to_callable()` via Python bindings)
- Python extension

This PR changes our axis ordering conventions to be consistent throughout:
- Constructing a Python buffer object from an `hl.Buffer` reverses axes.
- Constructing an `hl.Buffer` from any Python buffer object reverses axes.
  - Note that an `hl.Buffer` is also a Python buffer object. E.g., you can use `my_buffer[idx0, idx1, idx2]` to retrieve an element. Constructing an `hl.Buffer` from another `hl.Buffer` does *not* reverse axes because that would be surprising.
  - `hl.Buffer(my_array, "my_buffer_name", reverse_axes=False)` is explicitly supported.
- Calling a PyCallable with:
  - A `hl.Buffer` does not reverse axes (it's considered a Halide Buffer, not a Python buffer object)
  - A Python buffer object does reverse axes.
- Calling a pipeline function in a Python extension generated by a generator has the same conventions as a PyCallable.

Because of the new consistent conventions, we now permit:
- Non-contiguous views. Previously, Python-side C-contiguous buffers were reversed and F-contiguous buffers were not. Now we support views with arbitrary strides (so long as the Python-side stride in bytes divides the element size, and if the result is in `[INT_MIN, INT_MAX]`).
- This includes negatives strides.

- `hl.Buffer` also gained a `.numpy_view()` method that returns a `np.array` with an explicit `reverse_axes` argument (defaulting to `True`).

Author: jiawen

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(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.");
     }
@@ -351,33 +351,52 @@ void define_buffer(py::module &m) {
         py::class_<Buffer<>, PyBuffer>(m, "Buffer", py::buffer_protocol())
 
             // 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()
+            // most notably, we can convert to an ndarray by calling numpy.array(other_buffer). This
+            // always reverses the axes.
             .def_buffer([](Buffer<> &b) -> py::buffer_info {
                 return to_buffer_info(b, /*reverse_axes*/ true);
             })
 
-            // 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)
+            // Returns a NumPy array that is a view of (shares storage with) this Buffer.
+            .def(
+                "numpy_view", [](Buffer<> &self, bool reverse_axes) -> py::array {
+                    // base = py::cast(self) ensures that `self` outlives the returned view.
+                    return py::array(to_buffer_info(self, reverse_axes), /*base*/ py::cast(self));
+                },
+                py::arg("reverse_axes") = true, "Returns a NumPy array that is a view of this Buffer. By default the axes are "
+                                                "reversed. Pass reverse_axes=False to preserve the original axis order.")
+
             .def(py::init_alias<>())
-            .def(py::init_alias<const Buffer<> &>())
+            // keep_alive</*Nurse*/ 1, /*Patient*/ 2>: this copy shares the source's data pointer
+            // without taking ownership (the source may itself wrap external memory, e.g. a numpy
+            // array kept alive only by the source Buffer).
+            //
+            // In keep_alive<1, 2>, 1 refers to `this`, 2 refers to the argument `Buffer<> &` arg.
+            .def(py::init_alias<const Buffer<> &>(), py::keep_alive<1, 2>())
+
+            // This py::init_alias allows us to use any buffer-like Python entity to create a
+            // Buffer<> (most notably, an ndarray). reverse_axes = True by default.
+            //
+            // ! Note that the copy/default constructors are registered *above, before* the
+            // py::buffer one before. This is because hl.Buffer also satisfies the buffer protocol,
+            // and we want hl.Buffer(other_buffer) to be a "direct" copy (no axis reversal) rather
+            // than routing through the axis-reversing buffer-protocol path.
+            .def(py::init_alias<py::buffer, const std::string &, bool>(), py::arg("buffer"), py::arg("name") = "", py::arg("reverse_axes") = true)
+
             .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") = "")
 
+            // 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") = "")

python_bindings/src/halide/halide_/PyBuffer.h

@@ -20,14 +20,17 @@ Halide::Runtime::Buffer<T, Dims, InClassDimStorage> pybufferinfo_to_halidebuffer
     halide_dimension_t *dims = (halide_dimension_t *)alloca(info.ndim * sizeof(halide_dimension_t));
     _halide_user_assert(dims);
     for (int i = 0; i < info.ndim; i++) {
-        if (INT_MAX < info.shape[i] || INT_MAX < (info.strides[i] / t.bytes())) {
+        // Strides may be negative in both numpy and Halide. So check against both INT_MIN and
+        // INT_MAX.
+        const auto elem_stride = info.strides[i] / t.bytes();
+        if (info.shape[i] > INT_MAX || elem_stride < INT_MIN || elem_stride > INT_MAX) {
             throw py::value_error("Out of range dimensions in buffer conversion.");
         }
-        // Halide's default indexing convention is col-major (the most rapidly varying index comes first);
+        // Reverse axes if requested.
         // Numpy's default is row-major (most rapidly varying comes last).
-        // We usually want to reverse the order so that most-varying comes first.
+
         const int dst_axis = reverse_axes ? (info.ndim - i - 1) : i;
-        dims[dst_axis] = {0, (int32_t)info.shape[i], (int32_t)(info.strides[i] / t.bytes())};
+        dims[dst_axis] = {0, (int32_t)info.shape[i], (int32_t)elem_stride};
     }
     return Halide::Runtime::Buffer<T, Dims, InClassDimStorage>(t, info.ptr, (int)info.ndim, dims);
 }

python_bindings/src/halide/halide_/PyCallable.cpp

@@ -47,28 +47,6 @@ T cast_to(const py::handle &h) {
     }
 }
 
-std::pair<bool, bool> is_any_contiguous(const py::buffer_info &info) {
-    py::ssize_t c_stride = info.itemsize;
-    py::ssize_t f_stride = info.itemsize;
-    bool c_contig = true;
-    bool f_contig = true;
-
-    for (size_t i = 0; i < info.ndim; ++i) {
-        size_t c_idx = info.ndim - 1 - i;
-        if (info.strides[c_idx] != c_stride) {
-            c_contig = false;
-        }
-        c_stride *= info.shape[c_idx];
-
-        if (info.strides[i] != f_stride) {
-            f_contig = false;
-        }
-        f_stride *= info.shape[i];
-    }
-
-    return {c_contig, f_contig};
-}
-
 }  // namespace
 
 class PyCallable {
@@ -105,28 +83,25 @@ class PyCallable {
 
         const auto define_one_arg = [&argv, &scalar_storage, &buffers, &cci](const Argument &c_arg, py::handle value, size_t slot) {
             if (c_arg.is_buffer()) {
-                // If the argument is already a Halide Buffer of some sort,
-                // skip pybuffer_to_halidebuffer entirely, since the latter requires
-                // a non-null host ptr, but we might want such a buffer for bounds inference,
-                // and we don't need the intermediate HalideBuffer wrapper anyway.
                 halide_buffer_t *raw_buffer;
                 if (py::isinstance<Halide::Buffer<>>(value)) {
+                    // If the argument is already a Halide Buffer of some sort,
+                    // skip pybuffer_to_halidebuffer entirely, since the latter requires
+                    // a non-null host ptr, but we might want such a buffer for bounds inference,
+                    // and we don't need the intermediate HalideBuffer wrapper anyway.
+                    //
+                    // Do not reverse axes.
                     auto b = cast_to<Halide::Buffer<>>(value);
                     raw_buffer = b.raw_buffer();
                 } else {
+                    // If it's a buffer-protocol object (e.g. a NumPy array), the convention
+                    // is to always reverse axes.
+                    constexpr bool reverse_axes = true;
                     py::buffer py_buffer_value = cast_to<py::buffer>(value);
                     const bool writable = c_arg.is_output();
 
                     const py::buffer_info value_buffer_info = py_buffer_value.request(writable);
-                    auto [c_contig, f_contig] = is_any_contiguous(value_buffer_info);
-
-                    if (!c_contig && !f_contig) {
-                        throw Halide::Error("Invalid buffer: only C or F contiguous buffers are supported");
-                    }
 
-                    // It is possible for a buffer to be both C and F contiguous
-                    // (e.g., a scalar or a 1D buffer).
-                    const bool reverse_axes = c_contig && !f_contig;
                     buffers.buffers[slot] =
                         pybufferinfo_to_halidebuffer<void, AnyDims, MaxFastDimensions>(
                             value_buffer_info, reverse_axes);

python_bindings/test/correctness/addconstant_test.py

@@ -42,7 +42,7 @@ def test(addconstant_impl_func, offset):
     input_float = numpy.array([3.14, 2.718, 1.618], dtype=numpy.float32)
     input_double = numpy.array([3.14, 2.718, 1.618], dtype=numpy.float64)
     input_half = numpy.array([3.14, 2.718, 1.618], dtype=numpy.float16)
-    input_2d = numpy.array([[1, 2, 3], [4, 5, 6]], dtype=numpy.int8, order="F")
+    input_2d = numpy.array([[1, 2, 3], [4, 5, 6]], dtype=numpy.int8)
     input_3d = numpy.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=numpy.int8)
 
     output_u8 = numpy.zeros((3,), dtype=numpy.uint8)
@@ -56,7 +56,7 @@ def test(addconstant_impl_func, offset):
     output_float = numpy.zeros((3,), dtype=numpy.float32)
     output_double = numpy.zeros((3,), dtype=numpy.float64)
     output_half = numpy.zeros((3,), dtype=numpy.float16)
-    output_2d = numpy.zeros((2, 3), dtype=numpy.int8, order="F")
+    output_2d = numpy.zeros((2, 3), dtype=numpy.int8)
     output_3d = numpy.zeros((2, 2, 2), dtype=numpy.int8)
 
     addconstant_impl_func(

python_bindings/test/correctness/buffer.py

@@ -1,15 +1,17 @@
-import halide as hl
-import numpy as np
 import gc
 import sys
 
+import halide as hl
+import numpy as np
+
 
-def test_ndarray_to_buffer(reverse_axes=True):
+# Tests constructing a hl.Buffer from an np.array.
+# `reverse_axes`: passed to hl.Buffer (default False).
+def test_ndarray_to_buffer(reverse_axes):
     a0 = np.ones((200, 300), dtype=np.int32)
 
     # Buffer always shares data (when possible) by default,
-    # and maintains the shape of the data source. (note that
-    # the ndarray is col-major by default!)
+    # and maintains the shape of the data source.
     b0 = hl.Buffer(a0, "float32_test_buffer", reverse_axes)
     assert b0.type() == hl.Int(32)
     assert b0.name() == "float32_test_buffer"
@@ -49,7 +51,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)
@@ -58,13 +60,12 @@ def test_buffer_to_ndarray(reverse_axes=True):
 
     # This is subtle: the default behavior when converting
     # a Buffer to an np.array (or ndarray, etc) is to reverse the
-    # order of the axes, since Halide prefers column-major and
-    # the rest of Python prefers row-major. By calling reverse_axes()
-    # before that conversion, we end up doing a *double* reverse, i.e,
-    # not reversing at all. So the 'not' here is correct.
+    # order of the axes. By calling reverse_axes() before that conversion,
+    # we end up doing a *double* reverse, i.e, not reversing at all.
+    # So the 'not' here is correct.
     buf = buf0.reverse_axes() if not reverse_axes else buf0
 
-    # Should share storage with buf
+    # Should share storage with buf0.
     array_shared = np.array(buf, copy=False)
     assert array_shared.dtype == np.int16
     if reverse_axes:
@@ -74,7 +75,7 @@ def test_buffer_to_ndarray(reverse_axes=True):
         assert array_shared.shape == (4, 6)
         assert array_shared[1, 2] == 42
 
-    # Should *not* share storage with buf
+    # copy=True -> should *not* share storage with buf.
     array_copied = np.array(buf, copy=True)
     assert array_copied.dtype == np.int16
     if reverse_axes:
@@ -84,7 +85,8 @@ def test_buffer_to_ndarray(reverse_axes=True):
         assert array_copied.shape == (4, 6)
         assert array_copied[1, 2] == 42
 
-    # Should affect array_shared but not array_copied
+    # Modifying one element of buf0 should affect array_shared but not
+    # array_copied.
     buf0[1, 2] = 3
     if reverse_axes:
         assert array_shared[2, 1] == 3
@@ -94,7 +96,8 @@ def test_buffer_to_ndarray(reverse_axes=True):
         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
+    # 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
@@ -169,7 +172,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
@@ -539,13 +542,90 @@ def make_orig_buf():
             assert realized[row, col] == tag * 2
 
 
+def test_numpy_view():
+    # numpy_view() always reverses axes by default, regardless of the Buffer's
+    # storage order; reverse_axes=False keeps Halide's axis order. It works for
+    # any layout, contiguous or not.
+
+    # Default storage order: dim 0 is densest (stride 1).
+    buf = hl.Buffer(hl.Int(16), [4, 6])
+    buf.fill(0)
+    buf[1, 2] = 42
+
+    v = buf.numpy_view()  # default reverse_axes=True
+    assert v.dtype == np.int16
+    assert v.shape == (6, 4)  # axes reversed
+    assert v[2, 1] == 42  # numpy[a0, a1] <-> halide[a1, a0]
+
+    v_no = buf.numpy_view(reverse_axes=False)
+    assert v_no.shape == (4, 6)  # Halide order preserved
+    assert v_no[1, 2] == 42
+
+    # Both are views that share storage with the Buffer.
+    v[0, 0] = 7
+    assert buf[0, 0] == 7
+    assert v_no[0, 0] == 7
+
+    # Allocate hl.Buffer with axis 1 densest instead.
+    # numpy_view() doesn't care about storage order, just reverses indexing order.
+    buf_f = hl.Buffer(hl.Int(16), [4, 6], [1, 0])  # store axis 1 densest
+    buf_f.fill(0)
+    buf_f[1, 2] = 24
+    vf = buf_f.numpy_view()
+    assert vf.shape == (6, 4)
+    assert vf[2, 1] == 24
+
+    # Allocate a hl.Buffer that is neither C- nor F-contiguous.
+    # numpy_view() still works but is a strided view.
+    buf_nc = hl.Buffer(hl.Int(16), [4, 6, 8], [1, 0, 2])
+    buf_nc.fill(0)
+    buf_nc[1, 2, 3] = 99
+    vnc = buf_nc.numpy_view()
+    assert vnc.shape == (8, 6, 4)
+    assert vnc[3, 2, 1] == 99
+
+
+def test_buffer_copy_no_reverse():
+    # Constructing an hl.Buffer from another hl.Buffer must be a direct copy
+    # (no axis reversal), even though hl.Buffer also satisfies the buffer protocol.
+    b1 = hl.Buffer(hl.Int(16), [4, 6, 8], [1, 0, 2])  # non-square, non-default order
+    b1.fill(0)
+    b1[1, 2, 3] = 17
+
+    b2 = hl.Buffer(b1)
+    assert b2.dimensions() == b1.dimensions()
+    for i in range(b1.dimensions()):
+        assert b2.dim(i).extent() == b1.dim(i).extent(), i
+        assert b2.dim(i).stride() == b1.dim(i).stride(), i
+    assert b2[1, 2, 3] == 17
+
+    # A bounds-query buffer (null host) can be copied without error.
+    bq = hl.Buffer.make_bounds_query(hl.Int(16), [4, 6])
+    bq_copy = hl.Buffer(bq)
+    assert bq_copy.dim(0).extent() == 4
+    assert bq_copy.dim(1).extent() == 6
+
+
+def test_ndarray_negative_strides():
+    # NumPy arrays with negative strides should work.
+    a = np.arange(6, dtype=np.int32)[::-1]  # [5, 4, 3, 2, 1, 0], stride -1
+    b = hl.Buffer(a)
+    assert b.dim(0).extent() == 6
+    assert b.dim(0).stride() == -1
+    assert b[0] == 5
+    assert b[5] == 0
+
+
 if __name__ == "__main__":
     test_make_interleaved()
     test_interleaved_ndarray()
     test_ndarray_to_buffer(reverse_axes=True)
     test_ndarray_to_buffer(reverse_axes=False)
     test_buffer_to_ndarray(reverse_axes=True)
     test_buffer_to_ndarray(reverse_axes=False)
+    test_numpy_view()
+    test_buffer_copy_no_reverse()
+    test_ndarray_negative_strides()
     test_for_each_element()
     test_fill_all_equal()
     test_bufferinfo_sharing()