[bug] Matrix refactor bug fix: Fix cross scope matrix operations (#6822)

Issue: #5819

### Brief Summary

python/taichi/lang/matrix.py

@@ -421,6 +421,7 @@ class Matrix(TaichiOperations):
         0
     """
     _is_taichi_class = True
+    _is_matrix_class = True
     __array_priority__ = 1000
 
     def __init__(self, arr, dt=None, is_ref=False, ndim=None):
@@ -665,6 +666,9 @@ def _subscript(self, *indices):
         is_global_mat = isinstance(self, _MatrixFieldElement)
         return self._impl._subscript(is_global_mat, *indices)
 
+    def _make_matrix(self):
+        return make_matrix(self._impl.entries)
+
     def to_list(self):
         """Return this matrix as a 1D `list`.
 

python/taichi/lang/ops.py

@@ -7,7 +7,26 @@
 from taichi.lang import expr, impl
 from taichi.lang.exception import TaichiSyntaxError
 from taichi.lang.field import Field
-from taichi.lang.util import cook_dtype, is_taichi_class, taichi_scope
+from taichi.lang.util import (cook_dtype, is_matrix_class, is_taichi_class,
+                              taichi_scope)
+
+
+def uniform_matrix_inputs(*args):
+    has_real_matrix = False
+    for arg in args:
+        if is_taichi_expr(arg) and arg.ptr.is_tensor():
+            has_real_matrix = True
+            break
+
+    results = []
+    for arg in args:
+        if has_real_matrix and is_matrix_class(arg):
+            results.append(arg._make_matrix())
+        else:
+            results.append(arg)
+
+    return results
+
 
 unary_ops = []
 
@@ -52,6 +71,8 @@ def rev_foo(x, y):
 
     @functools.wraps(foo)
     def wrapped(a, b):
+        a, b = uniform_matrix_inputs(a, b)
+
         if isinstance(a, Field) or isinstance(b, Field):
             return NotImplemented
         if is_taichi_class(a):
@@ -82,6 +103,8 @@ def cab_foo(c, a, b):
 
     @functools.wraps(foo)
     def wrapped(a, b, c):
+        a, b, c = uniform_matrix_inputs(a, b, c)
+
         if isinstance(a, Field) or isinstance(b, Field) or isinstance(
                 c, Field):
             return NotImplemented
@@ -107,6 +130,8 @@ def imp_foo(x, y):
 
     @functools.wraps(foo)
     def wrapped(a, b):
+        a, b = uniform_matrix_inputs(a, b)
+
         if isinstance(a, Field) or isinstance(b, Field):
             return NotImplemented
         if is_taichi_class(a):

python/taichi/lang/util.py

@@ -68,6 +68,16 @@ def get_clangpp():
     return _clangpp_presence
 
 
+def is_matrix_class(rhs):
+    matrix_class = False
+    try:
+        if rhs._is_matrix_class:
+            matrix_class = True
+    except:
+        pass
+    return matrix_class
+
+
 def is_taichi_class(rhs):
     taichi_class = False
     try:

tests/python/test_matrix.py

@@ -1225,6 +1225,70 @@ def foo():
         foo()
 
 
+@test_utils.test(arch=[ti.cuda, ti.cpu],
+                 real_matrix=True,
+                 real_matrix_scalarize=True,
+                 debug=True)
+def test_cross_scope_matrix_binary_ops():
+    n = 128
+    x = ti.Vector.field(3, dtype=int, shape=(n, n))
+    spring_offsets = [ti.Vector([1, 2]), ti.Vector([2, 3])]
+
+    @ti.kernel
+    def test():
+        vec = ti.Vector([4, 5])
+        ind0 = vec + ti.static(spring_offsets)[0]
+        ind1 = ti.lang.ops.add(vec, ti.static(spring_offsets)[1])
+
+        x[ind0] = [100, 10, 1]
+        x[ind1] = [1, 10, 100]
+
+    test()
+
+    assert (x[5, 7] == [100, 10, 1]).all()
+    assert (x[6, 8] == [1, 10, 100]).all()
+
+
+@test_utils.test(arch=[ti.cuda, ti.cpu],
+                 real_matrix=True,
+                 real_matrix_scalarize=True,
+                 debug=True)
+def test_cross_scope_matrix_ternary_ops():
+    n = 128
+    x = ti.Vector.field(3, dtype=int, shape=(n, n))
+    spring_offsets = [ti.Vector([1, 2]), ti.Vector([2, 3])]
+
+    @ti.kernel
+    def test():
+        vec = ti.Vector([0, 1])
+        ind0 = ti.select(vec, vec, ti.static(spring_offsets)[0])
+        x[ind0] = [100, 10, 1]
+
+    test()
+
+    assert (x[1, 1] == [100, 10, 1]).all()
+
+
+@test_utils.test(arch=[ti.cuda, ti.cpu],
+                 real_matrix=True,
+                 real_matrix_scalarize=True,
+                 debug=True)
+def test_cross_scope_matrix_atomic_ops():
+    n = 128
+    x = ti.Vector.field(3, dtype=int, shape=(n, n))
+    spring_offsets = [ti.Vector([1, 2]), ti.Vector([2, 3])]
+
+    @ti.kernel
+    def test():
+        vec = ti.Vector([0, 1])
+        vec += ti.static(spring_offsets)[0]
+        x[vec] = [100, 10, 1]
+
+    test()
+
+    assert (x[1, 3] == [100, 10, 1]).all()
+
+
 @test_utils.test(require=ti.extension.dynamic_index,
                  dynamic_index=True,
                  debug=True)