Fix unnecessary calls to compute shape of trace tensor; fix #246 (#280)

This MR resolves the following:

1. In a trivial example like below, we were earlier computing shape of
dynamic shape trace tensor which can become very expensive if reshape
occurs somewhere in the middle of a big compute graph. The fix is a
simple hack to force the shape of trace tensor when its statically known
(in this case it is when we convert a shape scalar to 1d tensor).
```py
    a = tp.ones((2, 3, 4))
    s1, s2, s3 = a.shape
    out = tp.reshape(a, (s1, s2, s3 / 2, 2))
```

2. Fixes #246 (thanks to @yizhuoz004 for the suggested fix)

Author: parthchadha

tripy/pyproject.toml

@@ -57,6 +57,8 @@ test = [
   "pytest-xdist==3.6.1",
   "pytest-benchmark==4.0.0",
   "pytest-lazy-fixture==0.6.3",
+  "pytest-mock==3.14.0",
+  "path.py==12.5.0",
   # Triton is required for torch.compile
   "triton==3.0.0",
   "snakeviz==2.2.0",

tripy/tests/frontend/trace/ops/test_unsqueeze.py

@@ -16,15 +16,15 @@
 #
 
 import tripy as tp
-from tripy.frontend.trace.ops import Unsqueeze
+from tripy.frontend.trace.ops import Reshape
 
 
 class TestUnsqueeze:
     def test_func_op(self):
         a = tp.ones((2, 1))
         a = tp.unsqueeze(a, 0)
         assert isinstance(a, tp.Tensor)
-        assert isinstance(a.trace_tensor.producer, Unsqueeze)
+        assert isinstance(a.trace_tensor.producer, Reshape)
 
     def test_infer_rank(self):
         a = tp.ones((2, 1))

tripy/tests/integration/test_reshape.py

@@ -97,3 +97,13 @@ def test_flatten_with_unknown_dims(self):
         a = tp.ones((2, 3, 4, 5))
         b = tp.flatten(a, start_dim=1, end_dim=-1)
         assert np.array_equal(cp.from_dlpack(b).get(), np.ones((2, 60), dtype=np.float32))
+
+
+def test_reshape_with_no_trace_shape_inference(mocker):
+    # This test ensures that Tripy does not unnecessarily compute shape of a trace tensor.
+    mock_function = mocker.patch("tripy.backend.mlir.utils.ShapeContext.get_shape_of_dynamic_trace_tensor")
+    a = tp.ones((2, 3, 4))
+    s1, s2, s3 = a.shape
+    out = tp.reshape(a, (s1, s2, s3 / 2, 2))
+    out.eval()
+    mock_function.assert_not_called()

tripy/tests/integration/test_unsqueeze.py

@@ -35,3 +35,9 @@ def func(a):
         assert tp.allclose(out, tp.Tensor(ref_out))
 
         assert out.shape == ref_out.shape
+
+    def test_unsqueeze_compile(self):
+        def func(a):
+            return tp.unsqueeze(a, 3) == tp.Tensor(3, dtype=tp.float32)
+
+        c = tp.compile(func, args=[tp.InputInfo(((1, 2, 3), 2, 3), dtype=tp.float32)])

tripy/tripy/backend/mlir/utils.py

@@ -380,22 +380,14 @@ def traverse_backwards(tensor, visited_tensors, visited_producers):
             """
             Recurse back from tensor to the inputs to the graph and store the visited tensors and nodes.
             """
-            from tripy.frontend.trace.ops.unsqueeze import Unsqueeze
-
             if id(tensor) in visited_tensors:
                 return
 
             visited_tensors[id(tensor)] = tensor
             if tensor.producer is not None:
                 visited_producers[id(tensor.producer)] = tensor.producer
-                # Special recursion conditions op by op basis.
-                # Only recurse inputs which are used in output shape calculations.
-                if isinstance(tensor.producer, Unsqueeze):
-                    traverse_backwards(tensor.producer.inputs[1], visited_tensors, visited_producers)
-                else:
-                    # Naively recurse all the inputs until a constant or user input.
-                    for input_tensor in tensor.producer.inputs:
-                        traverse_backwards(input_tensor, visited_tensors, visited_producers)
+                for input_tensor in tensor.producer.inputs:
+                    traverse_backwards(input_tensor, visited_tensors, visited_producers)
 
         def find_inputs(graph_nodes):
             """

tripy/tripy/frontend/trace/ops/__init__.py

@@ -39,5 +39,4 @@
 from tripy.frontend.trace.ops.split import Split
 from tripy.frontend.trace.ops.storage import Storage
 from tripy.frontend.trace.ops.unary_elementwise import UnaryElementwise
-from tripy.frontend.trace.ops.unsqueeze import Unsqueeze
 from tripy.frontend.trace.ops.where import Where

tripy/tripy/frontend/trace/ops/unsqueeze.py

@@ -15,46 +15,7 @@
 # limitations under the License.
 #
 
-from dataclasses import dataclass
-
 from tripy import export, constraints
-from tripy.frontend.trace.ops.base import BaseTraceOp
-import tripy.frontend.trace.ops.utils as op_utils
-from tripy.common.datatype import DATA_TYPES
-
-
-@dataclass(repr=False)
-class Unsqueeze(BaseTraceOp):
-
-    dim: int
-
-    # the result will not be rank 1 and so can't be a shape but we may want to unsqueeze shapes
-    infer_shape_output_idxs = op_utils.ShapeOutputIdxPolicies.never_return_shape
-
-    def infer_dtypes(self):
-        self.outputs[0].dtype = self.inputs[0].dtype
-
-    def infer_rank(self):
-        self.outputs[0].rank = self.inputs[0].rank + 1
-
-    def to_flat_ir(self, inputs, outputs):
-        from tripy.flat_ir.ops import DynamicBroadcastOp
-
-        broadcast_dim = list(range(inputs[0].rank))
-        for idx in range(len(broadcast_dim)):
-            if idx >= self.dim:
-                broadcast_dim[idx] += 1
-
-        DynamicBroadcastOp.build(
-            [inputs[0], inputs[1]],
-            [outputs[0]],
-            broadcast_dim=broadcast_dim,
-        )
-
-
-# Two operand unsqueeze op to ensure that Trace op is 1:1 with Python code (for error messaging).
-def unsqueeze_two_operand(input, result_shape, dim):
-    return Unsqueeze.build([input, result_shape], dim)
 
 
 @export.public_api(document_under="operations/functions")
@@ -85,6 +46,7 @@ def unsqueeze(input: "tripy.Tensor", dim: int) -> "tripy.Tensor":
         assert np.array_equal(cp.from_dlpack(output).get(), np.expand_dims(cp.from_dlpack(input).get(), 1))
     """
     from tripy.frontend.trace.ops.concatenate import concatenate
+    from tripy.frontend.trace.ops.reshape import reshape
 
     from tripy.frontend import Shape
 
@@ -97,4 +59,4 @@ def unsqueeze(input: "tripy.Tensor", dim: int) -> "tripy.Tensor":
     else:
         input_shape = input.shape
         result_shape = concatenate([input_shape[:dim], Shape([1]), input_shape[dim:]], dim=0)
-    return unsqueeze_two_operand(input, result_shape, dim)
+    return reshape(input, result_shape)

tripy/tripy/frontend/utils.py

@@ -447,6 +447,9 @@ def convert_nontensor_arg(arg, list_index=None):
                     if member.rank != 0:
                         raise_error("Tensor in a shape argument must be a scalar.", [f"Got {member}"])
                     member = Shape(unsqueeze(member, 0))
+                    # Force the trace tensor shape to be (1,) since its known that we are reshaping a scalar to a 1D tensor.
+                    # If we don't force the shape below, Tripy might require computing the shape of this trace tensor which can be expensive.
+                    member.trace_tensor.shape = (1,)
                     shape_components.append(member)
                 if len(acc) > 0:
                     shape_components.append(convert_nontensor_arg(acc))
@@ -639,6 +642,7 @@ def pretty_print(data_list, shape, threshold=1000, linewidth=10, edgeitems=3):
     """
     Returns a pretty-print string of list format data.
     """
+
     def _data_str(data, summarize, linewidth, edgeitems, indent=0):
         if isinstance(data, (float, int)):
             return str(data)