Qualcomm AI Engine Direct - Support Gemma RMSNorm recomposition

backends/qualcomm/_passes/remove_redundancy.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 import torch
+from executorch.backends.qualcomm.builders.utils import 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
@@ -14,8 +15,9 @@ class RemoveRedundancy(ExportPass):
     Trim certain operators to reduce unnecessary overhead.
     """
 
-    def __init__(self, quantization_capture=False):
+    def __init__(self, quantization_capture=False, edge_program=None):
         super(RemoveRedundancy, self).__init__()
+        self.edge_program = edge_program
         self.redundant_ops_general = {
             torch.clone: self._default_condition,
             torch.ops.aten.clone.default: self._default_condition,
@@ -26,6 +28,7 @@ def __init__(self, quantization_capture=False):
             exir_ops.edge.aten.expand_copy.default: self._same_shape_condition,
             exir_ops.edge.aten.lift_fresh_copy.default: self._default_condition,
             exir_ops.edge.aten.repeat.default: self._same_shape_condition,
+            exir_ops.edge.aten.view_copy.default: self._rms_norm_weight_view_condition,
             # remove this target if '_skip_dim_order' is set to False
             exir_ops.edge.dim_order_ops._to_dim_order_copy.default: self._dim_order_op_condition,
             # remove channel_last / contiguous _to_copy if '_skip_dim_order' is set to True
@@ -48,6 +51,22 @@ def _dim_order_op_condition(self, node):
     def _same_shape_condition(self, node):
         return node.args[0].meta["val"].shape == node.meta["val"].shape
 
+    def _rms_norm_weight_view_condition(self, node):
+        source = node.args[0]
+        is_static = source.op == "get_attr" or (
+            self.edge_program is not None and is_parameter(source, self.edge_program)
+        )
+        if not is_static or not self._same_shape_condition(node):
+            return False
+        return all(
+            user.target == exir_ops.edge.aten.rms_norm.default
+            and (
+                (len(user.args) > 2 and user.args[2] is node)
+                or user.kwargs.get("weight") is node
+            )
+            for user in node.users
+        )
+
     def _to_copy_op_condition(self, node):
         return "memory_format" in node.kwargs
 

backends/qualcomm/builders/op_rms_norm.py

@@ -19,7 +19,7 @@
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 
-from .node_visitor import NodeVisitor
+from .node_visitor import NodeVisitor, q_ops
 from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpRmsNorm, QNN_OP_PACKAGE_NAME_QTI_AISW
 
@@ -51,7 +51,7 @@ def define_node(
         normalized_shapes = node.args[1]
         if (
             len(normalized_shapes) != 1
-            and normalized_shapes[0] != input_tensor.shape[-1]
+            or normalized_shapes[0] != input_tensor.shape[-1]
         ):
             warnings.warn(
                 "[QNN Delegate Op Builder]: Only supports normalization with last input dimension.",
@@ -61,9 +61,10 @@ def define_node(
         axes = [node.args[0].meta["val"].dim() - 1]
         axes_shape = [len(axes)]
 
-        has_weight = len(node.args) > 2 and node.args[2] is not None
+        weight = node.args[2] if len(node.args) > 2 else node.kwargs.get("weight")
+        has_weight = weight is not None
         if has_weight:
-            weight_node = self.get_node(node.args[2])
+            weight_node = self.get_node(weight)
             weight_tensor = get_parameter(weight_node, self.edge_program)
         else:
             # elementwise_affine=False: use all-ones weight as identity
@@ -110,9 +111,17 @@ def define_node(
             nodes_to_wrappers,
         )
 
-        epsilon = torch.finfo(torch.float32).eps
-        if len(node.args) > 3:
-            epsilon = node.args[3]
+        eps = node.args[3] if len(node.args) > 3 else node.kwargs.get("eps")
+        if eps is not None:
+            epsilon = eps
+        else:
+            logical_input = node.args[0]
+            input_dtype = logical_input.meta["val"].dtype
+            if logical_input.target in q_ops:
+                input_dtype = logical_input.args[0].meta["val"].dtype
+            if not input_dtype.is_floating_point:
+                input_dtype = node.meta["val"].dtype
+            epsilon = torch.finfo(input_dtype).eps
         output_tensor = self.get_tensor(node, node)
         output_tensor_wrapper = self.define_tensor(
             node,

backends/qualcomm/tests/models.py

@@ -2420,6 +2420,18 @@ def forward(self, x):
         return self.rms(x)
 
 
+class Fp32RmsNorm(torch.nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        input_dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6)
+        return x.to(input_dtype) * self.weight.to(input_dtype)
+
+
 class Roll(torch.nn.Module):
     def __init__(self, shifts, dims=None):
         super().__init__()

backends/qualcomm/tests/rework/passes/test.py

@@ -398,6 +398,12 @@ def test_recompose_rms_norm(request, kwargs):
     RecomposeRmsNorm.test(request, kwargs)  # noqa: F405
 
 
+@enumerate_backends_quantized()
+@repack_pass_fixtures
+def test_recompose_rms_norm_annotation(request, kwargs):
+    RecomposeRmsNorm.test_annotation(request, kwargs)  # noqa: F405
+
+
 @enumerate_backends()
 @repack_pass_fixtures
 def test_remove_0d_tensor(request, kwargs):