[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]]:
         """

nncf/experimental/torch/fx/nncf_graph_builder.py

@@ -24,6 +24,7 @@
 from nncf.experimental.torch.fx.node_utils import get_tensor_constant_from_node
 from nncf.torch.dynamic_graph.layer_attributes_handlers import apply_args_defaults
 from nncf.torch.graph.graph import PTNNCFGraph
+from nncf.torch.graph.operator_metatypes import FX_OPERATOR_METATYPES
 from nncf.torch.graph.operator_metatypes import PT_OPERATOR_METATYPES
 
 
@@ -65,22 +66,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 +75,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 +87,21 @@ 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)
+            # For FX specific metatypes not registered in PT operator metatype
+            node_metatype = (
+                FX_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
+                if node_metatype == UnknownMetatype
+                else node_metatype
+            )
         else:
             node_type = node.op
             node_metatype = UnknownMetatype
@@ -118,7 +112,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 +130,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
             )
@@ -148,7 +142,7 @@ def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
             source_nncf_node = nncf_graph.get_node_by_name(source_node.name)
             for idx, dist_node in enumerate(source_node.users):
                 dist_node_id = nncf_graph.get_node_by_name(dist_node.name).node_id
-                input_port_id, output_port_id, tensor_shape = GraphConverter.get_edge_params(
+                input_port_id, output_port_id, tensor_shape, tensor_dtype = GraphConverter.get_edge_params(
                     model, source_node, source_nncf_node, dist_node, idx
                 )
                 nncf_graph.add_edge_between_nncf_nodes(
@@ -157,7 +151,7 @@ def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
                     tensor_shape=tensor_shape,
                     input_port_id=input_port_id,
                     output_port_id=output_port_id,
-                    dtype=Dtype.FLOAT,
+                    dtype=tensor_dtype,
                 )
         return nncf_graph
 
@@ -182,8 +176,11 @@ def get_edge_params(
         """
         output_port_id = 0
         tensor_shape = None
+        tensor_dtype = Dtype.FLOAT
         if source_node.op in ("get_attr",):
