Add Tuple overloads for the non-arithmetic IROperator helpers

python_bindings/halide/src/halide_/PyIROperator.cpp

@@ -69,6 +69,29 @@ py::object py_select(const py::args &args) {
     return py_select_reduce<Tuple, Tuple>(args);  // Otherwise, the value must be a tuple, too.
 }
 
+// Shared by print and print_when. If the value being printed is a
+// Tuple, all of its elements are printed and the print is attached to
+// the first element.
+template<typename F>
+py::object py_print(const py::args &args, F make_print) {
+    if (args.empty()) {
+        throw py::value_error("print() must have at least 1 argument");
+    }
+    if (is_expr(args[0])) {
+        return py::cast(make_print(collect_print_args(args)));
+    }
+    Tuple t = args[0].cast<Tuple>();
+    py::tuple rest(args.size() - 1);
+    for (size_t i = 1; i < args.size(); i++) {
+        rest[i - 1] = args[i];
+    }
+    std::vector<Expr> v = t.as_vector();
+    std::vector<Expr> extra = collect_print_args(py::args(rest));
+    v.insert(v.end(), extra.begin(), extra.end());
+    t[0] = make_print(v);
+    return py::cast(t);
+}
+
 }  // namespace
 
 void define_operators(py::module &m) {
@@ -145,23 +168,35 @@ void define_operators(py::module &m) {
     m.def("reinterpret", static_cast<Expr (*)(Type, Expr)>(&reinterpret));
     m.def("cast", static_cast<Expr (*)(Type, Expr)>(&cast));
 
-    m.def("print", [](const py::args &args) -> Expr {
-        return print(collect_print_args(args));
+    m.def("print", [](const py::args &args) -> py::object {
+        return py_print(args, [](const std::vector<Expr> &v) { return print(v); });
     });
 
     m.def(
-        "print_when", [](const Expr &condition, const py::args &args) -> Expr {
-            return print_when(condition, collect_print_args(args));
+        "print_when", [](const Expr &condition, const py::args &args) -> py::object {
+            return py_print(args, [&](const std::vector<Expr> &v) { return print_when(condition, v); });
         },
         py::arg("condition"));
 
     m.def(
         "require", [](const Expr &condition, const Expr &value, const py::args &args) -> Expr {
-            auto v = args_to_vector<Expr>(args);
+            auto v = collect_print_args(args);
             v.insert(v.begin(), value);
             return require(condition, v);
         },
         py::arg("condition"), py::arg("value"));
+    m.def(
+        "require", [](const Expr &condition, const Tuple &value, const py::args &args) -> Tuple {
+            auto v = collect_print_args(args);
+            v.insert(v.begin(), Expr());
+            Tuple result = value;
+            for (Expr &e : result) {
+                v[0] = e;
+                e = require(condition, v);
+            }
+            return result;
+        },
+        py::arg("condition"), py::arg("value"));
 
     m.def("lerp", &lerp);
     m.def("popcount", &popcount);
@@ -182,11 +217,24 @@ void define_operators(py::module &m) {
             return Internal::memoize_tag_helper(result, args_to_vector<Expr>(cache_key_values));
         },
         py::arg("result"));
+    m.def(
+        "memoize_tag", [](const Tuple &result, const py::args &cache_key_values) -> Tuple {
+            auto v = args_to_vector<Expr>(cache_key_values);
+            Tuple tagged = result;
+            for (Expr &e : tagged) {
+                e = Internal::memoize_tag_helper(e, v);
+            }
+            return tagged;
+        },
+        py::arg("result"));
 
-    m.def("likely", &likely);
-    m.def("likely_if_innermost", &likely_if_innermost);
+    m.def("likely", static_cast<Expr (*)(Expr)>(&likely));
+    m.def("likely", static_cast<Tuple (*)(const Tuple &)>(&likely));
+    m.def("likely_if_innermost", static_cast<Expr (*)(Expr)>(&likely_if_innermost));
+    m.def("likely_if_innermost", static_cast<Tuple (*)(const Tuple &)>(&likely_if_innermost));
     m.def("saturating_cast", static_cast<Expr (*)(Type, Expr)>(&saturating_cast));
-    m.def("strict_float", &strict_float);
+    m.def("strict_float", static_cast<Expr (*)(const Expr &)>(&strict_float));
+    m.def("strict_float", static_cast<Tuple (*)(const Tuple &)>(&strict_float));
     m.def("scatter", static_cast<Expr (*)(const std::vector<Expr> &)>(&scatter));
     m.def("gather", static_cast<Expr (*)(const std::vector<Expr> &)>(&gather));
     m.def("extract_bits", static_cast<Expr (*)(Type, const Expr &, const Expr &)>(&extract_bits));

python_bindings/halide/src/halide_/PyTuple.cpp

@@ -28,6 +28,14 @@ Tuple to_halide_tuple(const py::object &o) {
     throw py::value_error("Expected an Expr or tuple-of-Expr.");
 }
 
+namespace {
+
+py::iterator tuple_iterator(const Tuple &t) {
+    return py::make_iterator(t.begin(), t.end());
+}
+
+}  // namespace
+
 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
@@ -62,6 +70,13 @@ void define_tuple(py::module &m) {
             .def(py::init([](const std::vector<Expr> &v) -> Tuple {
                 return Tuple(v);
             }))
+            .def("__len__", &Tuple::size)
+            .def("__getitem__", [](const Tuple &t, size_t i) -> Expr {
+                if (i >= t.size()) {
+                    throw py::index_error();
+                }
+                return t[i];
+            })
             .def("__repr__", [](const Tuple &t) -> std::string {
                 std::ostringstream o;
                 o << "<halide.Tuple of size " << t.size() << ">";
@@ -71,7 +86,8 @@ void define_tuple(py::module &m) {
                 std::ostringstream o;
                 o << t;
                 return o.str();
-            });
+            })
+            .def("__iter__", &tuple_iterator, py::keep_alive<0, 1>());
 
     py::implicitly_convertible<py::tuple, Tuple>();
 

python_bindings/halide/test/correctness/iroperator.py

@@ -127,9 +127,85 @@ def test_minmax():
     assert b[4] == 3
 
 
+def test_tuple_helpers():
+    x = hl.Var("x")
+    p = hl.Param(hl.Int(32), "p", 1)
+
+    # Helpers that don't imply any math distribute across Tuples.
+    f = hl.Func("f")
+    f[x] = hl.select(x < 10, (0, 0), hl.likely((x, x + 1)))
+    a, b = f.realize([20])
+    for i in range(20):
+        assert a[i] == (0 if i < 10 else i)
+        assert b[i] == (0 if i < 10 else i + 1)
+
+    # The result is an hl.Tuple, which can be indexed and iterated.
+    t = hl.likely((x, x + 1))
+    assert len(t) == 2
+    assert len(list(t)) == 2
+    f = hl.Func("f")
+    f[x] = t[0] + t[1]
+    a = f.realize([20])
+    for i in range(20):
+        assert a[i] == 2 * i + 1
+
+    f = hl.Func("f")
+    f[x] = hl.likely_if_innermost((x, x + 1))
+    a, b = f.realize([20])
+    for i in range(20):
+        assert a[i] == i
+        assert b[i] == i + 1
+
+    f = hl.Func("f")
+    f[x] = hl.strict_float((hl.f32(x) + 1.0, hl.f32(x) * 2.0))
+    a, b = f.realize([20])
+    for i in range(20):
+        assert a[i] == i + 1
+        assert b[i] == i * 2
+
+    f = hl.Func("f")
+    f[x] = hl.memoize_tag((x, x + 1), p)
+    a, b = f.realize([20])
+    for i in range(20):
+        assert a[i] == i
+        assert b[i] == i + 1
+
+    f = hl.Func("f")
+    f[x] = hl.require(p > 0, (x, x + 1), "p was", p)
+    a, b = f.realize([20])
+    for i in range(20):
+        assert a[i] == i
+        assert b[i] == i + 1
+
+    f = hl.Func("f")
+    f[x] = hl.print((x, x * 2), "at", x)
+    output = io.StringIO()
+    with _redirect_stdout(output):
+        a, b = f.realize([3])
+        expected = "0 0 at 0\n1 2 at 1\n2 4 at 2\n"
+        actual = output.getvalue()
+        assert expected == actual, f"Expected: {expected}, Actual: {actual}"
+    for i in range(3):
+        assert a[i] == i
+        assert b[i] == i * 2
+
+    f = hl.Func("f")
+    f[x] = hl.print_when(x == 1, (x, x * 2), "at", x)
+    output = io.StringIO()
+    with _redirect_stdout(output):
+        a, b = f.realize([3])
+        expected = "1 2 at 1\n"
+        actual = output.getvalue()
+        assert expected == actual, f"Expected: {expected}, Actual: {actual}"
+    for i in range(3):
+        assert a[i] == i
+        assert b[i] == i * 2
+
+
 if __name__ == "__main__":
     test_print_expr()
     test_print_when()
+    test_tuple_helpers()
     test_select()
     test_select_bad_argmax()
     test_mux()

src/Debug.h

@@ -13,9 +13,11 @@
 namespace Halide {
 
 struct Expr;
+class Tuple;
 struct Type;
 // Forward declare some things from IRPrinter, which we can't include yet.
 std::ostream &operator<<(std::ostream &stream, const Expr &);
+std::ostream &operator<<(std::ostream &stream, const Tuple &);
 std::ostream &operator<<(std::ostream &stream, const Type &);
 
 class Module;

src/Derivative.cpp

@@ -319,8 +319,8 @@ void ReverseAccumulationVisitor::propagate_adjoints(
             vector<Expr> zeros;
             Tuple rhs_tuple = func.values();
             zeros.reserve(rhs_tuple.size());
-            for (int i = 0; i < (int)rhs_tuple.size(); i++) {
-                zeros.push_back(make_zero(rhs_tuple[i].type()));
+            for (const Expr &e : rhs_tuple) {
+                zeros.push_back(make_zero(e.type()));
             }
             self_reference_adjoint = Tuple(zeros);
             self_reference_args.clear();
@@ -393,9 +393,8 @@ void ReverseAccumulationVisitor::propagate_adjoints(
                     // If the pure definition depends on any functions or buffers,
                     // there is no hope since we will overwrite something
                     Tuple rhs_tuple = func.values();
-                    for (int tuple_id = 0; tuple_id < (int)rhs_tuple.size();
-                         tuple_id++) {
-                        if (is_calling_function(rhs_tuple[tuple_id], let_var_mapping)) {
+                    for (const Expr &e : rhs_tuple) {
+                        if (is_calling_function(e, let_var_mapping)) {
                             error();
                         }
                     }
@@ -422,9 +421,8 @@ void ReverseAccumulationVisitor::propagate_adjoints(
 
             // Checking 2. here:
             bool all_zero_or_one_self_adjoint = true;
-            for (int i = 0; i < (int)self_reference_adjoint.size(); i++) {
-                if (!is_const(self_reference_adjoint[i], 0) &&
-                    !is_const(self_reference_adjoint[i], 1)) {
+            for (const Expr &e : self_reference_adjoint) {
+                if (!is_const(e, 0) && !is_const(e, 1)) {
                     all_zero_or_one_self_adjoint = false;
                     break;
                 }
@@ -449,9 +447,8 @@ void ReverseAccumulationVisitor::propagate_adjoints(
                     }
                 }
                 if (!r.defined()) {
-                    for (int tuple_id = 0; tuple_id < (int)update_tuple.size();
-                         tuple_id++) {
-                        r = extract_rdom(update_tuple[tuple_id]);
+                    for (const Expr &e : update_tuple) {
+                        r = extract_rdom(e);
                         if (r.defined()) {
                             break;
                         }

src/IROperator.cpp

@@ -1732,8 +1732,8 @@ Tuple select(const Expr &condition, const Tuple &true_value, const Tuple &false_
     return result;
 }
 
-Expr select(const Expr &condition, const FuncRef &true_value, const FuncRef &false_value) {
-    return select(condition, (Expr)true_value, (Expr)false_value);
+Tuple select(const Expr &condition, const FuncRef &true_value, const FuncRef &false_value) {
+    return select(condition, Tuple(true_value), Tuple(false_value));
 }
 
 Expr mux(const Expr &id, const std::vector<Expr> &values) {
@@ -2920,10 +2920,46 @@ Expr likely_if_innermost(Expr e) {
                       {std::move(e)}, Call::PureIntrinsic);
 }
 
+Tuple likely(const Tuple &t) {
+    Tuple result = t;
+    for (Expr &e : result) {
+        e = likely(e);
+    }
+    return result;
+}
+
+Tuple likely(const FuncRef &f) {
+    return likely(Tuple(f));
+}
+
+Tuple likely_if_innermost(const Tuple &t) {
+    Tuple result = t;
+    for (Expr &e : result) {
+        e = likely_if_innermost(e);
+    }
+    return result;
+}
+
+Tuple likely_if_innermost(const FuncRef &f) {
+    return likely_if_innermost(Tuple(f));
+}
+
 Expr strict_float(const Expr &e) {
     return strictify_float(e);
 }
 
+Tuple strict_float(const Tuple &t) {
+    Tuple result = t;
+    for (Expr &e : result) {
+        e = strict_float(e);
+    }
+    return result;
+}
+
+Tuple strict_float(const FuncRef &f) {
+    return strict_float(Tuple(f));
+}
+
 Expr undef(Type t) {
     return Call::make(t, Call::undef,
                       std::vector<Expr>(),