Revision e038765d17109d6f1e4226e88809a123aa3a4a07 authored by Steven Johnson on 18 April 2022, 18:20:51 UTC, committed by Steven Johnson on 18 April 2022, 18:20:51 UTC
Technically, we should check the return code of sprintf() for failure conditions. This is a minimal fix designed to quietly ensure that encoding errors aren't overlooked. (Since our "C" output actually requires C++11 at this point, I was tempted to replace it with `ostringstream`, but decided to avoid bringing that in...)
1 parent 60a909f
PyBuffer.cpp
#include "PyBuffer.h"
#include <utility>
#include "PyFunc.h"
#include "PyType.h"
namespace Halide {
namespace PythonBindings {
namespace {
// Standard stream output for halide_dimension_t
std::ostream &operator<<(std::ostream &stream, const halide_dimension_t &d) {
stream << "[" << d.min << "," << d.extent << "," << d.stride << "]";
return stream;
}
// Standard stream output for vector<halide_dimension_t>
std::ostream &operator<<(std::ostream &stream, const std::vector<halide_dimension_t> &shape) {
stream << "[";
bool need_comma = false;
for (const auto &d : shape) {
if (need_comma) {
stream << ",";
}
stream << d;
need_comma = true;
}
stream << "]";
return stream;
}
// Given a Buffer<>, return its shape in the form of a vector<halide_dimension_t>.
// (Oddly, Buffer<> has no API to do this directly.)
std::vector<halide_dimension_t> get_buffer_shape(const Buffer<> &b) {
if (!b.defined()) {
// Return an empty vector if the buffer is not defined.
return {};
}
std::vector<halide_dimension_t> s;
for (int i = 0; i < b.dimensions(); ++i) {
s.push_back(b.raw_buffer()->dim[i]);
}
return s;
}
// PyBind11 doesn't have baked-in support for float16, so add just
// enough special-case wrapping to support it.
template<typename T>
inline T value_cast(const py::object &value) {
return value.cast<T>();
}
template<>
inline float16_t value_cast<float16_t>(const py::object &value) {
return float16_t(value.cast<double>());
}
template<typename T>
inline std::string format_descriptor() {
return py::format_descriptor<T>::format();
}
template<>
inline std::string format_descriptor<float16_t>() {
return "e";
}
void call_fill(Buffer<> &b, const py::object &value) {
#define HANDLE_BUFFER_TYPE(TYPE) \
if (b.type() == type_of<TYPE>()) { \
b.as<TYPE>().fill(value_cast<TYPE>(value)); \
return; \
}
HANDLE_BUFFER_TYPE(bool)
HANDLE_BUFFER_TYPE(uint8_t)
HANDLE_BUFFER_TYPE(uint16_t)
HANDLE_BUFFER_TYPE(uint32_t)
HANDLE_BUFFER_TYPE(uint64_t)
HANDLE_BUFFER_TYPE(int8_t)
HANDLE_BUFFER_TYPE(int16_t)
HANDLE_BUFFER_TYPE(int32_t)
HANDLE_BUFFER_TYPE(int64_t)
HANDLE_BUFFER_TYPE(float16_t)
HANDLE_BUFFER_TYPE(float)
HANDLE_BUFFER_TYPE(double)
#undef HANDLE_BUFFER_TYPE
throw py::value_error("Unsupported Buffer<> type.");
}
bool call_all_equal(Buffer<> &b, const py::object &value) {
#define HANDLE_BUFFER_TYPE(TYPE) \
if (b.type() == type_of<TYPE>()) { \
return b.as<TYPE>().all_equal(value_cast<TYPE>(value)); \
}
HANDLE_BUFFER_TYPE(bool)
HANDLE_BUFFER_TYPE(uint8_t)
HANDLE_BUFFER_TYPE(uint16_t)
HANDLE_BUFFER_TYPE(uint32_t)
HANDLE_BUFFER_TYPE(uint64_t)
HANDLE_BUFFER_TYPE(int8_t)
HANDLE_BUFFER_TYPE(int16_t)
HANDLE_BUFFER_TYPE(int32_t)
HANDLE_BUFFER_TYPE(int64_t)
HANDLE_BUFFER_TYPE(float16_t)
HANDLE_BUFFER_TYPE(float)
HANDLE_BUFFER_TYPE(double)
#undef HANDLE_BUFFER_TYPE
throw py::value_error("Unsupported Buffer<> type.");
return false;
}
std::string type_to_format_descriptor(const Type &type) {
#define HANDLE_BUFFER_TYPE(TYPE) \
if (type == type_of<TYPE>()) return format_descriptor<TYPE>();
HANDLE_BUFFER_TYPE(bool)
HANDLE_BUFFER_TYPE(uint8_t)
HANDLE_BUFFER_TYPE(uint16_t)
HANDLE_BUFFER_TYPE(uint32_t)
HANDLE_BUFFER_TYPE(uint64_t)
HANDLE_BUFFER_TYPE(int8_t)
HANDLE_BUFFER_TYPE(int16_t)
HANDLE_BUFFER_TYPE(int32_t)
HANDLE_BUFFER_TYPE(int64_t)
HANDLE_BUFFER_TYPE(float16_t)
HANDLE_BUFFER_TYPE(float)
HANDLE_BUFFER_TYPE(double)
#undef HANDLE_BUFFER_TYPE
throw py::value_error("Unsupported Buffer<> type.");
return std::string();
}
Type format_descriptor_to_type(const std::string &fd) {
#define HANDLE_BUFFER_TYPE(TYPE) \
if (fd == format_descriptor<TYPE>()) return type_of<TYPE>();
HANDLE_BUFFER_TYPE(bool)
HANDLE_BUFFER_TYPE(uint8_t)
HANDLE_BUFFER_TYPE(uint16_t)
HANDLE_BUFFER_TYPE(uint32_t)
HANDLE_BUFFER_TYPE(uint64_t)
HANDLE_BUFFER_TYPE(int8_t)
HANDLE_BUFFER_TYPE(int16_t)
HANDLE_BUFFER_TYPE(int32_t)
HANDLE_BUFFER_TYPE(int64_t)
HANDLE_BUFFER_TYPE(float16_t)
HANDLE_BUFFER_TYPE(float)
HANDLE_BUFFER_TYPE(double)
#undef HANDLE_BUFFER_TYPE
// The string 'l' corresponds to np.int_, which is essentially
// a C 'long'; return a 32 or 64 bit int as appropriate.
if (fd == "l") {
return sizeof(long) == 8 ? type_of<int64_t>() : type_of<int32_t>();
}
throw py::value_error("Unsupported Buffer<> type.");
return Type();
}
py::object buffer_getitem_operator(Buffer<> &buf, const std::vector<int> &pos) {
if ((size_t)pos.size() != (size_t)buf.dimensions()) {
throw py::value_error("Incorrect number of dimensions.");
}
// TODO: add bounds checking?
#define HANDLE_BUFFER_TYPE(TYPE) \
if (buf.type() == type_of<TYPE>()) \
return py::cast(buf.as<TYPE>()(pos.data()));
HANDLE_BUFFER_TYPE(bool)
HANDLE_BUFFER_TYPE(uint8_t)
HANDLE_BUFFER_TYPE(uint16_t)
HANDLE_BUFFER_TYPE(uint32_t)
HANDLE_BUFFER_TYPE(uint64_t)
HANDLE_BUFFER_TYPE(int8_t)
HANDLE_BUFFER_TYPE(int16_t)
HANDLE_BUFFER_TYPE(int32_t)
HANDLE_BUFFER_TYPE(int64_t)
HANDLE_BUFFER_TYPE(float16_t)
HANDLE_BUFFER_TYPE(float)
HANDLE_BUFFER_TYPE(double)
#undef HANDLE_BUFFER_TYPE
throw py::value_error("Unsupported Buffer<> type.");
return py::object();
}
py::object buffer_setitem_operator(Buffer<> &buf, const std::vector<int> &pos, const py::object &value) {
if ((size_t)pos.size() != (size_t)buf.dimensions()) {
throw py::value_error("Incorrect number of dimensions.");
}
// TODO: add bounds checking?
#define HANDLE_BUFFER_TYPE(TYPE) \
if (buf.type() == type_of<TYPE>()) \
return py::cast(buf.as<TYPE>()(pos.data()) = value_cast<TYPE>(value));
HANDLE_BUFFER_TYPE(bool)
HANDLE_BUFFER_TYPE(uint8_t)
HANDLE_BUFFER_TYPE(uint16_t)
HANDLE_BUFFER_TYPE(uint32_t)
HANDLE_BUFFER_TYPE(uint64_t)
HANDLE_BUFFER_TYPE(int8_t)
HANDLE_BUFFER_TYPE(int16_t)
HANDLE_BUFFER_TYPE(int32_t)
HANDLE_BUFFER_TYPE(int64_t)
HANDLE_BUFFER_TYPE(float16_t)
HANDLE_BUFFER_TYPE(float)
HANDLE_BUFFER_TYPE(double)
#undef HANDLE_BUFFER_TYPE
throw py::value_error("Unsupported Buffer<> type.");
return py::object();
}
// Use an alias class so that if we are created via a py::buffer, we can
// keep the py::buffer_info class alive for the life of the Buffer<>,
// ensuring the data isn't collected out from under us.
class PyBuffer : public Buffer<> {
py::buffer_info info;
static std::vector<halide_dimension_t> make_dim_vec(const py::buffer_info &info) {
const Type t = format_descriptor_to_type(info.format);
std::vector<halide_dimension_t> dims;
dims.reserve(info.ndim);
for (int i = 0; i < info.ndim; i++) {
if (INT_MAX < info.shape[i] || INT_MAX < (info.strides[i] / t.bytes())) {
throw py::value_error("Out of range arguments to make_dim_vec.");
}
dims.emplace_back(0, (int32_t)info.shape[i], (int32_t)(info.strides[i] / t.bytes()));
}
return dims;
}
PyBuffer(py::buffer_info &&info, const std::string &name)
: Buffer<>(
format_descriptor_to_type(info.format),
info.ptr,
(int)info.ndim,
make_dim_vec(info).data(),
name),
info(std::move(info)) {
}
public:
PyBuffer()
: Buffer<>(), info() {
}
explicit PyBuffer(const Buffer<> &b)
: Buffer<>(b), info() {
}
PyBuffer(const py::buffer &buffer, const std::string &name)
: PyBuffer(buffer.request(/*writable*/ true), name) {
// Default to setting host-dirty on any PyBuffer we create from an existing py::buffer;
// this allows (e.g.) code like
//
// input = hl.Buffer(imageio.imread(image_path))
//
// to work as you'd expected when the buffer is used for input. (Without
// host_dirty being set, copying to GPU can be skipped and produce surprising
// results.)
//
// Note that this is a bit different from the C++ Buffer<> ctors that
// take pointer-to-existing data (which do *not* set host dirty); a crucial
// difference there is that those also do not take ownership, but here,
// we always take (shared) ownership of the underlying buffer.
this->set_host_dirty();
}
~PyBuffer() override = default;
};
} // namespace
void define_buffer(py::module &m) {
using BufferDimension = Halide::Runtime::Buffer<>::Dimension;
auto buffer_dimension_class =
py::class_<BufferDimension>(m, "BufferDimension")
.def("min", &BufferDimension::min)
.def("stride", &BufferDimension::stride)
.def("extent", &BufferDimension::extent)
.def("max", &BufferDimension::max);
auto buffer_class =
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()
.def_buffer([](Buffer<> &b) -> py::buffer_info {
if (b.data() == nullptr) {
throw py::value_error("Cannot convert a Buffer<> with null host ptr to a Python buffer.");
}
const int d = b.dimensions();
const int bytes = b.type().bytes();
std::vector<Py_ssize_t> shape, strides;
for (int i = 0; i < d; i++) {
shape.push_back((Py_ssize_t)b.raw_buffer()->dim[i].extent);
strides.push_back((Py_ssize_t)(b.raw_buffer()->dim[i].stride * bytes));
}
return py::buffer_info(
b.data(), // Pointer to buffer
bytes, // Size of one scalar
type_to_format_descriptor(b.type()), // Python struct-style format descriptor
d, // Number of dimensions
shape, // Buffer dimensions
strides // Strides (in bytes) for each index
);
})
// 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 &>(), py::arg("buffer"), py::arg("name") = "")
.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::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_scalar", (Buffer<>(*)(Type, const std::string &))Buffer<>::make_scalar, py::arg("type"), py::arg("name") = "")
.def_static("make_interleaved", (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") = "")
.def_static(
"make_with_shape_of", [](const Buffer<> &buffer, const std::string &name) -> Buffer<> {
return Buffer<>::make_with_shape_of(buffer, nullptr, nullptr, name);
},
py::arg("src"), py::arg("name") = "")
.def("set_name", &Buffer<>::set_name)
.def("name", &Buffer<>::name)
.def("same_as", (bool(Buffer<>::*)(const Buffer<> &other) const) & Buffer<>::same_as, py::arg("other"))
.def("defined", &Buffer<>::defined)
.def("type", &Buffer<>::type)
.def("channels", (int(Buffer<>::*)() const) & Buffer<>::channels)
.def("dimensions", (int(Buffer<>::*)() const) & Buffer<>::dimensions)
.def("width", (int(Buffer<>::*)() const) & Buffer<>::width)
.def("height", (int(Buffer<>::*)() const) & Buffer<>::height)
.def("top", (int(Buffer<>::*)() const) & Buffer<>::top)
.def("bottom", (int(Buffer<>::*)() const) & Buffer<>::bottom)
.def("left", (int(Buffer<>::*)() const) & Buffer<>::left)
.def("right", (int(Buffer<>::*)() const) & Buffer<>::right)
.def("number_of_elements", (size_t(Buffer<>::*)() const) & Buffer<>::number_of_elements)
.def("size_in_bytes", (size_t(Buffer<>::*)() const) & Buffer<>::size_in_bytes)
.def("has_device_allocation", (bool(Buffer<>::*)() const) & Buffer<>::has_device_allocation)
.def("host_dirty", (bool(Buffer<>::*)() const) & Buffer<>::host_dirty)
.def("device_dirty", (bool(Buffer<>::*)() const) & Buffer<>::device_dirty)
.def(
"set_host_dirty", [](Buffer<> &b, bool dirty) -> void {
b.set_host_dirty(dirty);
},
py::arg("dirty") = true)
.def(
"set_device_dirty", [](Buffer<> &b, bool dirty) -> void {
b.set_device_dirty(dirty);
},
py::arg("dirty") = true)
.def("copy", &Buffer<>::copy)
.def("copy_from", &Buffer<>::copy_from<void, Buffer<>::AnyDims>)
.def("add_dimension", (void(Buffer<>::*)()) & Buffer<>::add_dimension)
.def("allocate", [](Buffer<> &b) -> void {
b.allocate(nullptr, nullptr);
})
.def("deallocate", (void(Buffer<>::*)()) & Buffer<>::deallocate)
.def("device_deallocate", (void(Buffer<>::*)()) & Buffer<>::device_deallocate)
.def(
"crop", [](Buffer<> &b, int d, int min, int extent) -> void {
b.crop(d, min, extent);
},
py::arg("dimension"), py::arg("min"), py::arg("extent"))
.def(
"crop", [](Buffer<> &b, const std::vector<std::pair<int, int>> &rect) -> void {
b.crop(rect);
},
py::arg("rect"))
// Present in Runtime::Buffer but not Buffer
// .def("cropped", [](Buffer<> &b, int d, int min, int extent) -> Buffer<> {
// return b.cropped(d, min, extent);
// }, py::arg("dimension"), py::arg("min"), py::arg("extent"))
// .def("cropped", [](Buffer<> &b, const std::vector<std::pair<int, int>> &rect) -> Buffer<> {
// return b.cropped(rect);
// }, py::arg("rect"))
.def(
"embed", [](Buffer<> &b, int d, int pos) -> void {
b.embed(d, pos);
},
py::arg("dimension"), py::arg("pos"))
.def(
"embedded", [](Buffer<> &b, int d, int pos) -> Buffer<> {
return b.embedded(d, pos);
},
py::arg("dimension"), py::arg("pos"))
.def(
"embed", [](Buffer<> &b, int d) -> void {
b.embed(d);
},
py::arg("dimension"))
.def(
"embedded", [](Buffer<> &b, int d) -> Buffer<> {
return b.embedded(d);
},
py::arg("dimension"))
.def(
"slice", [](Buffer<> &b, int d, int pos) -> void {
b.slice(d, pos);
},
py::arg("dimension"), py::arg("pos"))
.def(
"sliced", [](Buffer<> &b, int d, int pos) -> Buffer<> {
return b.sliced(d, pos);
},
py::arg("dimension"), py::arg("pos"))
.def(
"slice", [](Buffer<> &b, int d) -> void {
b.slice(d);
},
py::arg("dimension"))
.def(
"sliced", [](Buffer<> &b, int d) -> Buffer<> {
return b.sliced(d);
},
py::arg("dimension"))
.def(
"translate", [](Buffer<> &b, int d, int dx) -> void {
b.translate(d, dx);
},
py::arg("dimension"), py::arg("dx"))
.def(
"translate", [](Buffer<> &b, const std::vector<int> &delta) -> void {
b.translate(delta);
},
py::arg("delta"))
// Present in Runtime::Buffer but not Buffer
// .def("translated", [](Buffer<> &b, int d, int dx) -> Buffer<> {
// return b.translated(d, dx);
// }, py::arg("dimension"), py::arg("dx"))
// .def("translated", [](Buffer<> &b, const std::vector<int> &delta) -> Buffer<> {
// return b.translated(delta);
// }, py::arg("delta"))
.def(
"transpose", [](Buffer<> &b, int d1, int d2) -> void {
b.transpose(d1, d2);
},
py::arg("d1"), py::arg("d2"))
.def(
"transposed", [](Buffer<> &b, int d1, int d2) -> Buffer<> {
return b.transposed(d1, d2);
},
py::arg("d1"), py::arg("d2"))
.def(
"transpose", [](Buffer<> &b, const std::vector<int> &order) -> void {
b.transpose(order);
},
py::arg("order"))
.def(
"transposed", [](Buffer<> &b, const std::vector<int> &order) -> Buffer<> {
return b.transposed(order);
},
py::arg("order"))
.def(
"dim", [](Buffer<> &b, int dimension) -> BufferDimension {
return b.dim(dimension);
},
py::arg("dimension"), py::keep_alive<0, 1>() // Keep Buffer<> alive while Dimension exists
)
.def(
"for_each_element", [](Buffer<> &b, py::function f) -> void {
const int d = b.dimensions();
std::vector<int> pos_v(d, 0);
b.for_each_element([&f, &pos_v](const int *pos) -> void {
for (size_t i = 0; i < pos_v.size(); ++i) {
pos_v[i] = pos[i];
}
f(pos_v);
});
},
py::arg("f"))
.def("fill", &call_fill, py::arg("value"))
.def("all_equal", &call_all_equal, py::arg("value"))
// TODO: for_each_value() needs to be rethought a bit for Python;
// for C++ is passes a by-reference to each value (for mutation),
// but Python doesn't allow mutable references to intrinsic types
// like int. Leaving unimplemented for now.
// .def("for_each_value", [](Buffer<> &b, py::args f, py::args other_buffers) -> void {
// }, py::arg("f"))
.def("copy_to_host", [](Buffer<> &b) -> int {
return b.copy_to_host(nullptr);
})
.def("device_detach_native", [](Buffer<> &b) -> int {
return b.device_detach_native(nullptr);
})
.def("device_free", [](Buffer<> &b) -> int {
return b.device_free(nullptr);
})
.def("device_sync", [](Buffer<> &b) -> int {
return b.device_sync(nullptr);
})
.def(
"copy_to_device", [](Buffer<> &b, const Target &t) -> int {
return b.copy_to_device(t);
},
py::arg("target") = get_jit_target_from_environment())
.def(
"copy_to_device", [](Buffer<> &b, const DeviceAPI &d, const Target &t) -> int {
return b.copy_to_device(d, t);
},
py::arg("device_api"), py::arg("target") = get_jit_target_from_environment())
.def(
"device_malloc", [](Buffer<> &b, const Target &t) -> int {
return b.device_malloc(t);
},
py::arg("target") = get_jit_target_from_environment())
.def(
"device_malloc", [](Buffer<> &b, const DeviceAPI &d, const Target &t) -> int {
return b.device_malloc(d, t);
},
py::arg("device_api"), py::arg("target") = get_jit_target_from_environment())
.def(
"set_min", [](Buffer<> &b, const std::vector<int> &mins) -> void {
if (mins.size() > (size_t)b.dimensions()) {
throw py::value_error("Too many arguments");
}
b.set_min(mins);
},
py::arg("mins"))
.def(
"contains", [](Buffer<> &b, const std::vector<int> &coords) -> bool {
if (coords.size() > (size_t)b.dimensions()) {
throw py::value_error("Too many arguments");
}
return b.contains(coords);
},
py::arg("coords"))
.def("__getitem__", [](Buffer<> &buf, const int &pos) -> py::object {
return buffer_getitem_operator(buf, {pos});
})
.def("__getitem__", [](Buffer<> &buf, const std::vector<int> &pos) -> py::object {
return buffer_getitem_operator(buf, pos);
})
.def("__getitem__", [](Buffer<> &buf, const Expr &pos) -> Expr {
return buf(std::vector<Expr>{pos});
})
.def("__getitem__", [](Buffer<> &buf, const std::vector<Expr> &pos) -> Expr {
return buf(pos);
})
.def("__setitem__", [](Buffer<> &buf, const int &pos, const py::object &value) -> py::object {
return buffer_setitem_operator(buf, {pos}, value);
})
.def("__setitem__", [](Buffer<> &buf, const std::vector<int> &pos, const py::object &value) -> py::object {
return buffer_setitem_operator(buf, pos, value);
})
.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) << ">";
} else {
o << "<undefined halide.Buffer>";
}
return o.str();
});
}
} // namespace PythonBindings
} // namespace Halide
Computing file changes ...
