Revert "Convert aten.view to ttnn.reshape (#221)" (#239)

This reverts commit 8135edd.

Author: ayerofieiev-tt

tests/lowering/tensor_manipulation/test_view.py

@@ -132,11 +132,6 @@ def is_function_call(node) -> bool:
 ]
 
 
-def can_be_tilized(node):
-    size = node.meta["val"].size()
-    return len(size) >= 2 and size[-1] % 32 == 0 and size[-2] % 32 == 0
-
-
 # For operations limitations
 # See https://github.com/tenstorrent-metal/tt-metal/blob/main/ttnn/README.md?plain=1#L19
 def is_tt_compute(node) -> bool:
@@ -269,7 +264,7 @@ def try_add_data_move_in(src_node, dst_idx, dst_node, device) -> torch.fx.node.N
     with g.inserting_before(dst_node):
         kwargs = {}
         if (
-            (dst_node.target == ttnn.reshape and not can_be_tilized(dst_node))
+            dst_node.target == ttnn.reshape
             or dst_node.target == ttnn.embedding
             or dst_node.target == ttnn.zeros_like
             or dst_node.target == target_wrappers.repeat
@@ -282,7 +277,9 @@ def try_add_data_move_in(src_node, dst_idx, dst_node, device) -> torch.fx.node.N
             kwargs["dtype"] = TtnnBfloat16()
 
         # For reshape only put tensor on device if rank is 4
-        if is_tt_compute(dst_node):
+        if (is_tt_compute(dst_node) and dst_node.target != ttnn.reshape) or (
+            dst_node.target == ttnn.reshape and len(dst_node.args[1]) == 4
+        ):
             kwargs["device"] = device
 
         new_nodes.append(g.call_function(ttnn.from_torch, (src_node,), kwargs))
@@ -305,12 +302,7 @@ def try_add_layout_change_before_node(src_node, dst_idx, dst_node) -> torch.fx.n
         return None
     if not is_function_call(dst_node):
         return None
-    if (
-        dst_node.target not in layout_change_ops
-        or dst_idx != 0
-        or not is_tt(src_node)
-        or (dst_node.target == ttnn.reshape and can_be_tilized(dst_node))
-    ):
+    if dst_node.target not in layout_change_ops or dst_idx != 0 or not is_tt(src_node):
         return None
 
     g = dst_node.graph
@@ -326,12 +318,7 @@ def try_add_layout_change_after_node(src_node, dst_idx, dst_node) -> torch.fx.no
     # Consider src_node is ttnn.repeat, and dst_node should be any tt_compute node that uses ttnn.repeat
     if not is_function_call(src_node):
         return None
-    if (
-        src_node.target not in layout_change_ops
-        or not is_tt_compute(dst_node)
-        or dst_node.target == ttnn.embedding
-        or (src_node.target == ttnn.reshape and can_be_tilized(src_node))
-    ):
+    if src_node.target not in layout_change_ops or not is_tt_compute(dst_node) or dst_node.target == ttnn.embedding:
         return None
 
     g = dst_node.graph

torch_ttnn/passes/lowering/add_data_move_pass.py

@@ -295,6 +295,10 @@ def call_function(self, target, args, kwargs):
         if target == torch.ops.aten.permute.default:
             return self.call_function_prop_meta(ttnn.permute, args, kwargs)
 
+        if target == torch.ops.aten.view.default:
+            # aten.reshape is more stable if the input nodes have changed
+            return self.call_function_prop_meta(torch.ops.aten.reshape.default, args, kwargs)
+
         ############################################################
         # Other ops
         ############################################################
@@ -619,56 +623,6 @@ def rewrite_node(node):
                     input = g.call_function(ttnn.to_layout, args=(input, TtnnRowMajorLayout()))
                 return g.call_function(ttnn.pad, args=(input, full_pad, value))
 
-            if node.target == torch.ops.aten.view.default:
-                source_shape = args[0].meta["val"].size()
-                out_shape = args[1]
-                source_rank = len(source_shape)
-                out_rank = len(out_shape)
-
-                # Allow lowering by default
-                can_reshape = True
-
-                # Unsupported:
-                # (1) -> (1, 1) or (1, 1, 1), etc
-                if (source_rank == 1) and (np.prod(source_shape) == 1):
-                    can_reshape = False
-                elif not has_valid_page_size(source_shape):
-                    can_reshape = False
-                # Same as ttnn.squeeze with dim = 0
-                # Supported:
-                # (1, 16, 256, 256) -> (16, 256, 256)
-                # (1, 256, 256) - > (256, 256)
-                elif (source_rank != 1) and (out_rank == (source_rank - 1)) and (source_shape[0] == 1):
-                    for i in range(0, out_rank):
-                        if source_shape[i + 1] != out_shape[i]:
-                            can_reshape = False
-                            break
-
-                # Same as ttnn.unsqueeze_to_4D
-                # Supported:
-                # (16, 256, 256) -> (1, 16, 256, 256)
-                # (256, 256) -> (1, 1, 256, 256)
-                elif (out_rank > 1) and (out_rank <= 4) and (source_rank > 0) and (source_rank <= 4):
-                    for i in range(0, out_rank):
-                        si = i + (source_rank - out_rank)
-                        if si < 0:
-                            if out_shape[i] != 1:
-                                can_reshape = False
-                                break
-                        else:
-                            if out_shape[i] != source_shape[si]:
-                                can_reshape = False
-                                break
-                else:
-                    can_reshape = False
-
-                # Transform to ttnn.reshape if possible
-                if can_reshape:
-                    return g.call_function(ttnn.reshape, (args[0], args[1]), {})
-                else:
-                    # Fallback: aten.reshape is more stable if the input nodes have changed
-                    return g.call_function(torch.ops.aten.reshape.default, args, kwargs)
-
         with g.inserting_before(node):
             new_node = rewrite_node(node)
             if new_node is not None: