[Lang] Demote pow() with integer exponent

taichi/ir/frontend_ir.cpp

@@ -192,7 +192,8 @@ void BinaryOpExpression::type_check(CompileConfig *config) {
     ret_type = PrimitiveType::i32;
     return;
   }
-  if (is_shift_op(type)) {
+  if (is_shift_op(type) ||
+      (type == BinaryOpType::pow && is_integral(rhs_type))) {
     ret_type = lhs_type;
     return;
   }

taichi/ir/ir_builder.h

@@ -136,6 +136,11 @@ class IRBuilder {
   ConstStmt *get_float32(float32 value);
   ConstStmt *get_float64(float64 value);
 
+  template <typename T>
+  ConstStmt *get_constant(DataType dt, const T &value) {
+    return insert(Stmt::make_typed<ConstStmt>(TypedConstant(dt, value)));
+  }
+
   RandStmt *create_rand(DataType value_type);
 
   // Load kernel arguments.

taichi/transforms/alg_simp.cpp

@@ -300,7 +300,7 @@ class AlgSimp : public BasicStmtVisitor {
                  exponent >= -max_weaken_exponent) {
         // a ** -n -> 1 / a ** n
         if (is_integral(stmt->lhs->ret_type)) {
-          TI_ERROR("negative exponent in integer pow is not allowed.");
+          TI_ERROR("Negative exponent in pow(int, int) is not allowed.");
         }
         auto one = Stmt::make<ConstStmt>(TypedConstant(1));
         auto one_raw = one.get();

taichi/transforms/check_out_of_bound.cpp

@@ -115,7 +115,7 @@ class CheckOutOfBound : public BasicStmtVisitor {
         auto compare = std::make_unique<BinaryOpStmt>(
             BinaryOpType::cmp_ge, stmt->rhs, compare_rhs.get());
         compare->ret_type = PrimitiveType::i32;
-        std::string msg = "Negative exponent for integer pows are not allowed";
+        std::string msg = "Negative exponent in pow(int, int) is not allowed.";
         msg += "\n" + stmt->tb;
         auto assert_stmt = std::make_unique<AssertStmt>(compare.get(), msg,
                                                         std::vector<Stmt *>());

taichi/transforms/constant_fold.cpp

@@ -144,7 +144,7 @@ class ConstantFold : public BasicStmtVisitor {
       if (is_integral(rhs->ret_type)) {
         auto rhs_val = rhs->val.val_int();
         if (rhs_val < 0 && is_integral(stmt->ret_type)) {
-          TI_ERROR("negative exponent in integer pow is not allowed.");
+          TI_ERROR("Negative exponent in pow(int, int) is not allowed.");
         }
       }
     }

taichi/transforms/demote_operations.cpp

@@ -1,4 +1,5 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/ir_builder.h"
 #include "taichi/ir/statements.h"
 #include "taichi/ir/transforms.h"
 #include "taichi/ir/visitors.h"
@@ -118,6 +119,68 @@ class DemoteOperations : public BasicStmtVisitor {
       modifier.insert_before(stmt, std::move(shift));
       modifier.insert_before(stmt, std::move(signed_cast));
       modifier.erase(stmt);
+    } else if (stmt->op_type == BinaryOpType::pow &&
+               is_integral(rhs->element_type())) {
+      // @ti.func
+      // def pow(lhs, rhs):
+      //     a = lhs
+      //     b = abs(rhs)
+      //     result = 1
+      //     while b > 0:
+      //         if b & 1:
+      //             result *= a
+      //         a *= a
+      //         b >>= 1
+      //     if rhs < 0:              # for real lhs
+      //         result = 1 / result  # for real lhs
+      //     return result
+      IRBuilder builder;
+      auto one_lhs = builder.get_constant(lhs->element_type(), 1);
+      auto one_rhs = builder.get_constant(rhs->element_type(), 1);
+      auto zero_rhs = builder.get_constant(rhs->element_type(), 0);
+      auto a = builder.create_local_var(lhs->element_type());
+      builder.create_local_store(a, lhs);
+      auto b = builder.create_local_var(rhs->element_type());
+      builder.create_local_store(b, builder.create_abs(rhs));
+      auto result = builder.create_local_var(lhs->element_type());
+      builder.create_local_store(result, one_lhs);
+      auto loop = builder.create_while_true();
+      {
+        auto loop_guard = builder.get_loop_guard(loop);
+        auto current_a = builder.create_local_load(a);
+        auto current_b = builder.create_local_load(b);
+        auto if_stmt =
+            builder.create_if(builder.create_cmp_le(current_b, zero_rhs));
+        {
+          auto _ = builder.get_if_guard(if_stmt, true);
+          builder.create_break();
+        }
+        auto bit_and = builder.create_and(current_b, one_rhs);
+        if_stmt = builder.create_if(builder.create_cmp_ne(bit_and, zero_rhs));
+        {
+          auto _ = builder.get_if_guard(if_stmt, true);
+          auto current_result = builder.create_local_load(result);
+          auto new_result = builder.create_mul(current_result, current_a);
+          builder.create_local_store(result, new_result);
+        }
+        auto new_a = builder.create_mul(current_a, current_a);
+        builder.create_local_store(a, new_a);
+        auto new_b = builder.create_sar(current_b, one_rhs);
+        builder.create_local_store(b, new_b);
+      }
+      if (is_real(lhs->element_type())) {
+        auto if_stmt = builder.create_if(builder.create_cmp_le(rhs, zero_rhs));
+        {
+          auto _ = builder.get_if_guard(if_stmt, true);
+          auto current_result = builder.create_local_load(result);
+          auto new_result = builder.create_div(one_lhs, current_result);
+          builder.create_local_store(result, new_result);
+        }
+      }
+      auto final_result = builder.create_local_load(result);
+      stmt->replace_usages_with(final_result);
+      modifier.insert_before(stmt, std::move(builder.extract_ir()->statements));
+      modifier.erase(stmt);
     }
   }
 

taichi/transforms/type_check.cpp

@@ -276,6 +276,11 @@ class TypeCheck : public IRVisitor {
     if (stmt->lhs->ret_type->is_primitive(PrimitiveTypeID::unknown) &&
         stmt->rhs->ret_type->is_primitive(PrimitiveTypeID::unknown))
       error();
+    if (stmt->op_type == BinaryOpType::pow &&
+        is_integral(stmt->rhs->ret_type)) {
+      stmt->ret_type = stmt->lhs->ret_type;
+      return;
+    }
 
     // lower truediv into div
 

tests/python/test_pow.py

@@ -15,7 +15,7 @@ def func(x: dt, y: dt):
     for x in [0.5, 1, 1.5, 2, 6.66]:
         for y in [-2, -1, -0.3, 0, 0.5, 1, 1.4, 2.6]:
             func(x, y)
-            assert abs(z[None] / x**y - 1) < 0.00001
+            assert z[None] == pytest.approx(x**y)
 
 
 def _test_pow_i(dt):
@@ -76,3 +76,55 @@ def test_ipow_negative_exp_i32():
                  exclude=[ti.vulkan, ti.opengl, ti.cc])
 def test_ipow_negative_exp_i64():
     _ipow_negative_exp(ti.i64)
+
+
+def _test_pow_int_base_int_exp(dt_base, dt_exp):
+    z = ti.field(dt_base, shape=())
+
+    @ti.kernel
+    def func(x: dt_base, y: dt_exp):
+        z[None] = x**y
+
+    for x in range(-5, 5):
+        for y in range(0, 10):
+            func(x, y)
+            assert z[None] == x**y
+
+
+@test_utils.test()
+def test_pow_int_base_int_exp_32():
+    _test_pow_int_base_int_exp(ti.i32, ti.i32)
+
+
+@pytest.mark.parametrize('dt_base, dt_exp',
+                         [(ti.i32, ti.i64), (ti.i64, ti.i64),
+                          (ti.i64, ti.i32)])
+@test_utils.test(require=ti.extension.data64)
+def test_pow_int_base_int_exp_64(dt_base, dt_exp):
+    _test_pow_int_base_int_exp(dt_base, dt_exp)
+
+
+def _test_pow_float_base_int_exp(dt_base, dt_exp):
+    z = ti.field(dt_base, shape=())
+
+    @ti.kernel
+    def func(x: dt_base, y: dt_exp):
+        z[None] = x**y
+
+    for x in [-6.66, -2, -1.5, -1, -0.5, 0.5, 1, 1.5, 2, 6.66]:
+        for y in range(-10, 10):
+            func(x, y)
+            assert z[None] == pytest.approx(x**y)
+
+
+@test_utils.test()
+def test_pow_float_base_int_exp_32():
+    _test_pow_float_base_int_exp(ti.f32, ti.i32)
+
+
+@pytest.mark.parametrize('dt_base, dt_exp',
+                         [(ti.f64, ti.i32), (ti.f32, ti.i64),
+                          (ti.f64, ti.i64)])
+@test_utils.test(require=ti.extension.data64)
+def test_pow_float_base_int_exp_64(dt_base, dt_exp):
+    _test_pow_float_base_int_exp(dt_base, dt_exp)