Get raw Runtime::Buffer from Buffer in Python rather than use PyBuffer (#8682)

The buffer protocol (PyBuffer) doesn't support all the same metadata that
Halide::Runtime::Buffer does. In particular, it does not support
coordinate offsets (as in set_min).

When an object is passed to an AOT Python pipeline, it is now checked
for a _get_raw_halide_buffer_t method. If it exists, it should
return a pointer to a halide_buffer_t whose lifetime is at least that
of the object.

Fixes #8652

Author: alexreinking

python_bindings/apps/CMakeLists.txt

@@ -31,16 +31,19 @@ set(TEST_IMAGES_DIR "$<SHELL_PATH:${CMAKE_CURRENT_SOURCE_DIR}/../../apps/images>
 set(APPS
     bilateral_grid
     blur
+    identity
     interpolate
     local_laplacian)
 
 set(GENERATORS_bilateral_grid   bilateral_grid bilateral_grid_Adams2019 bilateral_grid_Li2018 bilateral_grid_Mullapudi2016)
+set(GENERATORS_blur             blur)
+set(GENERATORS_identity         identity)
 set(GENERATORS_interpolate      interpolate interpolate_Mullapudi2016)
 set(GENERATORS_local_laplacian  local_laplacian local_laplacian_Mullapudi2016)
-set(GENERATORS_blur             blur)
 
 set(ARGS_bilateral_grid   ${TEST_IMAGES_DIR}/gray.png 0.1 ${TEST_TMPDIR}/out.png)
 set(ARGS_blur             ${TEST_IMAGES_DIR}/gray.png ${TEST_TMPDIR}/out.png)
+set(ARGS_identity         "")
 set(ARGS_interpolate      ${TEST_IMAGES_DIR}/rgba.png ${TEST_TMPDIR}/out.png)
 set(ARGS_local_laplacian  ${TEST_IMAGES_DIR}/rgba.png 8 1 1 ${TEST_TMPDIR}/out.png)
 

python_bindings/apps/identity_app.py

@@ -0,0 +1,38 @@
+import sys
+
+import halide as hl
+import numpy as np
+from identity import identity
+
+
+def main():
+    size = 100
+
+    ##
+    # First test using Numpy buffers
+
+    expected = np.arange(size, dtype=np.int32)
+    array = np.empty_like(expected)
+
+    identity(array)
+
+    if not np.array_equal(array, expected):
+        sys.exit("np.array failure!")
+
+    ##
+    # Second test using hl.Buffer as a wrapper
+
+    expected = np.arange(-size, 0, dtype=np.int32)
+    array = np.empty_like(expected)
+    array_hl = hl.Buffer(array)
+    array_hl.set_min([-size])
+    identity(array_hl)
+
+    if not np.array_equal(array, expected):
+        sys.exit("hl.Buffer failure!")
+
+    print("Success!")
+
+
+if __name__ == "__main__":
+    main()

python_bindings/apps/identity_generator.py

@@ -0,0 +1,14 @@
+import halide as hl
+
+
+@hl.generator(name="identity")
+class Identity:
+    output = hl.OutputBuffer(hl.Int(32), 1)
+
+    def generate(self):
+        x = hl.Var("x")
+        self.output[x] = x
+
+
+if __name__ == "__main__":
+    hl.main()

python_bindings/src/halide/halide_/PyBuffer.cpp

@@ -653,6 +653,10 @@ void define_buffer(py::module &m) {
                     o << "<undefined halide.Buffer>";
                 }
                 return o.str();  //
+            })
+
+            .def("_get_raw_halide_buffer_t", [](const Buffer<> &b) -> uintptr_t {
+                return reinterpret_cast<uintptr_t>(b.raw_buffer());  //
             });
 }
 

src/PythonExtensionGen.cpp

