[IR] Add type inference for loop variables (taichi-dev#3437)

zerolyj · Nov 10, 2021 · b0f66ab · b0f66ab
1 parent b170071
commit b0f66ab
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 4 deletions.
diff --git a/taichi/ir/frontend_ir.cpp b/taichi/ir/frontend_ir.cpp
@@ -53,6 +53,7 @@ FrontendForStmt::FrontendForStmt(const ExprGroup &loop_var,
   loop_var_id.resize(loop_var.size());
   for (int i = 0; i < (int)loop_var.size(); i++) {
     loop_var_id[i] = loop_var[i].cast<IdExpression>()->id;
+    loop_var[i].expr->ret_type = PrimitiveType::i32;
   }
 }
 
@@ -86,6 +87,7 @@ FrontendForStmt::FrontendForStmt(const Expr &loop_var,
     vectorize = 1;
   loop_var_id.resize(1);
   loop_var_id[0] = loop_var.cast<IdExpression>()->id;
+  loop_var.expr->ret_type = PrimitiveType::i32;
 }
 
 void ArgLoadExpression::type_check() {

diff --git a/tests/python/test_grouped.py b/tests/python/test_grouped.py
@@ -16,7 +16,7 @@ def test():
         for i in range(n):
             for j in range(m):
                 for k in range(p):
-                    I = ti.Vector([i, j, k], dt=ti.i32)
+                    I = ti.Vector([i, j, k])
                     val[I] = i + j * 2 + k * 3
 
     test()

diff --git a/tests/python/test_mpm88.py b/tests/python/test_mpm88.py
@@ -150,8 +150,7 @@ def substep(x: ti.any_arr(element_dim=1), v: ti.any_arr(element_dim=1),
             fx = x[p] * inv_dx - base.cast(float)
             w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
             stress = -dt * p_vol * (J[p] - 1) * 4 * inv_dx * inv_dx * E
-            affine = ti.Matrix([[stress, 0], [0, stress]],
-                               dt=ti.f32) + p_mass * C[p]
+            affine = ti.Matrix([[stress, 0], [0, stress]]) + p_mass * C[p]
             for i in ti.static(range(3)):
                 for j in ti.static(range(3)):
                     offset = ti.Vector([i, j])

diff --git a/tests/python/test_tensor_dimensionality.py b/tests/python/test_tensor_dimensionality.py
@@ -8,7 +8,7 @@ def _test_dimensionality(d):
     @ti.kernel
     def fill():
         for I in ti.grouped(x):
-            x[I] += ti.Vector([I.sum(), I[0]], dt=ti.i32)
+            x[I] += ti.Vector([I.sum(), I[0]])
 
     for i in range(2**d):
         indices = []