[Lang] [ir] Add short-circuit if-then-else operator (#5022)

* [Lang] [ir] Add a short-circuit if-then-else operator and use it to implement IfExp

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: re-xyr

python/taichi/lang/ast/ast_transformer.py

@@ -1108,17 +1108,7 @@ def build_IfExp(ctx, node):
                 node.ptr = build_stmt(ctx, node.orelse)
             return node.ptr
 
-        val = impl.expr_init(None)
-
-        impl.begin_frontend_if(ctx.ast_builder, node.test.ptr)
-        ctx.ast_builder.begin_frontend_if_true()
-        val._assign(node.body.ptr)
-        ctx.ast_builder.pop_scope()
-        ctx.ast_builder.begin_frontend_if_false()
-        val._assign(node.orelse.ptr)
-        ctx.ast_builder.pop_scope()
-
-        node.ptr = val
+        node.ptr = ti_ops.ifte(node.test.ptr, node.body.ptr, node.orelse.ptr)
         return node.ptr
 
     @staticmethod

python/taichi/lang/ops.py

@@ -1143,6 +1143,29 @@ def py_select(cond, x1, x2):
     return _ternary_operation(_ti_core.expr_select, py_select, cond, x1, x2)
 
 
+@ternary
+def ifte(cond, x1, x2):
+    """Evaluate and return `x1` if `cond` is true; otherwise evaluate and return `x2`. This operator guarantees
+    short-circuit semantics: exactly one of `x1` or `x2` will be evaluated.
+
+    Args:
+        cond (:mod:`~taichi.types.primitive_types`): \
+            The condition.
+        x1, x2 (:mod:`~taichi.types.primitive_types`): \
+            The outputs.
+
+    Returns:
+        `x1` if `cond` is true and `x2` otherwise.
+    """
+    # TODO: systematically resolve `-1 = True` problem by introducing u1:
+    cond = logical_not(logical_not(cond))
+
+    def py_ifte(cond, x1, x2):
+        return x1 if cond else x2
+
+    return _ternary_operation(_ti_core.expr_ifte, py_ifte, cond, x1, x2)
+
+
 @writeback_binary
 def atomic_add(x, y):
     """Atomically compute `x + y`, store the result in `x`,

taichi/ir/expression_ops.h

@@ -124,6 +124,7 @@ DEFINE_EXPRESSION_FUNC_BINARY(floordiv)
 DEFINE_EXPRESSION_FUNC_BINARY(bit_shr)
 
 DEFINE_EXPRESSION_FUNC_TERNARY(select)
+DEFINE_EXPRESSION_FUNC_TERNARY(ifte)
 
 }  // namespace lang
 }  // namespace taichi

taichi/ir/frontend_ir.cpp

@@ -263,6 +263,37 @@ void BinaryOpExpression::flatten(FlattenContext *ctx) {
   stmt = ctx->back_stmt();
 }
 
+void make_ifte(Expression::FlattenContext *ctx,
+               DataType ret_type,
+               Expr cond,
+               Expr true_val,
+               Expr false_val) {
+  auto result = ctx->push_back<AllocaStmt>(ret_type);
+  flatten_rvalue(cond, ctx);
+  auto if_stmt = ctx->push_back<IfStmt>(cond->stmt);
+
+  Expression::FlattenContext lctx;
+  lctx.current_block = ctx->current_block;
+  flatten_rvalue(true_val, &lctx);
+  lctx.push_back<LocalStoreStmt>(result, true_val->stmt);
+
+  Expression::FlattenContext rctx;
+  rctx.current_block = ctx->current_block;
+  flatten_rvalue(false_val, &rctx);
+  rctx.push_back<LocalStoreStmt>(result, false_val->stmt);
+
+  auto true_block = std::make_unique<Block>();
+  true_block->set_statements(std::move(lctx.stmts));
+  if_stmt->set_true_statements(std::move(true_block));
+
+  auto false_block = std::make_unique<Block>();
+  false_block->set_statements(std::move(rctx.stmts));
+  if_stmt->set_false_statements(std::move(false_block));
+
+  ctx->push_back<LocalLoadStmt>(LocalAddress(result, 0));
+  return;
+}
+
 void TernaryOpExpression::type_check(CompileConfig *) {
   TI_ASSERT_TYPE_CHECKED(op1);
   TI_ASSERT_TYPE_CHECKED(op2);
@@ -276,21 +307,27 @@ void TernaryOpExpression::type_check(CompileConfig *) {
                     ternary_type_name(type), op1->ret_type->to_string(),
                     op2->ret_type->to_string(), op3->ret_type->to_string()));
   };
-  if (!is_integral(op1_type) || !op2_type->is<PrimitiveType>() ||
-      !op3_type->is<PrimitiveType>())
+  if (op1_type != PrimitiveType::i32)
+    error();
+  if (!op2_type->is<PrimitiveType>() || !op3_type->is<PrimitiveType>())
     error();
   ret_type = promoted_type(op2_type, op3_type);
 }
 
 void TernaryOpExpression::flatten(FlattenContext *ctx) {
   // if (stmt)
   //  return;
-  flatten_rvalue(op1, ctx);
-  flatten_rvalue(op2, ctx);
-  flatten_rvalue(op3, ctx);
-  ctx->push_back(
-      std::make_unique<TernaryOpStmt>(type, op1->stmt, op2->stmt, op3->stmt));
+  if (type == TernaryOpType::select) {
+    flatten_rvalue(op1, ctx);
+    flatten_rvalue(op2, ctx);
+    flatten_rvalue(op3, ctx);
+    ctx->push_back(
+        std::make_unique<TernaryOpStmt>(type, op1->stmt, op2->stmt, op3->stmt));
+  } else if (type == TernaryOpType::ifte) {
+    make_ifte(ctx, ret_type, op1, op2, op3);
+  }
   stmt = ctx->back_stmt();
+  stmt->tb = tb;
 }
 
 void InternalFuncCallExpression::type_check(CompileConfig *) {

taichi/ir/stmt_op_types.cpp

@@ -77,6 +77,7 @@ std::string ternary_type_name(TernaryOpType type) {
     return #i;
 
     REGISTER_TYPE(select);
+    REGISTER_TYPE(ifte);
 
 #undef REGISTER_TYPE
     default:

taichi/ir/stmt_op_types.h

@@ -62,7 +62,7 @@ inline bool is_bit_op(BinaryOpType type) {
 
 std::string binary_op_type_symbol(BinaryOpType type);
 
-enum class TernaryOpType : int { select, undefined };
+enum class TernaryOpType : int { select, ifte, undefined };
 
 std::string ternary_type_name(TernaryOpType type);
 

taichi/python/export_lang.cpp

@@ -763,6 +763,7 @@ void export_lang(py::module &m) {
   DEFINE_EXPRESSION_OP(log)
 
   DEFINE_EXPRESSION_OP(select)
+  DEFINE_EXPRESSION_OP(ifte)
 
   DEFINE_EXPRESSION_OP(cmp_le)
   DEFINE_EXPRESSION_OP(cmp_lt)

tests/cpp/ir/frontend_type_inference_test.cpp

@@ -66,16 +66,16 @@ TEST(FrontendTypeInference, UnaryOp) {
 }
 
 TEST(FrontendTypeInference, TernaryOp) {
-  auto const_i16 = value<int16>(-(1 << 10));
-  const_i16->type_check(nullptr);
-  EXPECT_EQ(const_i16->ret_type, PrimitiveType::i16);
-  auto cast_i8 = cast(const_i16, PrimitiveType::i8);
+  auto const_i32 = value<int32>(-(1 << 10));
+  const_i32->type_check(nullptr);
+  EXPECT_EQ(const_i32->ret_type, PrimitiveType::i32);
+  auto cast_i8 = cast(const_i32, PrimitiveType::i8);
   cast_i8->type_check(nullptr);
   EXPECT_EQ(cast_i8->ret_type, PrimitiveType::i8);
   auto const_f32 = value<float32>(5.0);
   const_f32->type_check(nullptr);
   EXPECT_EQ(const_f32->ret_type, PrimitiveType::f32);
-  auto ternary_f32 = expr_select(const_i16, cast_i8, const_f32);
+  auto ternary_f32 = expr_select(const_i32, cast_i8, const_f32);
   ternary_f32->type_check(nullptr);
   EXPECT_EQ(ternary_f32->ret_type, PrimitiveType::f32);
 }

tests/python/test_type_check.py

@@ -41,7 +41,7 @@ def select():
         d = b if a else c
 
     with pytest.raises(TypeError,
-                       match="`if` conditions must be of type int32"):
+                       match="unsupported operand type\\(s\\) for 'ifte'"):
         select()