Add evaluate() and evaluate_may_gpu() to Python bindings

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()
@@ -364,6 +387,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()