[bug] Fixed type promotion rule for bit-shift operations

taichi/ir/frontend_ir.cpp

@@ -197,6 +197,11 @@ void BinaryOpExpression::type_check(CompileConfig *config) {
     ret_type = PrimitiveType::i32;
     return;
   }
+  if (is_shift_op(type)) {
+    ret_type = lhs_type;
+    return;
+  }
+
   if (type == BinaryOpType::truediv) {
     auto default_fp = config->default_fp;
     if (!is_real(lhs_type)) {

taichi/ir/stmt_op_types.h

@@ -47,6 +47,11 @@ inline bool binary_is_logical(BinaryOpType t) {
 
 std::string binary_op_type_name(BinaryOpType type);
 
+inline bool is_shift_op(BinaryOpType type) {
+  return type == BinaryOpType::bit_sar || type == BinaryOpType::bit_shl ||
+         type == BinaryOpType::bit_shr;
+}
+
 inline bool is_comparison(BinaryOpType type) {
   return starts_with(binary_op_type_name(type), "cmp");
 }

taichi/transforms/type_check.cpp

@@ -244,6 +244,29 @@ class TypeCheck : public IRVisitor {
     return stmt;
   }
 
+  void insert_shift_op_assertion_before(Stmt *stmt, Stmt *lhs, Stmt *rhs) {
+    int rhs_limit = data_type_bits(lhs->ret_type);
+    auto const_stmt =
+        Stmt::make<ConstStmt>(TypedConstant(rhs->ret_type, rhs_limit));
+    auto cond_stmt =
+        Stmt::make<BinaryOpStmt>(BinaryOpType::cmp_le, rhs, const_stmt.get());
+
+    std::string msg =
+        "Detected overflow for bit_shift_op with rhs = %d, exceeding limit of "
+        "%d.";
+    std::vector<Stmt *> args = {rhs, const_stmt.get()};
+    auto assert_stmt =
+        Stmt::make<AssertStmt>(cond_stmt.get(), msg, std::move(args));
+
+    const_stmt->accept(this);
+    cond_stmt->accept(this);
+    assert_stmt->accept(this);
+
+    stmt->insert_before_me(std::move(const_stmt));
+    stmt->insert_before_me(std::move(cond_stmt));
+    stmt->insert_before_me(std::move(assert_stmt));
+  }
+
   void cast(Stmt *&val, DataType dt) {
     auto cast_stmt = insert_type_cast_after(val, val, dt);
     val = cast_stmt;
@@ -287,7 +310,27 @@ class TypeCheck : public IRVisitor {
       };
       stmt->lhs = promote_custom_int_type(stmt, stmt->lhs);
       stmt->rhs = promote_custom_int_type(stmt, stmt->rhs);
-      auto ret_type = promoted_type(stmt->lhs->ret_type, stmt->rhs->ret_type);
+
+      DataType ret_type;
+      if (is_shift_op(stmt->op_type)) {
+        // shift_ops does not follow the same type promotion rule as numerical
+        // ops numerical ops: u8 + i32 = i32 shift_ops:     u8 << i32 = u8
+        // (return dtype follows that of the lhs)
+        //
+        // In the above example, while truncating rhs(i32) to u8 risks an
+        // overflow, the runtime value of rhs is very likely less than 8
+        // (otherwise meaningless). Nevertheless, we insert an AssertStmt here
+        // to warn user of this potential overflow.
+        ret_type = stmt->lhs->ret_type;
+
+        // Insert AssertStmt
+        if (config_.debug) {
+          insert_shift_op_assertion_before(stmt, stmt->lhs, stmt->rhs);
+        }
+      } else {
+        ret_type = promoted_type(stmt->lhs->ret_type, stmt->rhs->ret_type);
+      }
+
       if (ret_type != stmt->lhs->ret_type) {
         // promote lhs
         auto cast_stmt = insert_type_cast_before(stmt, stmt->lhs, ret_type);

tests/cpp/ir/ir_type_promotion_test.cpp

@@ -0,0 +1,31 @@
+#include "gtest/gtest.h"
+
+#include "taichi/ir/statements.h"
+#include "taichi/ir/ir_builder.h"
+#include "taichi/ir/transforms.h"
+#include "tests/cpp/program/test_program.h"
+
+namespace taichi {
+namespace lang {
+
+TEST(IRTypePromotionTest, ShiftOp) {
+  IRBuilder builder;
+
+  // (u8)x << (i32)1 -> (u8)res
+  auto *lhs = builder.create_arg_load(0, get_data_type<uint8>(), false);
+  auto *res = builder.create_shl(lhs, builder.get_int32(1));
+  auto ir = builder.extract_ir();
+
+  ASSERT_TRUE(ir->is<Block>());
+  auto *ir_block = ir->as<Block>();
+  irpass::type_check(ir_block, CompileConfig());
+
+  EXPECT_TRUE(ir_block->statements.back()->is<BinaryOpStmt>());
+  auto *binary_stmt = ir_block->statements.back()->as<BinaryOpStmt>();
+
+  auto ret_type = binary_stmt->ret_type;
+  EXPECT_TRUE(ret_type->is_primitive(PrimitiveTypeID::u8));
+}
+
+}  // namespace lang
+}  // namespace taichi