Add evaluate() and evaluate_may_gpu() to Python bindings (#7108)

* Add evaluate() and evaluate_may_gpu() to Python bindings

* pacify clang-tidy

python_bindings/src/halide/halide_/PyBuffer.cpp

@@ -195,8 +195,6 @@ Type format_descriptor_to_type(const std::string &fd) {
  return Type();
 }
 
-namespace {
-
 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.");
@@ -228,6 +226,8 @@ py::object buffer_getitem_operator(Buffer<> &buf, const std::vector<int> &pos) {
  return py::object();
 }
 
+namespace {
+
 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.");

python_bindings/src/halide/halide_/PyBuffer.h

@@ -10,6 +10,8 @@ void define_buffer(py::module &m);
 
 Type format_descriptor_to_type(const std::string &fd);
 
+py::object buffer_getitem_operator(Buffer<> &buf, const std::vector<int> &pos);
+
 template<typename T = void,
  int Dims = AnyDims,
  int InClassDimStorage = (Dims == AnyDims ? 4 : std::max(Dims, 1))>

python_bindings/src/halide/halide_/PyFunc.cpp

@@ -80,15 +80,38 @@ py::object realization_to_object(const Realization &r) {
  return to_python_tuple(r);
 }
 
+py::object evaluate_impl(const py::object &expr, bool may_gpu) {
+ Tuple t = to_halide_tuple(expr);
+ Func f("evaluate_func_" + std::to_string(t.size()));
+ f() = t;
+ if (may_gpu) {
+ Internal::schedule_scalar(f);
+ }
+
+ std::optional<Realization> r;
+ {
+ py::gil_scoped_release release;
+
+ r = f.realize();
+ }
+ if (r->size() == 1) {
+ return buffer_getitem_operator((*r)[0], {});
+ } else {
+ py::tuple result(r->size());
+ for (size_t i = 0; i < r->size(); i++) {
+ result[i] = buffer_getitem_operator((*r)[i], {});
+ }
+ return result;
+ }
+}
+
 } // namespace
 
 void define_func(py::module &m) {
  define_func_ref(m);
  define_var_or_rvar(m);
  define_loop_level(m);
 
- // TODO: add ParamMap support.
-
  // Deliberately not supported, because they don't seem to make sense for Python:
  // - set_custom_allocator()
  // - set_custom_do_task()
@@ -380,6 +403,18 @@ void define_func(py::module &m) {
  add_schedule_methods(func_class);
 
  define_stage(m);
+
+ m.def(
+ "evaluate", [](const py::object &expr) -> py::object {
+ return evaluate_impl(expr, false);
+ },
+ py::arg("expr"));
+
+ m.def(
+ "evaluate_may_gpu", [](const py::object &expr) -> py::object {
+ return evaluate_impl(expr, true);
+ },
+ py::arg("expr"));
 }
 
 } // namespace PythonBindings

python_bindings/src/halide/halide_/PyTuple.cpp

@@ -3,6 +3,31 @@
 namespace Halide {
 namespace PythonBindings {
 
+Tuple to_halide_tuple(const py::object &o) {
+ try {
+ Expr e = o.cast<Expr>();
+ return Tuple(e);
+ } catch (...) {
+ // fall thru
+ }
+
+ try {
+ py::tuple t = o.cast<py::tuple>();
+ if (t.empty()) {
+ throw py::value_error("Cannot use a zero-length tuple-of-Expr");
+ }
+ std::vector<Expr> v(t.size());
+ for (size_t i = 0; i < t.size(); i++) {
+ v[i] = t[i].cast<Expr>();
+ }
+ return Tuple(v);
+ } catch (...) {
+ // fall thru
+ }
+
+ throw py::value_error("Expected an Expr or tuple-of-Expr.");
+}
+
 void define_tuple(py::module &m) {
  // Halide::Tuple isn't surfaced to the user in Python;
  // we define it here to allow PyBind to do some automatic
@@ -42,6 +67,7 @@ void define_tuple(py::module &m) {
  o << "<halide.Tuple of size " << t.size() << ">";
  return o.str();
  });
+
  py::implicitly_convertible<py::tuple, Tuple>();
 
  // If we autoconvert from vector<Expr>, we must also special-case FuncRef, alas

python_bindings/src/halide/halide_/PyTuple.h

@@ -18,6 +18,11 @@ inline py::tuple to_python_tuple(const T &ht) {
  return pt;
 }
 
+// in: convertible-to-Expr, or tuple-of-convertible-to-Expr
+// out: Halide::Tuple
+// throws exception if not convertible
+Tuple to_halide_tuple(const py::object &o);
+
 } // namespace PythonBindings
 } // namespace Halide
 

python_bindings/test/correctness/division.py

@@ -1,24 +1,13 @@
 import halide as hl
 
-# TODO: Func.evaluate() needs a wrapper added;
-# this is a temporary equivalent for testing purposes
-def _evaluate(e):
- # TODO: support zero-dim Func, Buffers
- buf = hl.Buffer(type = e.type(), sizes = [1])
- f = hl.Func();
- x = hl.Var()
- f[x] = e;
- f.realize(buf)
- return buf[0]
-
-def test_division():
- f32 = hl.Param(hl.Float(32), 'f32', -32.0)
- f64 = hl.Param(hl.Float(64), 'f64', 64.0)
- i16 = hl.Param(hl.Int(16), 'i16', -16)
- i32 = hl.Param(hl.Int(32), 'i32', 32)
- u16 = hl.Param(hl.UInt(16), 'u16', 16)
- u32 = hl.Param(hl.UInt(32), 'u32', 32)
+f32 = hl.Param(hl.Float(32), 'f32', -32.0)
+f64 = hl.Param(hl.Float(64), 'f64', 64.0)
+i16 = hl.Param(hl.Int(16), 'i16', -16)
+i32 = hl.Param(hl.Int(32), 'i32', 32)
+u16 = hl.Param(hl.UInt(16), 'u16', 16)
+u32 = hl.Param(hl.UInt(32), 'u32', 32)
 
+def test_types():
  # Verify that the types match the rules in match_types()
  assert (f32 / f64).type() == hl.Float(64)
  assert (f32 // f64).type() == hl.Float(64)
@@ -41,28 +30,51 @@ def test_division():
  assert (i16 / f64).type() == hl.Float(64)
  assert (i16 // f64).type() == hl.Float(64)
 
+def test_division():
  # Verify that division semantics match those for Halide
  # (rather than python); this differs for int/int which
  # defaults to float (rather than floordiv) in Python3.
  # Also test that // always floors the result, even for float.
- assert _evaluate(f32 / f64) == -0.5
- assert _evaluate(f32 // f64) == -1.0
+ assert hl.evaluate(f32 / f64) == -0.5
+ assert hl.evaluate(f32 // f64) == -1.0
+
+ assert hl.evaluate(i16 / i32) == -1
+ assert hl.evaluate(i16 // i32) == -1
+ assert hl.evaluate(i32 / i16) == -2
+
+ assert hl.evaluate(u16 / u32) == 0
+ assert hl.evaluate(u16 // u32) == 0
 
- assert _evaluate(i16 / i32) == -1
- assert _evaluate(i16 // i32) == -1
- assert _evaluate(i32 / i16) == -2
+ assert hl.evaluate(u16 / i32) == 0
+ assert hl.evaluate(i32 // u16) == 2
 
- assert _evaluate(u16 / u32) == 0
- assert _evaluate(u16 // u32) == 0
+ assert hl.evaluate(u16 / f32) == -0.5
+ assert hl.evaluate(u16 // f32) == -1.0
 
- assert _evaluate(u16 / i32) == 0
- assert _evaluate(i32 // u16) == 2
+ assert hl.evaluate(i16 / f64) == -0.25
+ assert hl.evaluate(i16 // f64) == -1.0
 
- assert _evaluate(u16 / f32) == -0.5
- assert _evaluate(u16 // f32) == -1.0
+def test_division_tupled():
+ # Same as test_division, but using the tuple variant
+ assert hl.evaluate((f32 / f64, f32 // f64)) == (-0.5, -1.0)
+ assert hl.evaluate((i16 / i32, i16 // i32, i32 / i16)) == (-1, -1, -2)
+ assert hl.evaluate((u16 / u32, u16 // u32)) == (0, 0)
+ assert hl.evaluate((u16 / i32, i32 // u16)) == (0, 2)
+ assert hl.evaluate((u16 / f32, u16 // f32)) == (-0.5, -1.0)
+ assert hl.evaluate((i16 / f64, i16 // f64)) == (-0.25, -1.0)
 
- assert _evaluate(i16 / f64) == -0.25
- assert _evaluate(i16 // f64) == -1.0
+def test_division_gpu():
+ # Allow GPU usage -- don't use f64 since not all GPU backends support that
+ f = hl.cast(hl.Float(32), f64)
+ assert hl.evaluate_may_gpu((f32 / f, f32 // f)) == (-0.5, -1.0)
+ assert hl.evaluate_may_gpu((i16 / i32, i16 // i32, i32 / i16)) == (-1, -1, -2)
+ assert hl.evaluate_may_gpu((u16 / u32, u16 // u32)) == (0, 0)
+ assert hl.evaluate_may_gpu((u16 / i32, i32 // u16)) == (0, 2)
+ assert hl.evaluate_may_gpu((u16 / f32, u16 // f32)) == (-0.5, -1.0)
+ assert hl.evaluate_may_gpu((i16 / f, i16 // f)) == (-0.25, -1.0)
 
 if __name__ == "__main__":
+ test_types()
  test_division()
+ test_division_tupled()
+ test_division_gpu()