Repository: taichi-dev/taichi (language: python)

Issue:
[bug] TensorType arithmetic gives wrong results
As a follow-up of #7061, I found some vector arithmetics gave me wrong numerical results:

```Python
import taichi as ti

ti.init(arch=ti.cuda, print_kernel_llvm_ir_optimized=True, print_ir=True, real_matrix_scalarize=False)

f = ti.ndarray(ti.math.vec4, 4)

@ti.kernel
def fill(arr: ti.types.ndarray()):
    v0 = ti.math.vec4([0.0, 1.0, 2.0, 3.0])
    v1 = ti.math.vec4([1.0, 2.0, 3.0, 4.0])
    v2 = ti.math.vec4([2.0, 3.0, 4.0, 5.0])
    v3 = ti.math.vec4([4.0, 5.0, 6.0, 7.0])
    arr[0] = v0
    arr[1] = v1
    arr[2] = v2
    arr[3] = v3

@ti.kernel
def vec_test(arr : ti.types.ndarray()):
    v0 = arr[0]
    v1 = arr[1]
    v2 = arr[2]
    v3 = arr[3]

    arr[0] = v0 * v1 + v2
    arr[1] = v1 * v2 + v3
    arr[2] = v0 * v2 + v3

fill(f)
vec_test(f)
print(f.to_numpy())
```
Screen output was
```
[[ 2.  6.  4.  9.]
 [16. 31.  7.  0.]
 [ 0.  0.  0.  0.]
 [ 0.  0.  0.  0.]]
```
Which is apparently unexpected.


Discussion / hints:

Fix patch:
diff --git a/taichi/codegen/llvm/codegen_llvm.cpp b/taichi/codegen/llvm/codegen_llvm.cpp
index 6a036ae7c54c2..234ad2426e5e8 100644
--- a/taichi/codegen/llvm/codegen_llvm.cpp
+++ b/taichi/codegen/llvm/codegen_llvm.cpp
@@ -1897,14 +1897,15 @@ void TaskCodeGenLLVM::visit(ExternalPtrStmt *stmt) {
     auto address_offset = builder->CreateSExt(
         linear_index, llvm::Type::getInt64Ty(*llvm_context));
 
-    if (stmt->ret_type->is<TensorType>()) {
+    auto stmt_ret_type = stmt->ret_type.ptr_removed();
+    if (stmt_ret_type->is<TensorType>()) {
       // This case corresponds to outter indexing only
       // The stride for linear_index is num_elements() in TensorType.
       address_offset = builder->CreateMul(
           address_offset,
           tlctx->get_constant(
               get_data_type<int64>(),
-              stmt->ret_type->cast<TensorType>()->get_num_elements()));
+              stmt_ret_type->cast<TensorType>()->get_num_elements()));
     } else {
       // This case corresponds to outter + inner indexing
       // Since both outter and inner indices are linearized into linear_index,


Test patch:
diff --git a/tests/python/test_matrix.py b/tests/python/test_matrix.py
index 24e752bc52888..a9aa00b085ee9 100644
--- a/tests/python/test_matrix.py
+++ b/tests/python/test_matrix.py
@@ -1224,3 +1224,37 @@ def foo():
         assert a == 2.5
 
     foo()
+
+
+@test_utils.test(arch=[ti.cpu, ti.cuda], real_matrix_scalarize=False)
+def test_matrix_arithmatics():
+    f = ti.ndarray(ti.math.vec4, 4)
+
+    @ti.kernel
+    def fill(arr: ti.types.ndarray()):
+        v0 = ti.math.vec4([0.0, 1.0, 2.0, 3.0])
+        v1 = ti.math.vec4([1.0, 2.0, 3.0, 4.0])
+        v2 = ti.math.vec4([2.0, 3.0, 4.0, 5.0])
+        v3 = ti.math.vec4([4.0, 5.0, 6.0, 7.0])
+        arr[0] = v0
+        arr[1] = v1
+        arr[2] = v2
+        arr[3] = v3
+
+    @ti.kernel
+    def vec_test(arr: ti.types.ndarray()):
+        v0 = arr[0]
+        v1 = arr[1]
+        v2 = arr[2]
+        v3 = arr[3]
+
+        arr[0] = v0 * v1 + v2
+        arr[1] = v1 * v2 + v3
+        arr[2] = v0 * v2 + v3
+
+    fill(f)
+    vec_test(f)
+
+    assert (f.to_numpy() == np.array([[2., 5., 10., 17.], [6., 11., 18., 27.],
+                                      [4., 8., 14., 22.], [4., 5., 6.,
+                                                           7.]])).all()