Fix the implementation of the select() binding (#8862)

The new implementation corrects the select() grammar and gives
better error messages when nested expressions (inside tuples)
fail to cast / convert to the appropriate type.

Also add human-friendly __str__ overloads to Expr and Tuple

Author: alexreinking

python_bindings/src/halide/halide_/PyExpr.cpp

@@ -57,6 +57,11 @@ void define_expr(py::module &m) {
                     o << "<undefined halide.Expr>";
                 }
                 return o.str();
+            })
+            .def("__str__", [](const Expr &e) -> std::string {
+                std::ostringstream o;
+                o << e;
+                return o.str();
             });
 
     add_binary_operators(expr_class);

python_bindings/src/halide/halide_/PyIROperator.cpp

@@ -18,6 +18,59 @@ namespace Halide {
 
 namespace PythonBindings {
 
+namespace {
+
+bool is_expr(const py::handle &obj) {
+    try {
+        (void)obj.cast<Expr>();  // Check if casting succeeds
+        return true;
+    } catch (const py::cast_error &) {
+        return false;
+    }
+}
+
+template<typename T>
+T cast_arg(const py::handle &arg) {
+    try {
+        return arg.cast<T>();
+    } catch (const py::cast_error &) {
+        _halide_user_error
+            << "select(): Expected " << py::str(py::type::of<T>().attr("__name__"))
+            << " but got " << py::str(arg.get_type().attr("__name__")) << ": "
+            << py::str(arg);
+    }
+}
+
+template<typename TCond, typename TVal>
+py::object py_select_reduce(const py::args &args) {
+    auto false_case = cast_arg<TVal>(args[args.size() - 1]);
+    for (size_t pos = args.size() - 1; pos >= 2; pos -= 2) {
+        auto true_case = cast_arg<TVal>(args[pos - 1]);
+        auto condition = cast_arg<TCond>(args[pos - 2]);
+        false_case = select(condition, true_case, false_case);
+    }
+    return py::cast(false_case);
+}
+
+py::object py_select(const py::args &args) {
+    if (args.size() < 3) {
+        throw py::value_error("select() must have at least 3 arguments");
+    }
+
+    if (args.size() % 2 != 1) {
+        throw py::value_error("select() must have an odd number of arguments");
+    }
+
+    if (is_expr(args[0])) {        // If the condition is an Expr, then ...
+        return is_expr(args[1]) ?  // ... we need to check the value's kind.
+                   py_select_reduce<Expr, Expr>(args) :
+                   py_select_reduce<Expr, Tuple>(args);
+    }
+    return py_select_reduce<Tuple, Tuple>(args);  // Otherwise, the value must be a tuple, too.
+}
+
+}  // namespace
+
 void define_operators(py::module &m) {
     m.def("max", [](const py::args &args) -> Expr {
         if (args.size() < 2) {
@@ -48,85 +101,7 @@ void define_operators(py::module &m) {
     m.def("abs", &abs);
     m.def("absd", &absd);
 
-    m.def("select", [](const py::args &args) -> py::object {
-        if (args.size() < 3) {
-            throw py::value_error("select() must have at least 3 arguments");
-        }
-        if ((args.size() % 2) == 0) {
-            throw py::value_error("select() must have an odd number of arguments");
-        }
-
-        // Tricky set of options here:
-        //
-        // - (Expr, Expr, Expr, [Expr, Expr...]) -> Expr
-        // - (Expr, Tuple, Tuple, [Tuple, Tuple...]) -> Tuple   [Tuples must be same arity]
-        // - (Tuple, Tuple, Tuple, [Tuple, Tuple...]) -> Tuple  [Tuples must be same arity]
-        //
-        // It's made trickier by the fact that it's hard to do a reliable "is-a" check for Tuple here,
-        // so we'll do the slow-but-reliable approach of just trying to cast to Tuple and catching
-        // exceptions.
-
-        std::string tuple_error_msg;
-        try {
-            int pos = (int)args.size() - 1;
-            Tuple false_value = args[pos--].cast<Tuple>();
-            bool has_tuple_cond = false;
-            bool has_expr_cond = false;
-            while (pos > 0) {
-                Tuple true_value = args[pos--].cast<Tuple>();
-                // Note that 'condition' can be either Expr or Tuple, but must be consistent across all
-                py::object py_cond = args[pos--];
-                Expr expr_cond;
-                Tuple tuple_cond(expr_cond);
-                try {
-                    tuple_cond = py_cond.cast<Tuple>();
-                    has_tuple_cond = true;
-                } catch (...) {
-                    expr_cond = py_cond.cast<Expr>();
-                    has_expr_cond = true;
-                }
-
-                if (has_tuple_cond && has_expr_cond) {
-                    // We don't want to throw an error here, since the catch(...) would catch it,
-                    // and it would be hard to distinguish from other errors. Just set the string here
-                    // and jump to the error reporter outside the catch.
-                    tuple_error_msg = "select() on Tuples may not mix Expr and Tuple for the condition elements.";
-                    goto handle_tuple_error;
-                }
-
-                if (expr_cond.defined()) {
-                    false_value = select(expr_cond, true_value, false_value);
-                } else {
-                    if (tuple_cond.size() != true_value.size() || true_value.size() != false_value.size()) {
-                        // We don't want to throw an error here, since the catch(...) would catch it,
-                        // and it would be hard to distinguish from other errors. Just set the string here
-                        // and jump to the error reporter outside the catch.
-                        tuple_error_msg = "select() on Tuples requires all Tuples to have identical sizes.";
-                        goto handle_tuple_error;
-                    }
-                    false_value = select(tuple_cond, true_value, false_value);
-                }
-            }
-            return to_python_tuple(false_value);
-
-        } catch (...) {
-            // fall thru and try the Expr case
-        }
-
-    handle_tuple_error:
-        if (!tuple_error_msg.empty()) {
-            _halide_user_assert(false) << tuple_error_msg;
-        }
-
-        int pos = (int)args.size() - 1;
-        Expr false_expr_value = args[pos--].cast<Expr>();
-        while (pos > 0) {
-            Expr true_expr_value = args[pos--].cast<Expr>();
-            Expr condition_expr = args[pos--].cast<Expr>();
-            false_expr_value = select(condition_expr, true_expr_value, false_expr_value);
-        }
-        return py::cast(false_expr_value);
-    });
+    m.def("select", py_select);
 
     m.def("mux", static_cast<Expr (*)(const Expr &, const std::vector<Expr> &)>(&mux));
     m.def("mux", static_cast<Expr (*)(const Expr &, const Tuple &)>(&mux));

python_bindings/src/halide/halide_/PyTuple.cpp

@@ -66,6 +66,11 @@ void define_tuple(py::module &m) {
                 std::ostringstream o;
                 o << "<halide.Tuple of size " << t.size() << ">";
                 return o.str();
+            })
+            .def("__str__", [](const Tuple &t) -> std::string {
+                std::ostringstream o;
+                o << t;
+                return o.str();
             });
 
     py::implicitly_convertible<py::tuple, Tuple>();

python_bindings/test/correctness/iroperator.py

@@ -1,7 +1,9 @@
 import contextlib
-import halide as hl
-import sys
 import io
+import math
+import sys
+
+import halide as hl
 
 
 # redirect_stdout() requires Python3, alas
@@ -59,6 +61,30 @@ def test_select():
     assert b[3] == 3
 
 
+def test_select_bad_argmax():
+    x = hl.Var()
+    f = hl.Func()
+    f[x] = hl.sin(hl.f32(math.pi) * x / 16.0)
+
+    r = hl.RDom([(0, 10)])
+    g = hl.Func()
+
+    g[()] = (0, f.type().min())
+    try:
+        g[()] = hl.select(f[r] > g[()][1], (f[r], r), g[()])
+    except hl.HalideError as e:
+        assert (
+            "Error: The second and third arguments to a select do not have a matching type:"
+            in str(e)
+        )
+
+    g[()] = hl.select(f[r] > g[()][1], (r, f[r]), g[()])
+
+    idx, val = g.realize([])
+    assert idx[()] == 8
+    assert val[()] == 1.0
+
+
 def test_mux():
     c = hl.Var()
     f = hl.Func()
@@ -105,5 +131,6 @@ def test_minmax():
     test_print_expr()
     test_print_when()
     test_select()
+    test_select_bad_argmax()
     test_mux()
     test_minmax()

python_bindings/test/correctness/memoize.py

@@ -13,7 +13,7 @@ def test_memoize():
     output[x] = f[x]
 
     result = output.realize([3])
-    assert list(result) == [1., 1., 1.]
+    assert list(result) == [1.0, 1.0, 1.0]
 
 
 def main():

python_bindings/test/correctness/tuple_select.py

@@ -70,10 +70,7 @@ def test_tuple_select():
                                               (x-100, y-200))
         # fmt: on
     except hl.HalideError as e:
-        assert (
-            "select() on Tuples may not mix Expr and Tuple for the condition elements."
-            in str(e)
-        )
+        assert "select(): Expected Tuple but got Expr: ((x + y) < 100)" in str(e)
     else:
         assert False, "Did not see expected exception!"
 

src/IRPrinter.cpp

@@ -55,6 +55,7 @@ ostream &operator<<(ostream &out, const Type &type) {
     }
     return out;
 }
+
 ostream &operator<<(ostream &stream, const Expr &ir) {
     if (!ir.defined()) {
         stream << "(undefined)";
@@ -65,6 +66,14 @@ ostream &operator<<(ostream &stream, const Expr &ir) {
     return stream;
 }
 
+ostream &operator<<(ostream &stream, const Tuple &ir) {
+    stream << "(";
+    for (size_t i = 0; i < ir.size(); i++) {
+        stream << ir[i] << ", ";  // keep the trailing comma
+    }
+    return stream << ")";
+}
+
 ostream &operator<<(ostream &stream, const Buffer<> &buffer) {
     bool include_data = Internal::ends_with(buffer.name(), "_gpu_source_kernels");
     stream << "buffer " << buffer.name() << " = {";

src/IRPrinter.h

@@ -26,6 +26,10 @@ namespace Halide {
  * human-readable form */
 std::ostream &operator<<(std::ostream &stream, const Expr &);
 
+/** Emit a tuple on an output stream (such as std::cout) in
+ * human-readable form */
+std::ostream &operator<<(std::ostream &stream, const Tuple &);
+
 /** Emit a halide type on an output stream (such as std::cout) in
  * human-readable form */
 std::ostream &operator<<(std::ostream &stream, const Type &);