[Lang] Make tensor indicing automatically cast to int32 and raise an warning

python/taichi/lang/impl.py

@@ -99,9 +99,8 @@ def subscript(value, *indices):
     else:
         if isinstance(indices,
                       tuple) and len(indices) == 1 and indices[0] is None:
-            indices_expr_group = make_expr_group()
-        else:
-            indices_expr_group = make_expr_group(*indices)
+            indices = []
+        indices_expr_group = make_expr_group(*indices)
         tensor_dim = int(value.ptr.get_attribute("dim"))
         index_dim = indices_expr_group.size()
         assert tensor_dim == index_dim, f'Tensor with dim {tensor_dim} accessed with indices of dim {index_dim}'

taichi/transforms/type_check.cpp

@@ -94,10 +94,10 @@ class TypeCheck : public IRVisitor {
     }
     if (stmt->ptr->ret_type.data_type != common_container_type) {
       TI_WARN(
-          "[{}] Local store may lose precision (target = {}, value = {}, at",
+          "[{}] Local store may lose precision (target = {}, value = {}) at",
           stmt->name(), stmt->ptr->ret_data_type_name(),
           old_data->ret_data_type_name(), stmt->id);
-      fmt::print(stmt->tb);
+      TI_WARN("\n{}", stmt->tb);
     }
     stmt->ret_type = stmt->ptr->ret_type;
   }
@@ -127,13 +127,13 @@ class TypeCheck : public IRVisitor {
       }
     }
     for (int i = 0; i < stmt->indices.size(); i++) {
-      TI_ASSERT_INFO(
-          is_integral(stmt->indices[i]->ret_type.data_type),
-          "[{}] Taichi tensors must be accessed with integral indices (e.g., "
-          "i32/i64). It seems that you have used something else as "
-          "an index. You can cast a float-pointer number to an integer using "
-          "int(). Also note that ti.floor(ti.f32) returns f32.",
-          stmt->name());
+      if (!is_integral(stmt->indices[i]->ret_type.data_type)) {
+        TI_WARN(
+            "[{}] Tensor index {} not integral, casting into int32 implicitly",
+            stmt->name(), i);
+        stmt->indices[i] =
+            insert_type_cast_before(stmt, stmt->indices[i], DataType::i32);
+      }
       TI_ASSERT(stmt->indices[i]->ret_type.width == stmt->snodes.size());
     }
   }
@@ -148,8 +148,8 @@ class TypeCheck : public IRVisitor {
     }
     if (stmt->ptr->ret_type.data_type != promoted) {
       TI_WARN("[{}] Global store may lose precision: {} <- {}, at",
-              stmt->name(), stmt->ptr->ret_data_type_name(), input_type,
-              stmt->tb);
+              stmt->name(), stmt->ptr->ret_data_type_name(), input_type);
+      TI_WARN("\n{}", stmt->tb);
     }
     stmt->ret_type = stmt->ptr->ret_type;
   }
@@ -229,7 +229,7 @@ class TypeCheck : public IRVisitor {
       } else {
         TI_WARN("[{}] {} at", stmt->name(), comment);
       }
-      fmt::print(stmt->tb);
+      TI_WARN("\n{}", stmt->tb);
       TI_WARN("Compilation stopped due to type mismatch.");
       throw std::runtime_error("Binary operator type mismatch");
     };

tests/python/test_indices.py

@@ -19,3 +19,22 @@ def test_indices():
     assert mapping_b == {0: 1, 1: 0}
     # Note that b is column-major:
     # the virtual first index exposed to the user comes second in memory layout.
+
+
+@ti.host_arch_only
+def test_float_as_index():
+    a = ti.var(ti.f32, (8, 5))
+
+    @ti.kernel
+    def func():
+        i = 6.66
+        j = 3
+        I = ti.Vector([2, 1])
+        for _ in range(1):  # prevent constant fold
+            a[i, j] = 233
+            a[I + ti.Vector([1, 3.0])] = 666
+
+    func()
+
+    assert a[6, 3] == 233
+    assert a[3, 4] == 666