[PERF] [llama8B] Add subgraph rewrite rule to fuse Matmul and cast (#1021)

Unfused Graph:

![image](https://github.com/user-attachments/assets/766e5b8d-3882-400d-a398-e513046e5a63)

Original Behaviour:

![image](https://github.com/user-attachments/assets/67c52eed-145b-4cdf-afb3-6ffa492a29d5)

New Behaviour:

![image](https://github.com/user-attachments/assets/8a3b0588-9548-45f2-a0e9-8329a39c8533)

Originally, the CastOp was being left out of the fusion because after a
epilogue fusion pass, everything after matmul would be fused and CastOp
would be left as a prologue operator. Because prologue fusion is
disabled on Matmul, CastOp would never be fused.

Fixed this issue by adding a MultiInputCompositeElementWiseOp which
combines 2 UnaryOps with different inputs and 1 BinaryElementWiseOp,
allowing this shape to be fused with Matmul.

Author: PetersonGuo

python/hidet/graph/ops/arithmetic.py

@@ -135,6 +135,53 @@ def __init__(self, name: str, args: List[TensorNode], op: Callable[[Any], Any]):
         super().__init__(name=name, inputs=list(args), outputs=[out], inverse_map=inverse_map)
 
 
+class MultiInputCompositeElementwiseTask(Task):
+    def __init__(
+        self,
+        name: str,
+        x1: TensorNode,
+        x2: TensorNode,
+        left_unary_op: Callable[[Any], Any],
+        right_unary_op: Callable[[Any], Any],
+        binary_op: Callable[[Any, Any], Any],
+        attrs=None,
+    ):
+        def composite_op(binary_op, left_unary_op, right_unary_op, x1, x2):
+            if left_unary_op is None:
+                left_unary_op = lambda x1: x1
+            if right_unary_op is None:
+                right_unary_op = lambda x2: x2
+            return binary_op(left_unary_op(x1), right_unary_op(x2))
+
+        z_shape = broadcast_shape(x1.shape, x2.shape)
+        z = compute(
+            name='z',
+            shape=z_shape,
+            fcompute=lambda *indices: composite_op(
+                binary_op, left_unary_op, right_unary_op, x1.__getitem__(indices), x2.__getitem__(indices)
+            ),
+        )
+
+        inverse_map = {}
+        for inp, inp_shape in zip([x1, x2], [x1.shape, x2.shape]):
+            if same_list(inp_shape, z_shape):
+                inverse_map[inp] = InverseMap.from_lambda(lambda *indices: indices, num_args=len(inp_shape))
+            elif prod(inp_shape) == prod(z_shape):
+                inverse_map[inp] = InverseMap.from_lambda(
+                    lambda *indices: [0 for _ in range(len(z_shape) - len(inp_shape))] + list(indices),
+                    num_args=len(inp_shape),
+                )
+        # layout := 1:0 means identity layout
+        # It's obvious that arithmetic operations don't apply layout
+        # modification on tensors, so the tile mapping function can be
+        # considered as identity
+        for _, v in inverse_map.items():
+            v.tile_mapping = TensorLayout(1)
+        super().__init__(
+            name=name, inputs=[x1, x2], outputs=[z], inverse_map=inverse_map, attributes={} if attrs is None else attrs
+        )
+
+
 class CompositeElementwiseTask(Task):
     def __init__(
         self,
@@ -354,6 +401,29 @@ def get_dtype(scalar: Expr):
     return inferred_type
 
 
+class MultiInputCompositeElementwiseOp(Operator):
+    def __init__(
+        self,
+        x1: Tensor,
+        x2: Tensor,
+        left_unary_op: UnaryElementwiseOperation,
+        right_unary_op: UnaryElementwiseOperation,
+        binary_op: BinaryElementwiseOperation,
+    ):
+        name = 'multi_input_composite'
+        for op in [left_unary_op, right_unary_op, binary_op]:
+            if op is not None:
+                name += '_' + op.name
+        attributes = {'left_unary_op': left_unary_op, 'right_unary_op': right_unary_op, 'binary_op': binary_op}
+        super().__init__(
+            inputs=[x1, x2],
+            attributes=attributes,
+            task=MultiInputCompositeElementwiseTask(
+                name, input_like(x1, 'x1'), input_like(x2, 'x2'), left_unary_op, right_unary_op, binary_op
+            ),
+        )
+
+
 class CompositeElementwiseOp(Operator):
     def __init__(
         self,
@@ -1202,6 +1272,16 @@ def roll(x: Tensor, shifts: Union[int, Sequence[int]], dims: Union[int, Sequence
     return RollOp(x, shifts, dims).outputs[0]
 
 
+def composite_multi_input_elementwise(
+    x1: Tensor,
+    x2: Tensor,
+    left_unary_op: UnaryElementwiseOperation,
+    right_unary_op: UnaryElementwiseOperation,
+    binary_op: BinaryElementwiseOperation,
+) -> Tensor:
+    return MultiInputCompositeElementwiseOp(x1, x2, left_unary_op, right_unary_op, binary_op).outputs[0]
+
+
 # out = binary_op(left_unary_op(x), right_unary_op(x)); This allows more fusion opportunity.
 def composite_elementwise(
     x: Tensor,

python/hidet/graph/ops/transform.py

@@ -10,6 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import List, Optional, Tuple, Union, Sequence
+from hidet.graph.ops.arithmetic import UnaryElementwiseOp
 from hidet.ir.type import DataType, data_type
 from hidet.ir.expr import Expr, Constant, if_then_else, convert, cast as ir_cast, is_constant, logical_or
 from hidet.ir.expr import Int
@@ -697,15 +698,9 @@ def run_torch(self):
         return result
 
 
-class CastOp(Operator):
+class CastOp(UnaryElementwiseOp):
     def __init__(self, x: Tensor, dtype: DataType):
-        from .arithmetic import UnaryElementwiseTask
-
-        super().__init__(
-            inputs=[x],
-            attributes={'dtype': dtype},
-            task=UnaryElementwiseTask('cast', input_like(x, 'x'), op=lambda v: ir_cast(v, dtype)),
-        )
+        super().__init__(x, op=lambda v: ir_cast(v, dtype), name='cast', attributes={'dtype': dtype})
 
     def run_torch(self):
         x = self.inputs[0]

python/hidet/graph/transforms/graph_patterns/transform_patterns.py

@@ -171,6 +171,29 @@ def target(self, matched: MatchDict) -> Optional[List[Tensor]]:
         return [x]
 
 
+class MultiInputCompositeElementwiseLeftRightRewriteRule(SubgraphRewriteRule):
+    def __init__(self):
+        super().__init__('binaryOp(unaryOp(x), unaryOp(y)) => compositeOp(x, y)')
+        self.x1 = TensorPattern()
+        self.x2 = TensorPattern()
+        self.y1 = op_pattern(UnaryElementwiseOp, [self.x1])
+        self.y2 = op_pattern(UnaryElementwiseOp, [self.x2])
+        self.z = op_pattern(BinaryElementwiseOp, [self.y1, self.y2])
+
+    def source(self) -> List[TensorPattern]:
+        return [self.z]
+
+    def target(self, matched: MatchDict) -> Optional[List[Tensor]]:
+        x1, x2, y1, y2, z = [matched[v] for v in [self.x1, self.x2, self.y1, self.y2, self.z]]
+        left_unary_op: UnaryElementwiseOperation = y1.op.op
+        right_unary_op: UnaryElementwiseOperation = y2.op.op
+        if left_unary_op.name != right_unary_op.name:
+            return None
+        binary_op: BinaryElementwiseOperation = z.op.op
+        out = ops.arithmetic.composite_multi_input_elementwise(x1, x2, left_unary_op, right_unary_op, binary_op)
+        return [out]
+
+
 class CompositeElementwiseLeftRightRewriteRule(SubgraphRewriteRule):
     def __init__(self):
         super().__init__('binaryOp(unaryOp_left(x), unaryOp_right(x)) => compositeOp(x)')
@@ -235,6 +258,7 @@ def transform_patterns():
     register_rewrite_rule(FanoutTwoCast())
     register_rewrite_rule(FanoutThreeCast())
     register_rewrite_rule(DoubleCast())
+    register_rewrite_rule(MultiInputCompositeElementwiseLeftRightRewriteRule())
     register_rewrite_rule(CompositeElementwiseLeftRightRewriteRule())
     register_rewrite_rule(CompositeElementwiseLeftRewriteRule())
     register_rewrite_rule(CompositeElementwiseRightRewriteRule())