-            tensor_shape = tuple(get_tensor_constant_from_node(source_node, model).shape)
+            tensor = get_tensor_constant_from_node(source_node, model)
+            tensor_shape = tuple(tensor.shape)
+            tensor_dtype = Dtype.INTEGER if tensor.dtype == torch.int else tensor_dtype
         elif "val" in source_node.meta:
             if source_nncf_node.metatype is om.PTBatchNormMetatype and isinstance(
                 source_node.meta["val"], (tuple, list)
@@ -197,6 +194,7 @@ def get_edge_params(
                 tensor = source_node.meta["val"]
             if isinstance(tensor, torch.Tensor):
                 tensor_shape = tuple(-1 if isinstance(i, torch.SymInt) else i for i in tensor.shape)
+                tensor_dtype = Dtype.INTEGER if tensor.dtype == torch.int else tensor_dtype
             elif isinstance(tensor, torch.SymInt):
                 tensor_shape = (-1,)
 
@@ -205,4 +203,4 @@ def get_edge_params(
             nncf_logger.debug(f"Edge shape between {source_node.name} and {dist_node.name} is unknown.")
 
         input_port_id = dist_node.all_input_nodes.index(source_node)
-        return input_port_id, output_port_id, tensor_shape
+        return input_port_id, output_port_id, tensor_shape, tensor_dtype

nncf/experimental/torch2/function_hook/extractor.py

@@ -176,7 +176,7 @@ def extract_conv(
     if input_node == output_node:
         return conv_module
 
-    if output_node.metatype is not om.PTBatchNormMetatype:
+    if output_node.metatype != om.PT2BatchNormMetatype:
         msg = f"Support only PTBatchNormMetatype as output node, actual: {output_node.metatype}"
         raise nncf.InternalError(msg)
 

nncf/experimental/torch2/function_hook/graph/graph_utils.py

@@ -82,6 +82,14 @@ class FunctionMeta:
     def func_name(self) -> str:
         return self.func.__name__
 
+    @property
+    def func_namespace(self) -> str:
+        if self.func.__qualname__.split(".")[0] == "TensorBase":
+            return f"torch.tensor.{str(self.func.__name__)}"
+        elif self.func.__qualname__ == self.func.__name__:
+            return f"torch.nn.functional.{str(self.func.__name__)}"
+        return f"{str(self.func.__module__)}.{str(self.func.__name__)}"
+
 
 @dataclass
 class EdgeMeta:

nncf/experimental/torch2/function_hook/nncf_graph/nncf_graph_builder.py

@@ -23,6 +23,7 @@
 from nncf.common.graph.layer_attributes import BaseLayerAttributes
 from nncf.common.graph.layer_attributes import ConstantLayerAttributes
 from nncf.common.graph.layer_attributes import Dtype
+from nncf.common.graph.operator_metatypes import UnknownMetatype
 from nncf.experimental.torch2.function_hook.graph.build_graph_mode import build_graph
 from nncf.experimental.torch2.function_hook.graph.graph_utils import ConstMeta
 from nncf.experimental.torch2.function_hook.graph.graph_utils import EdgeMeta
@@ -48,7 +49,7 @@ def get_node_type(type: NodeType, meta: Union[ConstMeta, FunctionMeta, InOutMeta
     if isinstance(meta, ConstMeta):
         return "nncf_model_const"
     if isinstance(meta, FunctionMeta):
-        return meta.func_name
+        return meta.func_namespace
     msg = "Unexpected metadata type"
     raise nncf.InternalError(msg)
 
@@ -89,9 +90,14 @@ def get_meta_type(node_type: str, meta: Union[ConstMeta, FunctionMeta, InOutMeta
     :param meta: The metadata associated with the node.
     :return: The PTOperatorMetatype object.
     """
-    node_metatype = cast(
-        type[om.PTOperatorMetatype], om.PT_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
+    metatype = om.PT_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
+    metatype = (
+        om.PT2_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
+        if metatype == UnknownMetatype
+        else metatype
     )
+
+    node_metatype = cast(type[om.PTOperatorMetatype], metatype)
     node_sub_meta_type: Optional[type[om.PTOperatorMetatype]] = None
     if node_metatype.get_subtypes() and isinstance(meta, FunctionMeta):
         node_sub_meta_type = node_metatype.determine_subtype(function_args=meta.args, functions_kwargs=meta.kwargs)
@@ -187,7 +193,7 @@ def convert_to_nncf_graph(nx_graph: nx.MultiDiGraph) -> PTNNCFGraph:
             layer_name=node_name,
             node_metatype=meta_type,
             node_name=node_name,
-            node_type=node_type,
+            node_type=node_type.split(".")[-1],
         )
         map_nx_node_to_nncf_node[node] = nncf_node
 

nncf/quantization/algorithms/min_max/torch_backend.py

@@ -311,11 +311,13 @@ def get_ignored_metatypes(model_type: ModelType, device: TargetDevice) -> List[O
                 om.PTMaxMetatype,
                 om.PTSqueezeMetatype,
                 om.PTLayerNormMetatype,
+                om.PTAtenLayerNormMetatype,
                 om.PTModuleLayerNormMetatype,
                 om.PTGroupNormMetatype,
                 om.PTModuleGroupNormMetatype,
                 # Batchnorm
                 om.PTBatchNormMetatype,
+                om.PT2BatchNormMetatype,
                 om.PTModuleBatchNormMetatype,
                 # Comparison operations
                 om.PTGreaterEqualMetatype,

nncf/quantization/algorithms/min_max/torch_fx_backend.py

@@ -286,6 +286,7 @@ def get_ignored_metatypes(model_type: ModelType, device: TargetDevice) -> List[O
                 om.PTMaxMetatype,
                 om.PTSqueezeMetatype,
                 om.PTLayerNormMetatype,
+                om.PTAtenLayerNormMetatype,
                 om.PTModuleLayerNormMetatype,
                 om.PTGroupNormMetatype,
                 om.PTModuleGroupNormMetatype,

nncf/torch/dynamic_graph/layer_attributes_handlers.py

@@ -39,7 +39,7 @@
     om.PTConvTranspose1dMetatype, om.PTConvTranspose2dMetatype, om.PTConvTranspose3dMetatype
 )
 LINEAR_OP_NAMES = get_all_aliases(om.PTLinearMetatype)
-BATCHNORM_OP_NAMES = get_all_aliases(om.PTBatchNormMetatype)
+BATCHNORM_OP_NAMES = get_all_aliases(om.PTBatchNormMetatype, om.PT2BatchNormMetatype)
 EMBEDDING_OP_NAMES = get_all_aliases(om.PTEmbeddingMetatype, om.PTEmbeddingBagMetatype)
 GROUP_NORM_OP_NAMES = get_all_aliases(om.PTGroupNormMetatype)
 LAYER_NORM_OP_NAMES = get_all_aliases(om.PTLayerNormMetatype)