Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 24 additions & 10 deletions python_bindings/src/halide/halide_/PyBuffer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
}
Expand Down Expand Up @@ -351,33 +351,47 @@ 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<> &>())

// 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") = "")
Expand Down
11 changes: 7 additions & 4 deletions python_bindings/src/halide/halide_/PyBuffer.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
43 changes: 9 additions & 34 deletions python_bindings/src/halide/halide_/PyCallable.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
Expand Down
4 changes: 2 additions & 2 deletions python_bindings/test/correctness/addconstant_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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(
Expand Down
110 changes: 95 additions & 15 deletions python_bindings/test/correctness/buffer.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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)
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
Loading
Loading