Repository: Samsung/TICO

Issue:
Remove the Reshape surrounding Add → Softmax in the encoder
## What

Right before the attention soft-max, the exported graph still performs an **extra reshape round-trip**.

Eliminating those reshapes would let us feed `Softmax` directly with the `(B × H, T, S)` tensor.

Related: #124

## Why

* **Less data movement** per layer → lower latency on bandwidth-bound NPUs.
* **Cleaner graph** → makes subsequent pattern-matching / fusion passes simpler.

![Image](https://github.com/user-attachments/assets/0eb67bac-b9b7-4925-98b4-791a0f127bac)

Discussion:

Solution patch:
diff --git a/tico/passes/ops.py b/tico/passes/ops.py
index 6e71777..c371490 100644
--- a/tico/passes/ops.py
+++ b/tico/passes/ops.py
@@ -54,7 +54,10 @@ class AtenOps:
         self.reshape = [torch.ops.aten.reshape.default]
         self.select = [torch.ops.aten.select_copy.int, torch.ops.aten.select.int]
         self.slice = [torch.ops.aten.slice.Tensor, torch.ops.aten.slice_copy.Tensor]
-        self.softmax = [torch.ops.aten._softmax.default]
+        self.softmax = [
+            torch.ops.aten._softmax.default,
+            torch.ops.aten._safe_softmax.default,
+        ]
         self.squeeze = [torch.ops.aten.squeeze.dims, torch.ops.aten.squeeze_copy.dims]
         self.to_copy = [
             torch.ops.aten._to_copy.default,
diff --git a/tico/passes/remove_redundant_reshape.py b/tico/passes/remove_redundant_reshape.py
index b1e401d..1ea3718 100644
--- a/tico/passes/remove_redundant_reshape.py
+++ b/tico/passes/remove_redundant_reshape.py
@@ -24,11 +24,12 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
-from tico.utils.utils import set_new_meta_val
+from tico.utils.utils import broadcastable, set_new_meta_val
 from tico.utils.validate_args_kwargs import (
     AddTensorArgs,
     PermuteArgs,
     ReshapeArgs,
+    SafeSoftmaxArgs,
     SoftmaxArgs,
 )
 
@@ -253,15 +254,14 @@ class RemoveRedundantReshapePattern3(PassBase):
                 continue
             if not softmax.target in ops.aten.softmax:
                 continue
-            softmax_args = SoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type]
-            add, softmax_dim, softmax_half_to_float = (
+            if softmax.target == torch.ops.aten._softmax.default:
+                softmax_args = SoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type, assignment]
+            elif softmax.target == torch.ops.aten._safe_softmax.default:
+                softmax_args = SafeSoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type, assignment]
+            add, softmax_dim = (
                 softmax_args.input,
                 softmax_args.dim,
-                softmax_args.half_to_float,
             )
-            assert isinstance(add, torch.fx.Node), type(add)
-            assert isinstance(softmax_dim, int), type(softmax_dim)
-            assert isinstance(softmax_half_to_float, bool), type(softmax_half_to_float)
             softmax_shape = extract_shape(softmax)
             # TODO support other dimension
             if softmax_dim != -1 and softmax_dim != len(softmax_shape) - 1:
@@ -295,10 +295,16 @@ class RemoveRedundantReshapePattern3(PassBase):
             # Check condition
             reshape_2_input_shape = extract_shape(reshape_2_input)
             reshape_3_input_shape = extract_shape(reshape_3_input)
-            if reshape_2_input_shape != reshape_3_input_shape:
+            if not broadcastable(reshape_2_input_shape, reshape_3_input_shape):
                 continue
             reshape_1_shape = extract_shape(reshape_1)
-            if reshape_2_input_shape != reshape_1_shape:
+            if (
+                reshape_2_input_shape != reshape_1_shape
+                and reshape_3_input_shape != reshape_1_shape
+            ):
+                continue
+            # Make sure the softmax axis length is unchanged.
+            if softmax_shape[-1] != reshape_1_shape[-1]:
                 continue
             # Assume `aten.add` and `aten.softmax` have only one user.
             if len(add.users) != 1:
@@ -311,8 +317,7 @@ class RemoveRedundantReshapePattern3(PassBase):
             set_new_meta_val(add)
             # Update softmax
             if softmax_dim == len(softmax_shape) - 1:
-                updated_dim = len(extract_shape(reshape_2_input)) - 1
-                softmax.args = (add, updated_dim, softmax_half_to_float)
+                softmax.update_arg(1, -1)  # (index, last_dim)
             set_new_meta_val(softmax)
 
             reshape_1.replace_all_uses_with(softmax, propagate_meta=False)
