Fix based on comments

- Change to string based way to set up qconfig for submodule

Author: Chun-I Tsai

backends/qualcomm/_passes/utils.py

@@ -11,7 +11,6 @@
 from executorch.backends.qualcomm.utils.constants import (
     QCOM_DTYPE,
     QCOM_ENCODING,
-    QCOM_NN_MODULE_STACK,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch._subclasses import FakeTensor
@@ -129,8 +128,8 @@ def copy_nn_module_stack(src, target):
     """
     Copy meta["nn_module_stack"] from src node to target node if existing.
     """
-    if value := src.meta.get(QCOM_NN_MODULE_STACK):
-        target.meta[QCOM_NN_MODULE_STACK] = value
+    if value := src.meta.get("nn_module_stack"):
+        target.meta["nn_module_stack"] = value
 
 
 def is_float_tensor(node: torch.fx.Node) -> bool:

backends/qualcomm/quantizer/quantizer.py

@@ -3,11 +3,10 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-import importlib
 from dataclasses import dataclass
 from enum import IntEnum, unique
 from functools import partial
-from typing import Callable, Dict, Optional, Sequence, Set, Tuple
+from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple
 
 import torch
 from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager
@@ -140,9 +139,11 @@ def __post_init__(self):
             raise RuntimeError(
                 f"the quant config, (quant_dtype: {self.quant_dtype}, is_qat: {self.is_qat}) is not support"
             )
-        quant_config_func, per_channel_quant_config_func, per_block_quant_config_func = QUANT_CONFIG_DICT[
-            (self.quant_dtype, self.is_qat)
-        ]
+        (
+            quant_config_func,
+            per_channel_quant_config_func,
+            per_block_quant_config_func,
+        ) = QUANT_CONFIG_DICT[(self.quant_dtype, self.is_qat)]
         self.quant_config = (
             quant_config_func(act_observer=self.act_observer)
             if self.act_observer
@@ -184,7 +185,9 @@ def __init__(self):
         self.quant_ops: Set[OpOverload] = self.SUPPORTED_OPS.copy()
 
         self.default_quant_config = ModuleQConfig()
-        self.module_qconfig_dict: Dict[torch.nn.Module, ModuleQConfig] = {}
+        self.submodule_qconfig_list: List[
+            Tuple[Callable[[torch.fx.Node], bool], ModuleQConfig]
+        ] = []
         self.block_size_map = {}
 
         self.custom_quant_annotations: Sequence[Callable] = []
@@ -203,44 +206,30 @@ def _annotate_custom_annotation(self, gm: GraphModule) -> None:
         for annotation_func in self.custom_quant_annotations:
             annotation_func(gm)
 
-    def _get_submodule(self, node: torch.fx.Node):
-        """
-        An example of nn_module_stack
-        {
-            'L__self__': ('', 'executorch.backends.qualcomm.tests.models.SubModules'),
-            'L__self___add': ('add', 'executorch.backends.qualcomm.tests.models.Add')
-        }
-        """
-
-        nn_module_stack = node.meta.get("nn_module_stack")
-        if nn_module_stack:
-            module_source_str, module_str = list(nn_module_stack.values())[-1][
-                -1
-            ].rsplit(".", 1)
-            module_source = importlib.import_module(module_source_str)
-            return getattr(module_source, module_str)
-        return None
+    def _get_submodule_qconfig(self, node: torch.fx.Node):
+        for func, qconfig in self.submodule_qconfig_list:
+            if func(node):
+                return qconfig
+        return self.default_quant_config
 
     def _get_quant_config(self, node: torch.fx.Node) -> Optional[QuantizationConfig]:
         """
         How to pick:
-            1. is one of use_per_block_weight_quant_ops
-            2. Choose specific submodule config if given.
+            1. is one of per_block_quant_config
+            2. Pick specific submodule config if given.
             3. Pick one if op belongs to use_per_channel_weight_quant_ops
-            4. If not 2, pick normal quant config
+            4. If not 3, pick normal quant config
         """
         op = node.target
         if isinstance(op, str):
             return
 
-        if block_size := self.block_size_map.get(op.name):
+        if block_size := self.block_size_map.get(node.name):
             config = self.default_quant_config.per_block_quant_config
             config.block_size = block_size
             return config
 
-        config = self.module_qconfig_dict.get(
-            self._get_submodule(node), self.default_quant_config
-        )
+        config = self._get_submodule_qconfig(node)
 
         if op in config.use_per_channel_weight_quant_ops:
             return config.per_channel_quant_config
@@ -290,16 +279,55 @@ def set_default_quant_config(
     def set_block_size_map(self, block_size_map: Dict[str, Tuple]) -> None:
         self.block_size_map = block_size_map
 
-    def set_submodule_quant_config(
-        self, submodule: torch.nn.Module, module_qconfig: ModuleQConfig
+    def set_submodule_qconfig_list(
+        self, submodule_qconfig_list: List[Tuple[Callable, ModuleQConfig]]
     ) -> None:
         """
-        Set the quant config specific for a submodule
+        Set specific quant config from a callback function.
+        If a node fits more than one callback, only apply the first one.
         """
-        self.module_qconfig_dict[submodule] = module_qconfig
+        self.submodule_qconfig_list = submodule_qconfig_list
 
     def transform_for_annotation(self, model: GraphModule) -> GraphModule:
         return QnnPassManager().transform_for_annotation_pipeline(model)
 
     def validate(self, model: GraphModule) -> None:
         pass
+
+
+def get_submodule_type_predicate(module_type_str):
+    """
+    An example of nn_module_stack
+    {
+        'L__self__': ('', 'executorch.backends.qualcomm.tests.models.SubModules'),
+        'L__self___add': ('add', 'executorch.backends.qualcomm.tests.models.Add')
+    }
+    """
+
+    def predicate(node):
+        if nn_module_stack := node.meta.get("nn_module_stack"):
+            for _, type_name in nn_module_stack.values():
+                if module_type_str in type_name:
+                    return True
+        return False
+
+    return predicate
+
+
+def get_submodule_name_predicate(module_name_str):
+    """
+    An example of nn_module_stack
+    {
+        'L__self__': ('', 'executorch.backends.qualcomm.tests.models.SubModules'),
+        'L__self___add': ('add', 'executorch.backends.qualcomm.tests.models.Add')
+    }
+    """
+
+    def predicate(node):
+        if nn_module_stack := node.meta.get("nn_module_stack"):
+            for name in nn_module_stack.keys():
+                if module_name_str in name:
+                    return True
+        return False
+
+    return predicate

backends/qualcomm/tests/test_qnn_delegate.py

@@ -2131,11 +2131,20 @@ def test_qnn_backend_submodules(self):
             torch.rand(1, 3, 8, 8),
         )
 
-        submodule_quant_config = {
-            Add: ModuleQConfig(QuantDtype.use_16a16w)  # noqa: F405
-        }
+        from executorch.backends.qualcomm.quantizer.quantizer import (
+            get_submodule_type_predicate,
+        )
+
+        submodule_qconfig_list = [
+            (
+                get_submodule_type_predicate("Add"),
+                ModuleQConfig(QuantDtype.use_16a16w),
+            )  # noqa: F405
+        ]
         module = self.get_qdq_module(
-            module, sample_input, submodule_quant_config=submodule_quant_config
+            module,
+            sample_input,
+            submodule_qconfig_list=submodule_qconfig_list,
         )
         self.lower_module_and_test_output(module, sample_input)
 

backends/qualcomm/tests/utils.py

@@ -16,11 +16,7 @@
 
 from executorch import exir
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
-from executorch.backends.qualcomm.quantizer.quantizer import (
-    ModuleQConfig,
-    QnnQuantizer,
-    QuantDtype,
-)
+from executorch.backends.qualcomm.quantizer.quantizer import ModuleQConfig, QuantDtype
 from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset
 from executorch.backends.qualcomm.utils.constants import (
     QCOM_DTYPE,
@@ -508,8 +504,9 @@ def get_qdq_module(
         dynamic_shapes: Dict = None,
         bypass_check: bool = False,
         block_size_map: Dict[str, Tuple] = None,
-        submodule_quant_config: Optional[Dict[torch.nn.Module, ModuleQConfig]] = None,
+        submodule_qconfig_list: Optional[List[Tuple[Callable, ModuleQConfig]]] = None,
     ) -> torch.fx.GraphModule:
+        module = module.eval()
         m = torch.export.export(
             module, inputs, dynamic_shapes=dynamic_shapes, strict=True
         ).module()
@@ -519,7 +516,7 @@ def get_qdq_module(
             custom_annotations=custom_quant_annotations,
             per_channel_conv=is_conv_per_channel,
             per_channel_linear=is_linear_per_channel,
-            submodule_quant_config = submodule_quant_config,
+            submodule_qconfig_list=submodule_qconfig_list,
         )
         if block_size_map is not None:
             quantizer.set_block_size_map(block_size_map)
@@ -547,7 +544,7 @@ def get_prepared_qat_module(
         is_linear_per_channel: Optional[bool] = False,
         custom_quant_annotations: Tuple[Callable] = (),
         quant_dtype: QuantDtype = QuantDtype.use_8a8w,
-        submodule_quant_config: Optional[Dict[str, ModuleQConfig]] = None,
+        submodule_qconfig_list: Optional[List[Tuple[Callable, ModuleQConfig]]] = None,
     ) -> torch.fx.GraphModule:
         m = torch.export.export_for_training(module, inputs, strict=True).module()
 
@@ -557,12 +554,11 @@ def get_prepared_qat_module(
             per_channel_conv=is_conv_per_channel,
             per_channel_linear=is_linear_per_channel,
             is_qat=True,
-            submodule_quant_config=submodule_quant_config
+            submodule_qconfig_list=submodule_qconfig_list,
         )
 
-        submodule_quant_config = submodule_quant_config or {}
-        for submodule, module_qconfig in submodule_quant_config.items():
-            quantizer.set_submodule_quant_config(submodule, module_qconfig)
+        submodule_qconfig_list = submodule_qconfig_list or []
+        quantizer.set_submodule_qconfig_list(submodule_qconfig_list)
 
         prepared = prepare_qat_pt2e(m, quantizer)
         return torch.ao.quantization.move_exported_model_to_train(prepared)

backends/qualcomm/utils/constants.py

@@ -21,7 +21,6 @@
 QCOM_INSERTED_PERMUTE = "qnn_permute"
 QCOM_LAYOUT_CHANGE = "layout_change"
 QCOM_NUM_BLOCKS_PER_AXIS = "num_blocks_per_axis"
-QCOM_NN_MODULE_STACK = "nn_module_stack"
 QCOM_OFFSET = "offset"
 QCOM_ORIG_DTYPE = "orig_dtype"
 QCOM_QUANTIZED_IO = "q_tensor_io"

examples/qualcomm/utils.py

@@ -14,7 +14,7 @@
 import tempfile
 from pathlib import Path
 
-from typing import Callable, Dict, List, Optional
+from typing import Callable, List, Optional, Tuple
 
 import numpy as np
 
@@ -262,7 +262,7 @@ def make_quantizer(
     per_channel_linear=False,
     act_observer=MovingAverageMinMaxObserver,
     is_qat=False,
-    submodule_quant_config: Optional[Dict[str, ModuleQConfig]] = None,
+    callback_qconfig_list: Optional[List[Tuple[Callable, ModuleQConfig]]] = None,
 ):
     quantizer = QnnQuantizer()
     quantizer.add_custom_quant_annotations(custom_annotations)
@@ -273,9 +273,8 @@ def make_quantizer(
         is_linear_per_channel=per_channel_linear,
         act_observer=act_observer,
     )
-    submodule_quant_config = submodule_quant_config or {}
-    for submodule, module_qconfig in submodule_quant_config.items():
-        quantizer.set_submodule_quant_config(submodule, module_qconfig)
+    callback_qconfig_list = callback_qconfig_list or []
+    quantizer.set_submodule_qconfig_list(callback_qconfig_list)
     return quantizer