[Torch FX] Map Namespace Names to Metatype

nncf/common/graph/operator_metatypes.py

@@ -78,6 +78,7 @@ def __init__(self, name: str):
         """
         super().__init__(name)
         self._op_name_to_op_meta_dict: Dict[str, Type[OperatorMetatype]] = {}
+        self._func_name_to_op_meta_dict: Dict[str, Type[OperatorMetatype]] = {}
 
     def register(self, name: Optional[str] = None, is_subtype: bool = False) -> Callable[..., Type[OperatorMetatype]]:
         """
@@ -111,6 +112,17 @@ def wrap(obj: Type[OperatorMetatype]) -> Type[OperatorMetatype]:
                         )
                         raise nncf.InternalError(msg)
                     self._op_name_to_op_meta_dict[name] = obj
+                if hasattr(obj, "module_to_function_names"):
+                    for namespace, function_names in obj.module_to_function_names.items():
+                        for function_name in function_names:
+                            target_function_name = f"{namespace.value}.{function_name}"
+                            if target_function_name in self._func_name_to_op_meta_dict:
+                                msg = (
+                                    "Inconsistent operator metatype registry - single patched "
+                                    f"op name `{target_function_name}` maps to multiple metatypes!"
+                                )
+                                raise nncf.InternalError(msg)
+                            self._func_name_to_op_meta_dict[target_function_name] = obj
             return obj
 
         return wrap
@@ -126,6 +138,18 @@ def get_operator_metatype_by_op_name(self, op_name: str) -> Type[OperatorMetatyp
             return UnknownMetatype
         return self._op_name_to_op_meta_dict[op_name]
 
+    def get_operator_metatype_by_func(self, func_name: str) -> Type[OperatorMetatype]:
+        """
+        Returns the operator metatype by function name.
+
+        :param func_name: The function name.
+        :return: The operator metatype.
+        """
+        if func_name not in self._func_name_to_op_meta_dict:
+            return UnknownMetatype
+        obj = self._func_name_to_op_meta_dict[func_name]
+        return obj
+
 
 NOOP_METATYPES = Registry("noop_metatypes")
 INPUT_NOOP_METATYPES = Registry("input_noop_metatypes")

nncf/experimental/torch/fx/nncf_graph_builder.py

@@ -65,22 +65,6 @@ def _get_layer_attributes(
             )
         return None
 
-    def _map_fx_unique_metatypes(node: torch.fx.Node, metatype: om.OperatorMetatype) -> om.OperatorMetatype:
-        """
-        Attempts to retrieve correct subtype for the given node.
-
-        :param node: Given node.
-        :param metatype: Given node metatype.
-        :param model: Target GraphModule instance.
-        :return: Correct FX metatype of the given node if it is exist or the original node metatype otherwise.
-        """
-        if metatype in [om.PTEmbeddingMetatype]:
-            weight_node = node.args[0]
-            if weight_node.op == "get_attr":
-                return om.PTAtenEmbeddingMetatype
-
-        return metatype
-
     @staticmethod
     def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule) -> Tuple[str, om.OperatorMetatype]:
         """
@@ -90,6 +74,7 @@ def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule)
         :param model: Given GraphModule.
         :return: Node's type and metatype.
         """
+        node_type_name = None
         if node.op == "placeholder":
             node_type = "input"
             node_metatype = om.PTInputNoopMetatype
@@ -101,13 +86,15 @@ def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule)
             node_metatype = om.PTConstNoopMetatype
         elif node.op in ("call_function",):
             if hasattr(node.target, "overloadpacket"):
-                node_type = str(node.target.overloadpacket).split(".")[1]
+                node_type = str(node.target.overloadpacket)
+                node_type_name = node_type.split(".")[1]
             elif node.target.__name__ == "getitem":
-                node_type = "__getitem__"
+                node_type = "aten.__getitem__"
+                node_type_name = "__getitem__"
             else:
                 # TODO(dlyakhov): get correct nodes types from this nodes as well
                 node_type = str(node.target)
-            node_metatype = PT_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
+            node_metatype = PT_OPERATOR_METATYPES.get_operator_metatype_by_func(node_type)
         else:
             node_type = node.op
             node_metatype = UnknownMetatype
@@ -118,7 +105,8 @@ def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule)
             layer_attrs = GraphConverter._get_layer_attributes(node, node_metatype, model)
             node_subtype = node_metatype.determine_subtype(layer_attrs)
             node_metatype = node_subtype or node_metatype
-        return node_type, node_metatype
+        node_type_name = node_type_name or node_type
+        return node_type_name, node_metatype
 
     @staticmethod
     def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
@@ -135,7 +123,6 @@ def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
         const_targets_counter = Counter([node.target for node in model.graph.nodes if node.op == "get_attr"])
         for source_node in model.graph.nodes:
             node_type, node_metatype = GraphConverter.get_node_type_and_metatype(source_node, model)
-            node_metatype = GraphConverter._map_fx_unique_metatypes(source_node, node_metatype)
             is_shared_node = source_node.op in ("get_attr",) and (
                 const_targets_counter[source_node.target] > 1 or len(source_node.users) > 1
             )