[Qualcomm] Support native_layer_norm and affine-free LayerNorm in QNN backend

- add QNN layer norm support for aten.native_layer_norm.default
- handle missing weight/bias by creating identity weight and zero bias
- always provide bias tensor for QNN LayerNorm op
- add floating-point and quantized tests for native_layer_norm
- print generated pte filename after export

Author: KevinUW114514

backends/qualcomm/builders/op_layer_norm.py

@@ -11,7 +11,12 @@
 
 import numpy as np
 import torch
-from executorch.backends.qualcomm.utils.constants import QCOM_DATA
+from executorch.backends.qualcomm.utils.constants import (
+    QCOM_DATA,
+    QCOM_QUANT_ATTRS,
+    QCOM_ZERO_POINT,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
 
 from .node_visitor import NodeVisitor
 from .node_visitor_manager import register_node_visitor
@@ -31,6 +36,7 @@ def define_node(
         node: torch.fx.Node,
         nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
     ) -> PyQnnManager.PyQnnOpWrapper:
+        # args of node : ['input', 'normalized_shape', 'weight', 'bias', 'eps']
         input_node = self.get_node(node.args[0])
         input_tensor = self.get_tensor(input_node, node)
         input_tensor_wrapper = self.define_tensor(
@@ -54,37 +60,61 @@ def define_node(
         axis = [len(input_tensor.shape) - 1]
         axis_shape = [len(axis)]
 
-        weight_node = self.get_node(node.args[2])
-        if weight_node is not None:
+        has_weight = len(node.args) > 2 and node.args[2] is not None
+        if has_weight:
+            weight_node = self.get_node(node.args[2])
             weight_tensor = get_parameter(weight_node, self.edge_program)
-            weight_tensor_wrapper = self.define_tensor(
-                weight_node,
-                node,
-                weight_tensor,
-                PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
-                nodes_to_wrappers,
-            )
-            layer_norm_input_tensors = [input_tensor_wrapper, weight_tensor_wrapper]
         else:
-            warnings.warn(
-                "[QNN Delegate Op Builder]: LayerNorm weight is None, skipping",
-                stacklevel=1,
+            # elementwise_affine=False: use all-ones weight as identity
+            weight_tensor = torch.ones(normalized_shapes, dtype=torch.float32)
+            weight_node = torch.fx.Node(
+                node.graph,
+                node.name + "_runtime_weight",
+                "call_function",
+                exir_ops.edge.aten.tensor.default,
+                (),
+                {},
             )
-            layer_norm_input_tensors = [input_tensor_wrapper]
+            if quant_attrs := node.meta.get(QCOM_QUANT_ATTRS):
+                quant_attrs = quant_attrs.copy()
+                quant_attrs[QCOM_ZERO_POINT] = 0
+                weight_node.meta[QCOM_QUANT_ATTRS] = quant_attrs
+        weight_tensor_wrapper = self.define_tensor(
+            weight_node,
+            node,
+            weight_tensor,
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
+            nodes_to_wrappers,
+        )
 
-        bias_node = self.get_node(node.args[3])
-        if bias_node is not None:
+        # Fake node: even when original bias is absent, QNN still needs it
+        has_bias = len(node.args) > 3 and node.args[3] is not None
+        if has_bias:
+            bias_node = self.get_node(node.args[3])
             bias_tensor = get_parameter(bias_node, self.edge_program)
-            bias_tensor_wrapper = self.define_tensor(
-                bias_node,
-                node,
-                bias_tensor,
-                PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
-                nodes_to_wrappers,
+        else:
+            bias_tensor = torch.zeros(normalized_shapes, dtype=torch.float32)
+            bias_node = torch.fx.Node(
+                node.graph,
+                node.name + "_runtime_bias",
+                "call_function",
+                exir_ops.edge.aten.tensor.default,
+                (),
+                {},
             )
-            layer_norm_input_tensors.append(bias_tensor_wrapper)
+            if quant_attrs := node.meta.get(QCOM_QUANT_ATTRS):
+                quant_attrs = quant_attrs.copy()
+                quant_attrs[QCOM_ZERO_POINT] = 0
+                bias_node.meta[QCOM_QUANT_ATTRS] = quant_attrs
+        bias_tensor_wrapper = self.define_tensor(
+            bias_node,
+            node,
+            bias_tensor,
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
+            nodes_to_wrappers,
+        )
 
-        epsilon = node.args[4]
+        epsilon = node.args[4] if len(node.args) > 4 else 1e-05
 
         output_tensor = self.get_tensor(node, node, 0)
         output_tensor_wrapper = self.define_tensor(
@@ -100,7 +130,9 @@ def define_node(
             QNN_OP_PACKAGE_NAME_QTI_AISW,
             OpLayerNorm.op_name,
         )
-        layer_norm_op.AddInputTensors(layer_norm_input_tensors)
+        layer_norm_op.AddInputTensors(
+            [input_tensor_wrapper, weight_tensor_wrapper, bias_tensor_wrapper]
+        )
         layer_norm_op.AddOutputTensors([output_tensor_wrapper])
         layer_norm_op.AddScalarParam(
             OpLayerNorm.param_epsilon,

backends/qualcomm/export_utils.py

@@ -617,6 +617,7 @@ def build_executorch_binary(
     with open(pte_name, "wb") as file:
         exec_prog_mgr.write_to_file(file)
 
+    print(f"Successfully generated {pte_name}.")
     if qnn_config.compile_only:
         sys.exit(0)
 

backends/qualcomm/quantizer/annotators/htp_rules.py

@@ -828,7 +828,8 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
 
 
 @register_annotator(
-    [torch.ops.aten.layer_norm.default], QnnConstants.OpLayerNorm.op_name
+    [torch.ops.aten.layer_norm.default, torch.ops.aten.native_layer_norm.default],
+    QnnConstants.OpLayerNorm.op_name,
 )
 class LayerNorm(GeneralOpDef):
     @staticmethod

backends/qualcomm/tests/models.py

@@ -1388,6 +1388,26 @@ def forward(self, x):
         return self.linear(self.layer_norm(x))
 
 
+class NativeLayerNorm(torch.nn.Module):
+    def __init__(self, affine=True):
+        super().__init__()
+        self.affine = affine
+        self.weight = torch.nn.Parameter(torch.ones(768))
+        self.bias = torch.nn.Parameter(torch.zeros(768))
+        self.normalized_shape = [768]
+        self.eps = 1e-6
+
+    def forward(self, x):
+        if self.affine:
+            return torch.native_layer_norm(
+                x, self.normalized_shape, self.weight, self.bias, self.eps
+            )[0]
+        else:
+            return torch.native_layer_norm(
+                x, self.normalized_shape, self.weight, self.bias, self.eps
+            )[0]
+
+
 class LayerNormAdd(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1384,6 +1384,13 @@ def test_qnn_backend_layer_norm(self):
             with self.subTest(i=i):
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_native_layer_norm(self):
+        modules = [NativeLayerNorm(), NativeLayerNorm(affine=False)]  # noqa: F405
+        sample_input = (torch.randn(196, 768),)
+        for i, module in enumerate(modules):
+            with self.subTest(i=i):
+                self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_leaky_relu(self):
         torch.manual_seed(8)
         test_comb = [
@@ -3811,6 +3818,14 @@ def test_qnn_backend_layer_norm(self):
                 module = self.get_qdq_module(module, sample_input)
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_native_layer_norm(self):
+        modules = [NativeLayerNorm(), NativeLayerNorm(affine=False)]  # noqa: F405
+        sample_input = (torch.randn(196, 768),)
+        for i, module in enumerate(modules):
+            with self.subTest(i=i):
+                module = self.get_qdq_module(module, sample_input)
+                self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_leaky_relu(self):
         test_comb = [
             {