diff --git a/tico/utils/utils.py b/tico/utils/utils.py
index 1488cb9..53675e6 100644
--- a/tico/utils/utils.py
+++ b/tico/utils/utils.py
@@ -17,6 +17,7 @@ import subprocess
 import typing
 import warnings
 from functools import wraps
+from typing import List
 
 import torch
 from circle_schema import circle
@@ -341,3 +342,39 @@ def get_quant_dtype(qmin: int, qmax: int):
         return known_ranges[(qmin, qmax)]
     else:
         raise ValueError(f"Unsupported quantization range: ({qmin}, {qmax})")
+
+
+def broadcastable(
+    shape_a: List[int] | torch.Size, shape_b: List[int] | torch.Size
+) -> bool:
+    """
+    Return **True** if two shapes are broadcast-compatible under the standard
+    NumPy/PyTorch rules.
+
+    Broadcasting rule
+    --------------------------------
+    - Align the shapes **right-to-left**.
+    - For each aligned dimension `(a, b)` one of the following must hold
+      - `a == b`  (sizes match)
+      - `a == 1`  (shape-A can repeat along that dim)
+      - `b == 1`  (shape-B can repeat along that dim)
+    - When one shape is shorter, treat its missing leading dims as `1`.
+
+    Examples
+    --------
+    >>> _broadcastable([8, 16, 32], [16, 32])
+    True
+    >>> _broadcastable([8, 16, 32], [1, 32])
+    True
+    >>> _broadcastable([8, 16, 32], [8, 32, 16])
+    False
+    """
+    # Walk from the last dim to the front
+    len_a, len_b = len(shape_a), len(shape_b)
+    max_len = max(len_a, len_b)
+    for i in range(1, max_len + 1):
+        dim_a = shape_a[-i] if i <= len_a else 1
+        dim_b = shape_b[-i] if i <= len_b else 1
+        if dim_a != 1 and dim_b != 1 and dim_a != dim_b:
+            return False
+    return True


Test patch:
diff --git a/test/unit_test/pass_test/test_remove_redundant_reshape.py b/test/unit_test/pass_test/test_remove_redundant_reshape.py
index ad0476e..63fbbde 100644
--- a/test/unit_test/pass_test/test_remove_redundant_reshape.py
+++ b/test/unit_test/pass_test/test_remove_redundant_reshape.py
@@ -126,14 +126,45 @@ class RemoveRedundantReshapePattern3Test(SinglePassValueTest):
     def test_pass(self):
         self.setup(RedundantReshapePattern3())
 
-        if Version(torch.__version__) <= Version("2.6.0.dev20241015"):
-            self.run_value_test(ConvertLayoutOpToReshape())
+        from tico.utils.utils import SuppressWarning
 
+        # torch.ops.aten.softmax.int -> torch.ops.aten._softmax.default
+        with SuppressWarning(UserWarning, ".*quantize_per_tensor"):
+            # Warning details:
+            #   ...site-packages/torch/_subclasses/functional_tensor.py:364
+            #   UserWarning: At pre-dispatch tracing, we assume that any custom op marked with
+            #     CompositeImplicitAutograd and have functional schema are safe to not decompose.
+            self.ep = self.exported_program().run_decompositions()
+        self.run_value_test(ConvertLayoutOpToReshape())
         self.assertEqual(num_of_ops(self.exported_program(), ops.aten.reshape), 3)
 
