Qualcomm AI Engine Direct - Keep the dtype of a recomposed RMSNorm

backends/qualcomm/_passes/recompose_rms_norm.py

@@ -5,7 +5,14 @@
 # LICENSE file in the root directory of this source tree.
 import torch
 from executorch.backends.qualcomm._passes.utils import find_patterns
-from executorch.backends.qualcomm.builders.node_visitor import dq_ops
+from executorch.backends.qualcomm.builders.node_visitor import (
+    dq_ops,
+    QNN_TENSOR_TYPE_MAP,
+)
+from executorch.backends.qualcomm.builders.utils import (
+    is_mutable_buffer_input,
+    is_parameter,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 from executorch.exir.passes import dead_code_elimination_pass
@@ -74,23 +81,133 @@ class RecomposeRmsNorm(ExportPass):
     Merge decomposed operators back to one super node.
     """
 
-    def __init__(self, quantization_capture=False):
+    def __init__(self, quantization_capture=False, edge_program=None):
         super(RecomposeRmsNorm, self).__init__()
+        self.edge_program = edge_program
         self.rms_norm_target = exir_ops.edge.aten.rms_norm.default
         self.skip_targets = [
             exir_ops.edge.aten.to.dtype,
+            exir_ops.edge.aten._to_copy.default,
+            exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
         ]
         self.quantization_capture = quantization_capture
         if quantization_capture:
             self.rms_norm_target = torch.ops.aten.rms_norm.default
             self.skip_targets = [
                 torch.ops.aten.to.dtype,
+                torch.ops.aten.to.device,
+                torch.ops.aten.type_as.default,
             ]
 
-    def _get_input_node(self, node):
-        input_node = node.args[0]
-        while input_node.target in self.skip_targets:
+    def _is_cast(self, node, dtype):
+        """
+        Whether the node is a cast that changes nothing but the dtype, to a
+        dtype at least as wide as `dtype`. A narrower cast rounds the values,
+        which the rebuilt rms_norm in `dtype` would not do.
+        """
+        if not _is_call(node) or node.target not in self.skip_targets:
+            return False
+        source, cast = node.args[0].meta["val"], node.meta["val"]
+        return torch.promote_types(cast.dtype, dtype) == cast.dtype and (
+            cast.layout,
+            cast.device,
+            cast.dim_order(),
+        ) == (source.layout, source.device, source.dim_order())
+
+    def _strip_casts(self, node, dtype):
+        while self._is_cast(node, dtype):
+            node = node.args[0]
+        return node
+
+    def _squares_the_scaled_tensor(self, pattern, dtype):
+        square, scale_mul = pattern[6], pattern[2]
+        base = square.args[0]
+        if _is_pow(square):
+            exponent = square.args[1]
+            if _is_node(exponent):
+                # LiftConstantScalarOperands may have turned it into a constant.
+                constant = exponent.meta["val"].constant
+                exponent = None if constant is None else constant.item()
+            squares = exponent == 2
+        else:
+            squares = self._strip_casts(square.args[1], dtype) is self._strip_casts(
+                base, dtype
+            )
+        scaled = [arg for arg in scale_mul.args if arg is not pattern[3]]
+        # T5's norm squares a float32 copy of the tensor it scales.
+        return (
+            squares
+            and len(scaled) == 1
+            and self._strip_casts(scaled[0], dtype) is self._strip_casts(base, dtype)
+        )
+
+    def _get_input(self, pattern, dtype):
+        """
+        The tensor to rebuild a norm in `dtype` on, or None if there is none.
+        rms_norm returns its input's dtype, so the input has to be in `dtype`.
+        Unless the norm is about to be quantized, QNN also needs a type for it.
+        """
+        if not self._squares_the_scaled_tensor(pattern, dtype):
+            return None
+        input_node = pattern[6].args[0]
+        while self._is_cast(input_node, dtype) and (
+            input_node.meta["val"].dtype != dtype
+            or input_node.args[0].meta["val"].dtype == dtype
+        ):
             input_node = input_node.args[0]
+        if input_node.meta["val"].dtype != dtype or (
+            not self.quantization_capture and dtype not in QNN_TENSOR_TYPE_MAP
+        ):
+            return None
+        return input_node
+
+    def _get_normalization_input(self, pattern, dtype):
+        """
+        The tensor to rebuild only the normalization on, or None. It has to
+        read something at runtime: the pass cannot see the partitions, HTP
+        drops one that reads only constants, and rms_norm has no CPU kernel
+        to fall back on.
+        """
+        input_node = self._get_input(pattern, dtype)
+        if input_node is None or not self._reads_runtime_input(input_node):
+            return None
+        return input_node
+
+    def _reads_runtime_input(self, node):
+        # Placeholders count as QNN's partitioner counts them. Without the
+        # program, every placeholder counts.
+        nodes, seen = [node], set()
+        while nodes:
+            node = nodes.pop()
+            if node in seen:
+                continue
+            seen.add(node)
+            if node.op == "placeholder" and (
+                self.edge_program is None
+                or not is_parameter(node, self.edge_program)
+                or is_mutable_buffer_input(node, self.edge_program)
+            ):
+                return True
+            nodes.extend(node.all_input_nodes)
+        return False
+
+    def _get_norm_input(self, pattern, gamma_node):
+        last_mul_node = pattern[0]
+        dtype = last_mul_node.meta["val"].dtype
+        # The "*" in the patterns can match any op, but only casts may sit
+        # between the normalization and the weight.
+        normalized = next(
+            (arg for arg in last_mul_node.args if arg is not gamma_node), None
+        )
+        if self._strip_casts(normalized, dtype) is not pattern[2]:
+            return None
+        input_node = self._get_input(pattern, dtype)
+        if (
+            input_node is None
+            or list(gamma_node.meta["val"].shape)
+            != list(input_node.meta["val"].shape)[-1:]
+        ):
+            return None
         return input_node
 
     def call(self, graph_module: torch.fx.GraphModule):
@@ -124,30 +241,51 @@ def call(self, graph_module: torch.fx.GraphModule):
                     if _is_get_attr(arg) or _is_placeholder(arg) or _is_dq(arg):
                         gamma_node = arg
 
-                if gamma_node is None:
+                input_node = (
+                    self._get_norm_input(rms_norm_pattern, gamma_node)
+                    if gamma_node is not None
+                    else None
+                )
+                dtype = rms_norm_pattern[2].meta["val"].dtype
+                if (
+                    input_node is None
+                    and not self.quantization_capture
+                    and dtype == torch.float32
+                ):
+                    # Rebuild only the normalization and keep the casts and the
+                    # weight after it. rms_norm computes in float32, so this is
+                    # exact. Left decomposed, QNN's float16 mode would square in
+                    # float16, which overflows above 65504.
+                    output_node, gamma_node = rms_norm_pattern[2], None
+                    input_node = self._get_normalization_input(rms_norm_pattern, dtype)
+                else:
+                    output_node = last_mul_node
+                if input_node is None:
                     continue
-
-                input_node = self._get_input_node(rms_norm_pattern[6])
-                # Only the last-dim form the RmsNorm builder supports is recomposed.
-                # Anything else, e.g. a 0-d scale, is left decomposed.
+                # Only the form QNN's RmsNorm takes is recomposed: the last of
+                # 1 to 4 dims, all with static sizes. Anything else is left
+                # decomposed, since rms_norm has no CPU kernel to fall back on.
                 input_shape = list(input_node.meta["val"].shape)
                 rank = len(input_shape)
                 mean_node = rms_norm_pattern[5]
                 mean_dims = mean_node.args[1] or range(rank)
                 keepdim = len(mean_node.args) > 2 and mean_node.args[2]
                 if (
-                    rank == 0
+                    not 0 < rank <= 4
+                    or not all(isinstance(dim, int) for dim in input_shape)
                     or [dim % rank for dim in mean_dims] != [rank - 1]
                     or not (keepdim or rank == 1)
-                    or list(gamma_node.meta["val"].shape) != input_shape[-1:]
                 ):
                     continue
 
                 eps = rms_norm_pattern[4].args[1]
                 if isinstance(eps, torch.fx.Node):
-                    eps = eps.meta["val"].constant.item()
+                    eps = eps.meta["val"].constant
+                    if eps is None:
+                        continue
+                    eps = eps.item()
 
-                with graph.inserting_before(last_mul_node):
+                with graph.inserting_before(output_node):
                     # args schema
                     # (Tensor input, int[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor
                     rms_node = graph.create_node(
@@ -160,11 +298,11 @@ def call(self, graph_module: torch.fx.GraphModule):
                             eps,
                         ),
                     )
-                    users = last_mul_node.users.copy()
+                    users = output_node.users.copy()
                     for user in users:
-                        user.replace_input_with(last_mul_node, rms_node)
+                        user.replace_input_with(output_node, rms_node)
                     # copy metadata
-                    rms_node.meta = last_mul_node.meta
+                    rms_node.meta = output_node.meta
 
         dead_code_elimination_pass(graph_module)
         return PassResult(graph_module, True)

backends/qualcomm/_passes/remove_redundancy.py

@@ -49,7 +49,7 @@ def _same_shape_condition(self, node):
         return node.args[0].meta["val"].shape == node.meta["val"].shape
 
     def _to_copy_op_condition(self, node):
-        return "memory_format" in node.kwargs
+        return "memory_format" in node.kwargs and self._dim_order_op_condition(node)
 
     def _default_condition(self, ndoe):
         return True