@@ -319,22 +319,57 @@ template<int dimensions>
 struct PyHalideBuffer {
     // Must allocate at least 1, even if d=0
     static constexpr int dims_to_allocate = (dimensions < 1) ? 1 : dimensions;
+    static constexpr const char* get_raw_halide_runtime_buffer_fn = "_get_raw_halide_buffer_t";
 
     Py_buffer py_buf;
-    halide_dimension_t halide_dim[dims_to_allocate];
-    halide_buffer_t halide_buf;
+    halide_buffer_t* halide_buf = nullptr;
     bool py_buf_needs_release = false;
     bool needs_device_free = false;
 
+    bool unpack_from_halide_buffer(PyObject *py_obj) {
+        if (!PyObject_HasAttrString(py_obj, get_raw_halide_runtime_buffer_fn)) {
+            return false;
+        }
+
+        PyObject *py_raw_buffer = PyObject_CallMethod(py_obj, get_raw_halide_runtime_buffer_fn, NULL);
+        if (!py_raw_buffer) {
+            PyErr_Clear();
+            return false;
+        }
+
+        if (!PyLong_Check(py_raw_buffer)) {
+            Py_DECREF(py_raw_buffer);
+            return false;
+        }
+
+        uintptr_t py_raw_buffer_ptr = (uintptr_t)PyLong_AsUnsignedLongLong(py_raw_buffer);
+        Py_DECREF(py_raw_buffer);
+
+        if (py_raw_buffer_ptr == 0) {
+            return false;
+        }
+
+        halide_buf = reinterpret_cast<halide_buffer_t *>(py_raw_buffer_ptr);
+        return true;
+    }
+
     bool unpack(PyObject *py_obj, int py_getbuffer_flags, const char *name) {
-        return Halide::PythonRuntime::unpack_buffer(py_obj, py_getbuffer_flags,
-            name, dimensions, py_buf, halide_dim, halide_buf, py_buf_needs_release,
-            needs_device_free);
+        if (unpack_from_halide_buffer(py_obj)) {
+            return true;
+        }
+        if (Halide::PythonRuntime::unpack_buffer(
+                py_obj, py_getbuffer_flags, name, dimensions, py_buf,
+                unpacked_dim, unpacked_buf, py_buf_needs_release,
+                needs_device_free)) {
+            halide_buf = &unpacked_buf;
+            return true;
+        }
+        return false;
     }
 
     ~PyHalideBuffer() {
         if (needs_device_free) {
-            halide_device_free(nullptr, &halide_buf);
+            halide_device_free(nullptr, halide_buf);
         }
         if (py_buf_needs_release) {
             PyBuffer_Release(&py_buf);
@@ -346,6 +381,10 @@ struct PyHalideBuffer {
     PyHalideBuffer &operator=(const PyHalideBuffer &other) = delete;
     PyHalideBuffer(PyHalideBuffer &&other) = delete;
     PyHalideBuffer &operator=(PyHalideBuffer &&other) = delete;
+
+private:
+    halide_dimension_t unpacked_dim[dims_to_allocate];
+    halide_buffer_t unpacked_buf;
 };
 
 }  // namespace
@@ -470,7 +509,7 @@ void PythonExtensionGen::compile(const LoweredFunc &f) {
     // do a lazy-copy-to-GPU if needed.
     for (size_t i = 0; i < args.size(); i++) {
         if (args[i].is_buffer() && args[i].is_input()) {
-            dest << indent << "b_" << arg_names[i] << ".halide_buf.set_host_dirty();\n";
+            dest << indent << "b_" << arg_names[i] << ".halide_buf->set_host_dirty();\n";
         }
     }
     dest << indent << "int result;\n";
@@ -479,7 +518,7 @@ void PythonExtensionGen::compile(const LoweredFunc &f) {
     indent.indent += 2;
     for (size_t i = 0; i < args.size(); i++) {
         if (args[i].is_buffer()) {
-            dest << indent << "&b_" << arg_names[i] << ".halide_buf";
+            dest << indent << "b_" << arg_names[i] << ".halide_buf";
         } else {
             dest << indent << "py_" << arg_names[i] << "";
         }
@@ -496,7 +535,7 @@ void PythonExtensionGen::compile(const LoweredFunc &f) {
     // random garbage. (We need a better solution for this, see https://github.com/halide/Halide/issues/6868)
     for (size_t i = 0; i < args.size(); i++) {
         if (args[i].is_buffer() && args[i].is_output()) {
-            dest << indent << "if (result == 0) result = halide_copy_to_host(nullptr, &b_" << arg_names[i] << ".halide_buf);\n";
+            dest << indent << "if (result == 0) result = halide_copy_to_host(nullptr, b_" << arg_names[i] << ".halide_buf);\n";
         }
     }
     dest << indent << "if (result != 0) {\n";