-        for node in self.exported_program().graph.nodes:
-            if not node.op == "call_function":
-                continue
+        self.run_value_test(RemoveRedundantReshapePattern3())
+        self.assertEqual(num_of_ops(self.exported_program(), ops.aten.reshape), 0)
+
+
+class RedundantReshapePattern3Broadcasted(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        # AxBxC -> 1xAxBxC
+        x_reshape = torch.reshape(x, [1, *x.size()])
+        y_reshape = torch.reshape(y, [1, *y.size()])
+        # add
+        add = torch.add(x_reshape, y_reshape)
+        # softmax
+        softmax = torch.softmax(add, dim=3)
+        # 1xAxBxC -> AxBxC
+        reshape = torch.reshape(softmax, y.size())
+        return reshape
+
+    def get_example_inputs(self):
+        return (torch.randn(16, 1, 64), torch.randn(16, 64, 64))
+
+
+class RemoveRedundantReshapePattern3BroadcastedTest(SinglePassValueTest):
+    def test_pass(self):
+        self.setup(RedundantReshapePattern3Broadcasted())
 
         from tico.utils.utils import SuppressWarning
 
@@ -144,11 +175,51 @@ class RemoveRedundantReshapePattern3Test(SinglePassValueTest):
             #   UserWarning: At pre-dispatch tracing, we assume that any custom op marked with
             #     CompositeImplicitAutograd and have functional schema are safe to not decompose.
             self.ep = self.exported_program().run_decompositions()
+        self.run_value_test(ConvertLayoutOpToReshape())
+        self.assertEqual(num_of_ops(self.exported_program(), ops.aten.reshape), 3)
 
         self.run_value_test(RemoveRedundantReshapePattern3())
         self.assertEqual(num_of_ops(self.exported_program(), ops.aten.reshape), 0)
 
 
+class RedundantReshapePattern3DifferentSoftmaxLength(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        x_reshape = torch.reshape(x, [1, 2, 3])
+        y_reshape = torch.reshape(y, [1, 2, 3])
+        # add
+        add = torch.add(x_reshape, y_reshape)
+        # softmax
+        softmax = torch.softmax(add, dim=-1)
+        reshape = torch.reshape(softmax, y.size())
+        return reshape
+
+    def get_example_inputs(self):
+        return (torch.randn(1, 1, 6), torch.randn(1, 1, 6))
+
+
+class RemoveRedundantReshapePattern3DifferentSoftmaxLengthTest(SinglePassValueTest):
+    def test_pass(self):
+        self.setup(RedundantReshapePattern3DifferentSoftmaxLength())
+
+        from tico.utils.utils import SuppressWarning
+
+        # torch.ops.aten.softmax.int -> torch.ops.aten._softmax.default
+        with SuppressWarning(UserWarning, ".*quantize_per_tensor"):
+            # Warning details:
+            #   ...site-packages/torch/_subclasses/functional_tensor.py:364
+            #   UserWarning: At pre-dispatch tracing, we assume that any custom op marked with
+            #     CompositeImplicitAutograd and have functional schema are safe to not decompose.
+            self.ep = self.exported_program().run_decompositions()
+        self.run_value_test(ConvertLayoutOpToReshape())
+        self.assertEqual(num_of_ops(self.exported_program(), ops.aten.reshape), 3)
+
+        self.run_value_test(RemoveRedundantReshapePattern3())
+        self.assertEqual(num_of_ops(self.exported_program(), ops.aten.reshape), 3)
+
+
 class RedundantReshapePattern4(torch.nn.Module):
     def __init__(self):
         super().__init__()
diff --git a/test/unit_test/utils_test/test_utils.py b/test/unit_test/utils_test/test_utils.py
index 2a28d90..b459803 100644
--- a/test/unit_test/utils_test/test_utils.py
+++ b/test/unit_test/utils_test/test_utils.py
@@ -14,7 +14,7 @@
 
 import unittest
 
-from tico.utils.utils import get_quant_dtype
+from tico.utils.utils import broadcastable, get_quant_dtype
 
 
 class TestGetQuantDtype(unittest.TestCase):
@@ -35,3 +35,38 @@ class TestGetQuantDtype(unittest.TestCase):
             get_quant_dtype(256, 512)
         with self.assertRaises(ValueError):
             get_quant_dtype(-32768, 32768)
+
+
+class BroadcastableTest(unittest.TestCase):
+    def test_true_cases(self):
+        true_pairs = [
+            # identical
+            ([4, 8, 16], [4, 8, 16]),
+            # trailing singleton
+            ([4, 8, 16], [4, 8, 1]),
+            # head & tail singleton
+            ([4, 8, 1], [1, 8, 1]),
+            # rank-mismatch, typical bias broadcast
+            ([32, 64, 64], [32, 1, 64]),
+            # rank-mismatch, leading dims implicit
+            ([8, 16, 32], [16, 32]),  # (1, 16, 32)
+            ([8, 16, 32], [1, 32]),  # (1, 1, 32)
+            ([8, 16, 32], [32]),  # (1, 1, 32)
+        ]
+        for a, b in true_pairs:
+            with self.subTest(shape_a=a, shape_b=b):
+                self.assertTrue(broadcastable(a, b))
+
+    def test_false_cases(self):
+        false_pairs = [
+            # mismatch in non-singleton dim
+            ([4, 8, 16], [4, 7, 16]),
+            # trailing dim mismatch
+            ([4, 8, 16], [4, 8, 32]),
+            # rank-mismatch but incompatible
+            ([4, 8, 16], [4, 4]),  # align (8 vs 4)
+            ([4, 8, 16, 2], [4, 8, 16]),  # align first diff 2 vs 1
+        ]
+        for a, b in false_pairs:
+            with self.subTest(shape_a=a, shape_b=b):
+                self.assertFalse(broadcastable(a, b))