[Lang] Support Python-scope matrix/vector operations (#1051)

README.md

@@ -4,7 +4,7 @@
   <h3> <a href="https://taichi.readthedocs.io/en/stable/"> Documentation </a> | <a href="https://taichi.readthedocs.io/zh_CN/latest/"> 简体中文文档 </a> </h3>
 </div>
 
-[![Build Status](https://img.shields.io/travis/taichi-dev/taichi?logo=Travis)](https://travis-ci.com/taichi-dev/taichi)
+[![Build Status](https://img.shields.io/travis/taichi-dev/taichi?logo=Travis&branch=master)](https://travis-ci.com/taichi-dev/taichi)
 [![Build Status](https://img.shields.io/appveyor/build/yuanming-hu/taichi?logo=AppVeyor)](https://ci.appveyor.com/project/yuanming-hu/taichi/branch/master)
 [![Latest Release](https://img.shields.io/github/v/release/taichi-dev/taichi?color=blue)](https://github.com/taichi-dev/taichi/releases/latest)
 

docs/matrix.rst

@@ -78,7 +78,7 @@ As a temporary local variable
     ::
 
         # Taichi-scope
-        a = ti.Matrix([[2, 3], [4, 5])
+        a = ti.Matrix([[2, 3], [4, 5]])
 
 
 .. function:: ti.Matrix(rows=[v0, v1, v2, ...])

python/taichi/lang/common_ops.py

@@ -1,3 +1,11 @@
+def numpy_or_constant(x):
+    import taichi as ti
+    if ti.is_taichi_class(x):
+        return x.to_numpy()
+    else:
+        return x
+
+
 class TaichiOperations:
     def __neg__(self):
         import taichi as ti
@@ -8,18 +16,34 @@ def __abs__(self):
         return ti.abs(self)
 
     def __add__(self, other):
+        if self.is_pyconstant():
+            return self.make_from_numpy(self.to_numpy() +
+                                        numpy_or_constant(other))
+
         import taichi as ti
         return ti.add(self, other)
 
     def __radd__(self, other):
+        if self.is_pyconstant():
+            return self.make_from_numpy(
+                numpy_or_constant(other) + self.to_numpy())
+
         import taichi as ti
         return ti.add(other, self)
 
     def __sub__(self, other):
+        if self.is_pyconstant():
+            return self.make_from_numpy(self.to_numpy() -
+                                        numpy_or_constant(other))
+
         import taichi as ti
         return ti.sub(self, other)
 
     def __rsub__(self, other):
+        if self.is_pyconstant():
+            return self.make_from_numpy(
+                numpy_or_constant(other) - self.to_numpy())
+
         import taichi as ti
         return ti.sub(other, self)
 
@@ -51,6 +75,10 @@ def __mod__(self, other):
         import taichi as ti
         return ti.mod(self, other)
 
+    def __rmod__(self, other):
+        import taichi as ti
+        return ti.mod(other, self)
+
     def __pow__(self, other, modulo=None):
         import taichi as ti
         return ti.pow(self, other)
@@ -205,3 +233,10 @@ def __ti_int__(self):
     def __ti_float__(self):
         import taichi as ti
         return ti.cast(self, ti.get_runtime().default_fp)
+
+    def is_pyconstant(self):  # overrided by ti.Matrix
+        return False
+
+    def make_from_numpy(self):
+        raise NotImplementedError(
+            f'Python-scope operation for {type(self)} not implemented yet')

python/taichi/lang/matrix.py

@@ -3,7 +3,7 @@
 import copy
 import numbers
 import numpy as np
-from .util import *
+from .util import taichi_scope, python_scope, deprecated, to_numpy_type
 from .common_ops import TaichiOperations
 from collections.abc import Iterable
 
@@ -28,6 +28,7 @@ def __init__(self,
         # TODO: refactor: use multiple initializer like `ti.Matrix.cols([a, b, c])`
         # and `ti.Matrix.empty(n, m)` instead of ad-hoc `ti.Matrix(cols=[a, b, c])`.
         self.grad = None
+        # construct from rows or cols
         if rows is not None or cols is not None:
             if rows is not None and cols is not None:
                 raise Exception("cannot specify both rows and columns")
@@ -64,7 +65,7 @@ def __init__(self,
                 raise Exception(
                     'cols/rows required when using list of vectors')
             elif not isinstance(n[0], Iterable):
-                if impl.get_runtime().inside_kernel:
+                if impl.inside_kernel():
                     # wrap potential constants with Expr
                     if keep_raw:
                         mat = [list([x]) for x in n]
@@ -80,6 +81,7 @@ def __init__(self,
             else:
                 self.m = 1
             self.entries = [x for row in mat for x in row]
+        # construct global matrix
         else:
             self.entries = []
             self.n = n
@@ -140,10 +142,17 @@ def is_global(self):
             if isinstance(e, expr.Expr):
                 if e.ptr.is_global_var():
                     results[i] = True
-            assert results[i] == results[
-                0], "Matrices with mixed global/local entries are not allowed"
+            assert results[i] == results[0], \
+                "Matrices with mixed global/local entries are not allowed"
         return results[0]
 
+    def is_pyconstant(self):
+        return all([not isinstance(e, expr.Expr) for e in self.entries])
+
+    @staticmethod
+    def make_from_numpy(nparray):
+        return Matrix(nparray)
+
     @taichi_scope
     def element_wise_binary(self, foo, other):
         ret = self.empty_copy()
@@ -155,7 +164,7 @@ def element_wise_binary(self, foo, other):
                 f'taichi class {type(a)}, maybe you want to use `a.fill(b)` instead?'
             )
         if isinstance(other, Matrix):
-            assert self.m == other.m and self.n == other.n
+            assert self.m == other.m and self.n == other.n, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
             for i in range(self.n * self.m):
                 ret.entries[i] = foo(self.entries[i], other.entries[i])
         else:  # assumed to be scalar
@@ -171,9 +180,12 @@ def element_wise_unary(self, foo):
             ret.entries[i] = foo(self.entries[i])
         return ret
 
-    @taichi_scope
     def __matmul__(self, other):
-        assert self.m == other.n
+        # TODO: move to common_ops.py, redirect to `ti.matmul` too?
+        if self.is_pyconstant():
+            return self.make_from_numpy(self.to_numpy() @ other.to_numpy())
+
+        assert self.m == other.n, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
         ret = Matrix(self.n, other.m)
         for i in range(self.n):
             for j in range(other.m):
@@ -289,6 +301,7 @@ def w(self, value):
 
     class Proxy:
         def __init__(self, mat, index):
+            """Proxy when a tensor of Matrices is accessed by host."""
             self.mat = mat
             self.index = index
 
@@ -338,11 +351,6 @@ def w(self, value):
     @python_scope
     def __getitem__(self, index):
         return Matrix.Proxy(self, index)
-        ret = [[] for _ in range(self.n)]
-        for i in range(self.n):
-            for j in range(self.m):
-                ret[i].append(self(i, j)[index])
-        return ret
 
     # host access
     @python_scope
@@ -353,6 +361,14 @@ def __setitem__(self, index, item):
             for j in range(self.m):
                 self(i, j)[index] = item[i][j]
 
+    # host access, return a complete Matrix instead of Proxy
+    def at(self, index):
+        ret = self.empty_copy()
+        for i in range(self.n):
+            for j in range(self.m):
+                ret.entries[i * self.m + j] = self[index][i, j]
+        return ret
+
     def empty_copy(self):
         return Matrix(self.n, self.m, empty=True)
 
@@ -616,6 +632,10 @@ def to_numpy(self, keep_dims=False, as_vector=None):
                 stacklevel=3)
         as_vector = self.m == 1 and not keep_dims
         dim_ext = (self.n, ) if as_vector else (self.n, self.m)
+
+        if self.is_pyconstant():
+            return np.array(self.entries).reshape(dim_ext)
+
         ret = np.empty(self.loop_range().shape() + dim_ext,
                        dtype=to_numpy_type(
                            self.loop_range().snode().data_type()))
@@ -677,6 +697,10 @@ def __ti_repr__(self):
         if self.m != 1:
             yield ']'
 
+    def __repr__(self):
+        """Python scope object print support."""
+        return str(self.to_numpy())
+
     @staticmethod
     @taichi_scope
     def zero(dt, n, m=1):

tests/python/test_matrix.py

@@ -0,0 +1,97 @@
+import taichi as ti
+import numpy as np
+import operator
+
+operation_types = (operator.add, operator.sub, operator.matmul)
+test_matrix_arrays = (np.array([[1, 2], [3, 4]]), np.array([[5, 6], [7, 8]]))
+
+vector_operation_types = (operator.add, operator.sub)
+test_vector_arrays = (np.array([42, 42]), np.array([24, 24]))
+
+
+@ti.host_arch_only
+def test_python_scope_vector_operations():
+    for ops in vector_operation_types:
+        a, b = test_vector_arrays
+        m1, m2 = ti.Vector(a), ti.Vector(b)
+        c = ops(m1, m2)
+        assert np.allclose(c.to_numpy(), ops(a, b))
+
+
+@ti.host_arch_only
+def test_python_scope_matrix_operations():
+    for ops in operation_types:
+        a, b = test_matrix_arrays
+        m1, m2 = ti.Matrix(a), ti.Matrix(b)
+        c = ops(m1, m2)
+        assert np.allclose(c.to_numpy(), ops(a, b))
+
+
+# TODO: Loops inside the function will cause AssertionError:
+# No new variables can be declared after kernel invocations
+# or Python-scope tensor accesses.
+# ideally we should use pytest.fixture to parameterize the tests
+# over explicit loops
+@ti.host_arch_only
+def test_python_scope_vector_tensor_add():
+    t1 = ti.Vector(2, dt=ti.i32, shape=())
+    t2 = ti.Vector(2, dt=ti.i32, shape=())
+    a, b = test_vector_arrays
+    t1[None], t2[None] = a, b
+
+    # TODO: hook Matrix.Proxy to redirect to at + Matrix.__add__
+    c = t1.at(None) + t2.at(None)
+    assert np.allclose(c.to_numpy(), a + b)
+
+
+@ti.host_arch_only
+def test_python_scope_vector_tensor_sub():
+    t1 = ti.Vector(2, dt=ti.i32, shape=())
+    t2 = ti.Vector(2, dt=ti.i32, shape=())
+    a, b = test_vector_arrays
+    t1[None], t2[None] = a, b
+
+    # TODO: hook Matrix.Proxy to redirect to at + Matrix.__sub__
+    c = t1.at(None) - t2.at(None)
+    assert np.allclose(c.to_numpy(), a - b)
+
+
+@ti.host_arch_only
+def test_python_scope_matrix_tensor_add():
+    t1 = ti.Matrix(2, 2, dt=ti.i32, shape=())
+    t2 = ti.Matrix(2, 2, dt=ti.i32, shape=())
+    a, b = test_matrix_arrays
+    # ndarray not supported here
+    t1[None], t2[None] = a.tolist(), b.tolist()
+
+    # TODO: hook Matrix.Proxy to redirect to at + Matrix.__add__
+    c = t1.at(None) + t2.at(None)
+    print(c)
+
+    assert np.allclose(c.to_numpy(), a + b)
+
+
+@ti.host_arch_only
+def test_python_scope_matrix_tensor_sub():
+    t1 = ti.Matrix(2, 2, dt=ti.i32, shape=())
+    t2 = ti.Matrix(2, 2, dt=ti.i32, shape=())
+    a, b = test_matrix_arrays
+    # ndarray not supported here
+    t1[None], t2[None] = a.tolist(), b.tolist()
+
+    # TODO: hook Matrix.Proxy to redirect to at + Matrix.__sub__
+    c = t1.at(None) - t2.at(None)
+    assert np.allclose(c.to_numpy(), a - b)
+
+
+@ti.host_arch_only
+def test_python_scope_matrix_tensor_matmul():
+    t1 = ti.Matrix(2, 2, dt=ti.i32, shape=())
+    t2 = ti.Matrix(2, 2, dt=ti.i32, shape=())
+    a, b = test_matrix_arrays
+    # ndarray not supported here
+    t1[None], t2[None] = a.tolist(), b.tolist()
+
+    # TODO: hook Matrix.Proxy to redirect to at + Matrix.__matmul__
+    c = t1.at(None) @ t2.at(None)
+    assert np.allclose(c.to_numpy(), a @ b)

tests/python/test_numpy_io.py

@@ -3,7 +3,7 @@
 
 
 @ti.all_archs
-def test_from_numpy_2d():
+def test_to_numpy_2d():
     val = ti.var(ti.i32)
 
     n = 4
@@ -24,29 +24,7 @@ def test_from_numpy_2d():
 
 
 @ti.all_archs
-def test_to_numpy_2d():
-    val = ti.var(ti.i32)
-
-    n = 4
-    m = 7
-
-    ti.root.dense(ti.ij, (n, m)).place(val)
-
-    arr = np.empty(shape=(n, m), dtype=np.int32)
-
-    for i in range(n):
-        for j in range(m):
-            arr[i, j] = i + j * 3
-
-    val.from_numpy(arr)
-
-    for i in range(n):
-        for j in range(m):
-            assert val[i, j] == i + j * 3
-
-
-@ti.all_archs
-def test_to_numpy_2d():
+def test_from_numpy_2d():
     val = ti.var(ti.i32)
 
     n = 4