Update reshape() implementation to avoid unnecessary deepcopy (#543)

* Update reshape() implementation to avoid unnecessary deepcopy

* Optimize reshape function and update tests

* Refactor reshaping logic and add type hints to test

Author: shijie-nv

cupynumeric/_thunk/deferred.py

@@ -14,6 +14,7 @@
 #
 from __future__ import annotations
 
+import math
 import weakref
 from collections import Counter
 from collections.abc import Iterable
@@ -1091,134 +1092,64 @@ def reshape(self, newshape: NdShape, order: OrderType) -> NumPyThunk:
         # performance issues, but we will revisit this decision later once
         # we have enough evidence that that's not the case.
 
-        in_dim = 0
-        out_dim = 0
-
         in_shape = self.shape
         out_shape = newshape
 
-        in_ndim = len(in_shape)
-        out_ndim = len(out_shape)
-
-        groups = []
-
-        while in_dim < in_ndim and out_dim < out_ndim:
-            prev_in_dim = in_dim
-            prev_out_dim = out_dim
-
-            in_prod = 1
+        needs_copy = False
+        out_pos = 0
+        for in_elem in in_shape:
             out_prod = 1
-
-            while True:
-                if in_prod < out_prod:
-                    in_prod *= in_shape[in_dim]
-                    in_dim += 1
-                else:
-                    out_prod *= out_shape[out_dim]
-                    out_dim += 1
-                if in_prod == out_prod:
-                    if in_dim < in_ndim and in_shape[in_dim] == 1:
-                        in_dim += 1
-                    break
-
-            in_group = in_shape[prev_in_dim:in_dim]
-            out_group = out_shape[prev_out_dim:out_dim]
-            groups.append((in_group, out_group))
-
-        while in_dim < in_ndim:
-            assert in_shape[in_dim] == 1
-            groups.append(((1,), ()))
-            in_dim += 1
-
-        while out_dim < out_ndim:
-            assert out_shape[out_dim] == 1
-            groups.append(((), (1,)))
-            out_dim += 1
-
-        needs_linearization = any(len(src_g) > 1 for src_g, _ in groups)
-        needs_delinearization = any(len(tgt_g) > 1 for _, tgt_g in groups)
-        needs_copy = needs_linearization or needs_delinearization
+            while out_prod < in_elem and out_pos < len(out_shape):
+                out_prod *= out_shape[out_pos]
+                out_pos += 1
+            if out_prod != in_elem:
+                needs_copy = True
+                break
 
         if needs_copy:
-            tmp_shape: NdShape = ()
-            for src_g, tgt_g in groups:
-                if len(src_g) > 1 and len(tgt_g) > 1:
-                    tmp_shape += (_prod(tgt_g),)
-                else:
-                    tmp_shape += tgt_g
-
-            result = runtime.create_empty_thunk(
-                tmp_shape, dtype=self.base.type, inputs=[self]
+            flat_size = math.prod(in_shape)
+            flat_array = runtime.create_empty_thunk(
+                (flat_size,), dtype=self.base.type, inputs=[self]
             )
+            in_shape_store = flat_array.base.delinearize(0, in_shape)
+            out_shape_store = flat_array.base.delinearize(0, out_shape)
+            in_shape_array = DeferredArray(in_shape_store)
+            in_shape_array.copy(self, deep=True)
+            result = DeferredArray(out_shape_store)
+            return result
+
+        src = self.base
+
+        # Process each dimension from right to left
+        out_pos = len(out_shape) - 1
+        for src_dim, elem_in in zip(
+            range(len(in_shape) - 1, -1, -1), reversed(in_shape)
+        ):
+            if out_pos >= 0 and elem_in == out_shape[out_pos]:
+                # Case 1: Dimensions match exactly
+                out_pos -= 1
+                continue
+
+            if elem_in == 1:
+                # Case 2: Input dimension is 1 (projection)
+                src = src.project(src_dim, 0)
+                continue
+
+            # Case 3: Delinearize operation
+            new_sizes = []
+            out_prod = 1
+            while out_prod < elem_in and out_pos >= 0:
+                out_prod *= out_shape[out_pos]
+                new_sizes.append(out_shape[out_pos])
+                out_pos -= 1
 
-            src = self.base
-            tgt = result.base  # type: ignore
-
-            src_dim = 0
-            tgt_dim = 0
-            for src_g, tgt_g in groups:
-                diff = 1
-                if src_g == tgt_g:
-                    assert len(src_g) == 1
-                elif len(src_g) == 0:
-                    assert tgt_g == (1,)
-                    src = src.promote(src_dim, 1)
-                elif len(tgt_g) == 0:
-                    assert src_g == (1,)
-                    tgt = tgt.promote(tgt_dim, 1)
-                elif len(src_g) == 1:
-                    src = src.delinearize(src_dim, tgt_g)
-                    diff = len(tgt_g)
-                else:
-                    tgt = tgt.delinearize(tgt_dim, src_g)
-                    diff = len(src_g)
-
-                src_dim += diff
-                tgt_dim += diff
-
-            assert src.shape == tgt.shape
-
-            src_array = DeferredArray(src)
-            tgt_array = DeferredArray(tgt)
-            tgt_array.copy(src_array, deep=True)
-
-            if needs_delinearization and needs_linearization:
-                src = result.base  # type: ignore
-                src_dim = 0
-                for src_g, tgt_g in groups:
-                    if len(src_g) > 1 and len(tgt_g) > 1:
-                        src = src.delinearize(src_dim, tgt_g)
-                        src_dim += len(tgt_g)
-
-                assert src.shape == newshape
-                src_array = DeferredArray(src)
-                result = runtime.create_empty_thunk(
-                    newshape, dtype=self.base.type, inputs=[self]
-                )
-                result.copy(src_array, deep=True)
-
-        else:
-            src = self.base
-            src_dim = 0
-            for src_g, tgt_g in groups:
-                diff = 1
-                if src_g == tgt_g:
-                    assert len(src_g) == 1
-                elif len(src_g) == 0:
-                    assert tgt_g == (1,)
-                    src = src.promote(src_dim, 1)
-                elif len(tgt_g) == 0:
-                    assert src_g == (1,)
-                    src = src.project(src_dim, 0)
-                    diff = 0
-                else:
-                    # unreachable
-                    assert False
-
-                src_dim += diff
+            src = src.delinearize(src_dim, tuple(reversed(new_sizes)))
 
-            result = DeferredArray(src)
+        # Process remaining output dimensions by adding new dimensions (promotion)
+        for _ in range(out_pos + 1):
+            src = src.promote(0, 1)
 
+        result = DeferredArray(src)
         return result
 
     def squeeze(self, axis: int | tuple[int, ...] | None) -> DeferredArray:

tests/integration/test_reshape_copy.py

@@ -0,0 +1,42 @@
+# Copyright 2024 NVIDIA Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pytest
+
+import cupynumeric as num
+
+
+@pytest.mark.parametrize(
+    "in_shape, out_shape",
+    [
+        ((1, 1, 10), (10,)),
+        ((6, 1, 1), (1, 1, 6)),
+        ((12, 1), (1, 3, 4)),
+        ((12, 1, 4), (2, 3, 2, 4)),
+    ],
+)
+def test_reshape_no_copy(in_shape: tuple, out_shape: tuple) -> None:
+    x = num.zeros(in_shape, dtype=num.int32)
+    y = num.reshape(x, out_shape)
+    x.fill(1)
+
+    assert y.shape == out_shape
+    assert num.sum(y) != 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(sys.argv))