From c8b1166ed5f48f7fb0b223c1e09b45b1ddf09eac Mon Sep 17 00:00:00 2001 From: Bo Qiao Date: Tue, 22 Mar 2022 10:47:53 +0800 Subject: [PATCH 1/4] [lang] Fix matrix ndarray to numpy --- python/taichi/_kernels.py | 16 ++++++++++++---- python/taichi/lang/_ndarray.py | 10 ++++++---- python/taichi/lang/matrix.py | 8 ++++---- tests/python/test_ndarray.py | 31 +++++++++++++++++++++++++++++++ 4 files changed, 53 insertions(+), 12 deletions(-) diff --git a/python/taichi/_kernels.py b/python/taichi/_kernels.py index afbb8dc68f021..eb7825b3d086d 100644 --- a/python/taichi/_kernels.py +++ b/python/taichi/_kernels.py @@ -44,14 +44,21 @@ def ndarray_to_ext_arr(ndarray: any_arr(), arr: ext_arr()): @kernel def ndarray_matrix_to_ext_arr(ndarray: any_arr(), arr: ext_arr(), + layout_is_aos: template(), as_vector: template()): for I in grouped(ndarray): for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): - arr[I, p] = ndarray[I][p] + if static(layout_is_aos): + arr[I, p] = ndarray[I][p] + else: + arr[p, I] = ndarray[I][p] else: - arr[I, p, q] = ndarray[I][p, q] + if static(layout_is_aos): + arr[I, p, q] = ndarray[I][p, q] + else: + arr[p, q, I] = ndarray[I][p, q] @kernel @@ -124,14 +131,15 @@ def ext_arr_to_ndarray(arr: ext_arr(), ndarray: any_arr()): @kernel def ext_arr_to_ndarray_matrix(arr: ext_arr(), ndarray: any_arr(), + layout_is_aos: template(), as_vector: template()): for I in grouped(ndarray): for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): - ndarray[I][p] = arr[I, p] + ndarray[I][p] = arr[I, p] if layout_is_aos else arr[p, I] else: - ndarray[I][p, q] = arr[I, p, q] + ndarray[I][p, q] = arr[I, p, q] if layout_is_aos else arr[p, q, I] @kernel diff --git a/python/taichi/lang/_ndarray.py b/python/taichi/lang/_ndarray.py index a1bd856dabf39..7be640a2c45dc 100644 --- a/python/taichi/lang/_ndarray.py +++ b/python/taichi/lang/_ndarray.py @@ -90,7 +90,7 @@ def _ndarray_to_numpy(self): impl.get_runtime().sync() return arr - def _ndarray_matrix_to_numpy(self, as_vector): + def _ndarray_matrix_to_numpy(self, layout, as_vector): """Converts matrix ndarray to a numpy array. Returns: @@ -99,7 +99,8 @@ def _ndarray_matrix_to_numpy(self, as_vector): arr = np.zeros(shape=self.arr.shape, dtype=to_numpy_type(self.dtype)) from taichi._kernels import \ ndarray_matrix_to_ext_arr # pylint: disable=C0415 - ndarray_matrix_to_ext_arr(self, arr, as_vector) + layout_is_aos = 1 if layout == Layout.AOS else 0 + ndarray_matrix_to_ext_arr(self, arr, layout_is_aos, as_vector) impl.get_runtime().sync() return arr @@ -122,7 +123,7 @@ def _ndarray_from_numpy(self, arr): ext_arr_to_ndarray(arr, self) impl.get_runtime().sync() - def _ndarray_matrix_from_numpy(self, arr, as_vector): + def _ndarray_matrix_from_numpy(self, arr, layout, as_vector): """Loads all values from a numpy array. Args: @@ -139,7 +140,8 @@ def _ndarray_matrix_from_numpy(self, arr, as_vector): from taichi._kernels import \ ext_arr_to_ndarray_matrix # pylint: disable=C0415 - ext_arr_to_ndarray_matrix(arr, self, as_vector) + layout_is_aos = 1 if layout == Layout.AOS else 0 + ext_arr_to_ndarray_matrix(arr, self, layout_is_aos, as_vector) impl.get_runtime().sync() @python_scope diff --git a/python/taichi/lang/matrix.py b/python/taichi/lang/matrix.py index 7f442afa0001e..40549b64d17db 100644 --- a/python/taichi/lang/matrix.py +++ b/python/taichi/lang/matrix.py @@ -1416,11 +1416,11 @@ def __getitem__(self, key): @python_scope def to_numpy(self): - return self._ndarray_matrix_to_numpy(as_vector=0) + return self._ndarray_matrix_to_numpy(self.layout, as_vector=0) @python_scope def from_numpy(self, arr): - self._ndarray_matrix_from_numpy(arr, as_vector=0) + self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=0) def __deepcopy__(self, memo=None): ret_arr = MatrixNdarray(self.n, self.m, self.dtype, self.shape, @@ -1474,11 +1474,11 @@ def __getitem__(self, key): @python_scope def to_numpy(self): - return self._ndarray_matrix_to_numpy(as_vector=1) + return self._ndarray_matrix_to_numpy(self.layout, as_vector=1) @python_scope def from_numpy(self, arr): - self._ndarray_matrix_from_numpy(arr, as_vector=1) + self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=1) def __deepcopy__(self, memo=None): ret_arr = VectorNdarray(self.n, self.dtype, self.shape, self.layout) diff --git a/tests/python/test_ndarray.py b/tests/python/test_ndarray.py index 4cb6ebd70b3d2..33ab1ba083d43 100644 --- a/tests/python/test_ndarray.py +++ b/tests/python/test_ndarray.py @@ -335,6 +335,37 @@ def test_ndarray_numpy_io(): _test_ndarray_numpy_io() +def _test_ndarray_matrix_numpy_io(): + n = 5 + m = 2 + + x = ti.Vector.ndarray(n, ti.i32, (m,), ti.Layout.AOS) + x_np = 1 + np.arange(n * m).reshape(m, n).astype(np.int32) + x.from_numpy(x_np) + assert (x_np.flatten() == x.to_numpy().flatten()).all() + + x = ti.Vector.ndarray(n, ti.i32, (m,), ti.Layout.SOA) + x_np = 1 + np.arange(n * m).reshape(n, m).astype(np.int32) + x.from_numpy(x_np) + assert (x_np.flatten() == x.to_numpy().flatten()).all() + + k = 2 + x = ti.Matrix.ndarray(m, k, ti.i32, n, ti.Layout.AOS) + x_np = 1 + np.arange(m * k * n).reshape(n, m, k).astype(np.int32) + x.from_numpy(x_np) + assert (x_np.flatten() == x.to_numpy().flatten()).all() + + x = ti.Matrix.ndarray(m, k, ti.i32, n, ti.Layout.SOA) + x_np = 1 + np.arange(m * k * n).reshape(m, k, n).astype(np.int32) + x.from_numpy(x_np) + assert (x_np.flatten() == x.to_numpy().flatten()).all() + + +@test_utils.test(arch=supported_archs_taichi_ndarray) +def test_ndarray_matrix_numpy_io(): + _test_ndarray_matrix_numpy_io() + + def _test_matrix_ndarray_python_scope(layout): a = ti.Matrix.ndarray(2, 2, ti.i32, 5, layout=layout) for i in range(5): From 5fd26def2b03c983afd968adf09864629f916afb Mon Sep 17 00:00:00 2001 From: Bo Qiao Date: Tue, 22 Mar 2022 10:49:21 +0800 Subject: [PATCH 2/4] Format --- python/taichi/_kernels.py | 6 ++++-- python/taichi/lang/_ndarray.py | 4 ++-- tests/python/test_ndarray.py | 4 ++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/python/taichi/_kernels.py b/python/taichi/_kernels.py index eb7825b3d086d..12a875083d72c 100644 --- a/python/taichi/_kernels.py +++ b/python/taichi/_kernels.py @@ -137,9 +137,11 @@ def ext_arr_to_ndarray_matrix(arr: ext_arr(), ndarray: any_arr(), for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): - ndarray[I][p] = arr[I, p] if layout_is_aos else arr[p, I] + ndarray[I][p] = arr[I, p] if layout_is_aos else arr[p, I] else: - ndarray[I][p, q] = arr[I, p, q] if layout_is_aos else arr[p, q, I] + ndarray[I][p, q] = arr[I, p, + q] if layout_is_aos else arr[p, q, + I] @kernel diff --git a/python/taichi/lang/_ndarray.py b/python/taichi/lang/_ndarray.py index 7be640a2c45dc..9c04d03a68e87 100644 --- a/python/taichi/lang/_ndarray.py +++ b/python/taichi/lang/_ndarray.py @@ -99,7 +99,7 @@ def _ndarray_matrix_to_numpy(self, layout, as_vector): arr = np.zeros(shape=self.arr.shape, dtype=to_numpy_type(self.dtype)) from taichi._kernels import \ ndarray_matrix_to_ext_arr # pylint: disable=C0415 - layout_is_aos = 1 if layout == Layout.AOS else 0 + layout_is_aos = 1 if layout == Layout.AOS else 0 ndarray_matrix_to_ext_arr(self, arr, layout_is_aos, as_vector) impl.get_runtime().sync() return arr @@ -140,7 +140,7 @@ def _ndarray_matrix_from_numpy(self, arr, layout, as_vector): from taichi._kernels import \ ext_arr_to_ndarray_matrix # pylint: disable=C0415 - layout_is_aos = 1 if layout == Layout.AOS else 0 + layout_is_aos = 1 if layout == Layout.AOS else 0 ext_arr_to_ndarray_matrix(arr, self, layout_is_aos, as_vector) impl.get_runtime().sync() diff --git a/tests/python/test_ndarray.py b/tests/python/test_ndarray.py index 33ab1ba083d43..4a006230e2f42 100644 --- a/tests/python/test_ndarray.py +++ b/tests/python/test_ndarray.py @@ -339,12 +339,12 @@ def _test_ndarray_matrix_numpy_io(): n = 5 m = 2 - x = ti.Vector.ndarray(n, ti.i32, (m,), ti.Layout.AOS) + x = ti.Vector.ndarray(n, ti.i32, (m, ), ti.Layout.AOS) x_np = 1 + np.arange(n * m).reshape(m, n).astype(np.int32) x.from_numpy(x_np) assert (x_np.flatten() == x.to_numpy().flatten()).all() - x = ti.Vector.ndarray(n, ti.i32, (m,), ti.Layout.SOA) + x = ti.Vector.ndarray(n, ti.i32, (m, ), ti.Layout.SOA) x_np = 1 + np.arange(n * m).reshape(n, m).astype(np.int32) x.from_numpy(x_np) assert (x_np.flatten() == x.to_numpy().flatten()).all() From e3ccb57ed9f3093918e78e84650ebb8789d223d4 Mon Sep 17 00:00:00 2001 From: Bo Qiao Date: Tue, 22 Mar 2022 13:25:32 +0800 Subject: [PATCH 3/4] Use static --- python/taichi/_kernels.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/taichi/_kernels.py b/python/taichi/_kernels.py index 12a875083d72c..d202aa85b86a0 100644 --- a/python/taichi/_kernels.py +++ b/python/taichi/_kernels.py @@ -137,11 +137,15 @@ def ext_arr_to_ndarray_matrix(arr: ext_arr(), ndarray: any_arr(), for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): - ndarray[I][p] = arr[I, p] if layout_is_aos else arr[p, I] + if static(layout_is_aos): + ndarray[I][p] = arr[I, p] + else: + ndarray[I][p] = arr[p, I] else: - ndarray[I][p, q] = arr[I, p, - q] if layout_is_aos else arr[p, q, - I] + if static(layout_is_aos): + ndarray[I][p, q] = arr[I, p, q] + else: + ndarray[I][p, q] = arr[p, q, I] @kernel From 2645af7232a3ee5f86254f8e54d4d79db9c12bd7 Mon Sep 17 00:00:00 2001 From: Bo Qiao Date: Tue, 22 Mar 2022 13:36:09 +0800 Subject: [PATCH 4/4] Use parametrize --- tests/python/test_ndarray.py | 31 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/tests/python/test_ndarray.py b/tests/python/test_ndarray.py index 4a006230e2f42..50419e17b56c3 100644 --- a/tests/python/test_ndarray.py +++ b/tests/python/test_ndarray.py @@ -335,35 +335,32 @@ def test_ndarray_numpy_io(): _test_ndarray_numpy_io() -def _test_ndarray_matrix_numpy_io(): +def _test_ndarray_matrix_numpy_io(layout): n = 5 m = 2 - x = ti.Vector.ndarray(n, ti.i32, (m, ), ti.Layout.AOS) - x_np = 1 + np.arange(n * m).reshape(m, n).astype(np.int32) - x.from_numpy(x_np) - assert (x_np.flatten() == x.to_numpy().flatten()).all() - - x = ti.Vector.ndarray(n, ti.i32, (m, ), ti.Layout.SOA) - x_np = 1 + np.arange(n * m).reshape(n, m).astype(np.int32) + x = ti.Vector.ndarray(n, ti.i32, (m, ), layout) + if layout == ti.Layout.AOS: + x_np = 1 + np.arange(n * m).reshape(m, n).astype(np.int32) + else: + x_np = 1 + np.arange(n * m).reshape(n, m).astype(np.int32) x.from_numpy(x_np) assert (x_np.flatten() == x.to_numpy().flatten()).all() k = 2 - x = ti.Matrix.ndarray(m, k, ti.i32, n, ti.Layout.AOS) - x_np = 1 + np.arange(m * k * n).reshape(n, m, k).astype(np.int32) - x.from_numpy(x_np) - assert (x_np.flatten() == x.to_numpy().flatten()).all() - - x = ti.Matrix.ndarray(m, k, ti.i32, n, ti.Layout.SOA) - x_np = 1 + np.arange(m * k * n).reshape(m, k, n).astype(np.int32) + x = ti.Matrix.ndarray(m, k, ti.i32, n, layout) + if layout == ti.Layout.AOS: + x_np = 1 + np.arange(m * k * n).reshape(n, m, k).astype(np.int32) + else: + x_np = 1 + np.arange(m * k * n).reshape(m, k, n).astype(np.int32) x.from_numpy(x_np) assert (x_np.flatten() == x.to_numpy().flatten()).all() +@pytest.mark.parametrize('layout', layouts) @test_utils.test(arch=supported_archs_taichi_ndarray) -def test_ndarray_matrix_numpy_io(): - _test_ndarray_matrix_numpy_io() +def test_ndarray_matrix_numpy_io(layout): + _test_ndarray_matrix_numpy_io(layout) def _test_matrix_ndarray_python_scope(layout):