Refactor weights layouts handling

nncf/common/graph/layer_attributes.py

@@ -13,34 +13,14 @@
 from abc import abstractmethod
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, List, Tuple, Union
 
 
 class Dtype(Enum):
     FLOAT = "float"
     INTEGER = "int"
 
 
-class ConvLayoutElem(Enum):
-    """
-    Layout elements descriptor for convolutional and linear layers:
-        C_IN: Input channels dimension.
-        C_OUT: Output channels dimension.
-        SPATIAL: Spatial dimension.
-        GROUPS: Groups dimention.
-    In case a backend is not using a separate dimension for groups and
-    has layout C_IN // GROUPS for input channels of a convolution operation
-    and C_OUT // GROUPS for output channels of a transpose convolution operation
-    input/output channel dimentions are interpretated as
-    input/output channel dimension but with size divided by groups number.
-    """
-
-    C_IN = "channels_in"
-    C_OUT = "channels_out"
-    SPATIAL = "spatial"
-    GROUPS = "groups"
-
-
 class BaseLayerAttributes(ABC):
     """
     This class stores base useful for some algorithms attributes
@@ -138,7 +118,6 @@ def __init__(
         in_features: int,
         out_features: int,
         with_bias: bool = True,
-        weights_layout: Optional[Tuple[ConvLayoutElem, ...]] = None,
     ):
         """
 
@@ -150,7 +129,6 @@ def __init__(
         super().__init__(weight_requires_grad, with_bias=with_bias)
         self.in_features = in_features
         self.out_features = out_features
-        self.weights_layout = weights_layout
 
     def get_weight_shape(self) -> List[int]:
         return [self.out_features, self.in_features]
@@ -175,7 +153,6 @@ def __init__(
         transpose: bool,
         padding_values: Tuple[int, ...],
         with_bias: bool = False,
-        weights_layout: Optional[Tuple[ConvLayoutElem, ...]] = None,
     ):
         """
 
@@ -199,7 +176,6 @@ def __init__(
         self.groups = groups
         self.transpose = transpose
         self.padding_values = padding_values
-        self.weights_layout = weights_layout
 
     def get_weight_shape(self) -> List[int]:
         if not self.transpose:

nncf/openvino/graph/layer_attributes.py

@@ -9,16 +9,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import openvino.runtime as ov
 
+from nncf.common.graph.graph import NNCFNode
 from nncf.common.graph.layer_attributes import BaseLayerAttributes
-from nncf.common.graph.layer_attributes import ConvLayoutElem
 from nncf.common.graph.layer_attributes import ConvolutionLayerAttributes
 from nncf.common.graph.layer_attributes import GenericWeightedLayerAttributes
 from nncf.common.graph.layer_attributes import LinearLayerAttributes
 from nncf.common.graph.layer_attributes import WeightedLayerAttributes
+from nncf.openvino.graph.layout import OVConvLayoutElem
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVConvolutionBackpropDataMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVConvolutionMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVDepthwiseConvolutionMetatype
@@ -70,6 +71,62 @@ def get_backend_agnostic_attributes(self):
         return self._layer_attributes
 
 
+OV_CONV_METATYPES = [
+    OVConvolutionMetatype,
+    OVDepthwiseConvolutionMetatype,
+    OVGroupConvolutionMetatype,
+    OVConvolutionBackpropDataMetatype,
+    OVGroupConvolutionBackpropDataMetatype,
+]
+
+
+OVConvLayout = List[OVConvLayoutElem]
+
+
+def get_conv_weights_layout_from_node(node: NNCFNode) -> OVConvLayout:
+    layer_attributes = node.layer_attributes
+    port_id = _get_port_id_from_layer_attributes(layer_attributes)
+    return get_conv_weights_layout(
+        ov_metatype=node.metatype, weights_shape=layer_attributes.constant_attributes[port_id]["shape"]
+    )
+
+
+def get_linear_weights_layout_from_node(node: NNCFNode) -> OVConvLayout:
+    layer_attributes = node.layer_attributes
+    port_id = _get_port_id_from_layer_attributes(layer_attributes)
+    constant_layer_attrs = layer_attributes.constant_attributes[port_id]
+    return get_linear_weights_layout(
+        weights_shape=constant_layer_attrs["shape"],
+        transpose=constant_layer_attrs.get("transpose", False),
+        port_id=port_id,
+    )
+
+
+def _get_port_id_from_layer_attributes(layer_attributes):
+    port_ids = list(layer_attributes.constant_attributes.keys())
+    assert len(port_ids) == 1
+    return port_ids[0]
+
+
+def get_conv_weights_layout(ov_metatype: OVOpMetatype, weights_shape: Tuple[int, ...]) -> OVConvLayout:
+    weights_layout = ov_metatype.const_layout
+    kernel_size = weights_shape[len(weights_layout) :]
+    weights_layout += [OVConvLayoutElem.SPATIAL] * len(kernel_size)
+    return tuple(weights_layout)
+
+
+def get_linear_weights_layout(weights_shape: Tuple[int, ...], transpose: bool, port_id: int) -> OVConvLayout:
+    weights_layout = [OVConvLayoutElem.SPATIAL] * (len(weights_shape) - 2)
+    if len(weights_shape) > 1:
+        if (transpose and port_id == 0) or (not transpose and port_id == 1):
+            weights_layout += [OVConvLayoutElem.C_IN, OVConvLayoutElem.C_OUT]
+        else:
+            weights_layout += [OVConvLayoutElem.C_OUT, OVConvLayoutElem.C_IN]
+    else:
+        weights_layout += [OVConvLayoutElem.C_IN]
+    return tuple(weights_layout)
+
+
 def get_weighted_layer_attributes(
     ov_node: ov.Node, ov_metatype: OVOpMetatype, constant_attributes: Dict[int, Any]
 ) -> WeightedLayerAttributes:
@@ -85,13 +142,7 @@ def get_weighted_layer_attributes(
         return None
 
     port_id, attrs = constant_attributes.copy().popitem()
-    if ov_metatype in [
-        OVConvolutionMetatype,
-        OVDepthwiseConvolutionMetatype,
-        OVGroupConvolutionMetatype,
-        OVConvolutionBackpropDataMetatype,
-        OVGroupConvolutionBackpropDataMetatype,
-    ]:
+    if ov_metatype in OV_CONV_METATYPES:
         node_attrs = ov_node.get_attributes()
         kwargs = {
             "weight_requires_grad": False,
@@ -101,43 +152,35 @@ def get_weighted_layer_attributes(
             # TODO: ticket 114378: unify pad attribute
             "padding_values": tuple(node_attrs["pads_begin"] + node_attrs["pads_end"]),
         }
-
-        weights_layout = ov_metatype.const_layout
         weights_shape = attrs["shape"]
+        weights_layout = get_conv_weights_layout(ov_metatype=ov_metatype, weights_shape=weights_shape)
         kwargs.update(
             {
-                "in_channels": weights_shape[weights_layout.index(ConvLayoutElem.C_IN)],
-                "out_channels": weights_shape[weights_layout.index(ConvLayoutElem.C_OUT)],
-                "kernel_size": tuple(weights_shape[len(weights_layout) :]),
-                "groups": weights_shape[weights_layout.index(ConvLayoutElem.GROUPS)]
-                if ConvLayoutElem.GROUPS in weights_layout
+                "in_channels": weights_shape[weights_layout.index(OVConvLayoutElem.C_IN)],
+                "out_channels": weights_shape[weights_layout.index(OVConvLayoutElem.C_OUT)],
+                "kernel_size": tuple(
+                    dim for dim, elem in zip(weights_shape, weights_layout) if elem == OVConvLayoutElem.SPATIAL
+                ),
+                "groups": weights_shape[weights_layout.index(OVConvLayoutElem.GROUPS)]
+                if OVConvLayoutElem.GROUPS in weights_layout
                 else 1,
             }
         )
-        kwargs.update({"weights_layout": tuple(weights_layout + len(kwargs["kernel_size"]) * [ConvLayoutElem.SPATIAL])})
 
         return ConvolutionLayerAttributes(**kwargs)
     if ov_metatype == OVMatMulMetatype:
         weights_shape = attrs["shape"]
-
-        weights_layout = [ConvLayoutElem.SPATIAL] * (len(weights_shape) - 2)
-        if len(weights_shape) > 1:
-            transpose = attrs.get("transpose", False)
-            if (transpose and port_id == 0) or (not transpose and port_id == 1):
-                weights_layout += [ConvLayoutElem.C_IN, ConvLayoutElem.C_OUT]
-            else:
-                weights_layout += [ConvLayoutElem.C_OUT, ConvLayoutElem.C_IN]
-        else:
-            weights_layout += [ConvLayoutElem.C_IN]
+        weights_layout = get_linear_weights_layout(
+            weights_shape=weights_shape, transpose=attrs.get("transpose", False), port_id=port_id
+        )
 
         kwargs = {
             "weight_requires_grad": False,
-            "in_features": weights_shape[weights_layout.index(ConvLayoutElem.C_IN)],
-            "out_features": weights_shape[weights_layout.index(ConvLayoutElem.C_OUT)]
-            if ConvLayoutElem.C_OUT in weights_layout
+            "in_features": weights_shape[weights_layout.index(OVConvLayoutElem.C_IN)],
+            "out_features": weights_shape[weights_layout.index(OVConvLayoutElem.C_OUT)]
+            if OVConvLayoutElem.C_OUT in weights_layout
             else None,
             "with_bias": False,
-            "weights_layout": weights_layout,
         }
         return LinearLayerAttributes(**kwargs)
     return GenericWeightedLayerAttributes(weight_requires_grad=False, weight_shape=attrs.get("shape", None))

nncf/openvino/graph/layout.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2023 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+
+
+class OVConvLayoutElem(Enum):
+    """
+    Layout elements descriptor for convolutional and linear openvino layers:
+        C_IN: Input channels dimension.
+        C_OUT: Output channels dimension.
+        SPATIAL: Spatial dimension.
+        GROUPS: Groups dimention.
+    """
+
+    C_IN = "channels_in"
+    C_OUT = "channels_out"
+    SPATIAL = "spatial"
+    GROUPS = "groups"

nncf/openvino/graph/metatypes/openvino_metatypes.py

@@ -14,13 +14,13 @@
 
 import openvino.runtime as ov
 
-from nncf.common.graph.layer_attributes import ConvLayoutElem
 from nncf.common.graph.operator_metatypes import INPUT_NOOP_METATYPES
 from nncf.common.graph.operator_metatypes import OUTPUT_NOOP_METATYPES
 from nncf.common.graph.operator_metatypes import OperatorMetatype
 from nncf.common.graph.operator_metatypes import OperatorMetatypeRegistry
 from nncf.common.graph.operator_metatypes import UnknownMetatype
 from nncf.common.hardware.opset import HWConfigOpName
+from nncf.openvino.graph.layout import OVConvLayoutElem
 
 OV_OPERATOR_METATYPES = OperatorMetatypeRegistry("openvino_operator_metatypes")
 
@@ -60,7 +60,7 @@ class OVConvolutionMetatype(OVOpMetatype):
     op_names = ["Convolution"]
     hw_config_names = [HWConfigOpName.CONVOLUTION]
     const_channel_axis = [0]
-    const_layout = [ConvLayoutElem.C_OUT, ConvLayoutElem.C_IN]
+    const_layout = [OVConvLayoutElem.C_OUT, OVConvLayoutElem.C_IN]
     output_channel_axis = 1
 
 
@@ -70,7 +70,7 @@ class OVConvolutionBackpropDataMetatype(OVOpMetatype):
     op_names = ["ConvolutionBackpropData"]
     hw_config_names = [HWConfigOpName.CONVOLUTION]
     const_channel_axis = [1]
-    const_layout = [ConvLayoutElem.C_IN, ConvLayoutElem.C_OUT]
+    const_layout = [OVConvLayoutElem.C_IN, OVConvLayoutElem.C_OUT]
     output_channel_axis = 1
 
 
@@ -80,7 +80,7 @@ class OVDepthwiseConvolutionMetatype(OVOpMetatype):
     op_names = ["GroupConvolution"]
     hw_config_names = [HWConfigOpName.DEPTHWISECONVOLUTION]
     const_channel_axis = [0, 1]
-    const_layout = [ConvLayoutElem.GROUPS, ConvLayoutElem.C_OUT, ConvLayoutElem.C_IN]
+    const_layout = [OVConvLayoutElem.GROUPS, OVConvLayoutElem.C_OUT, OVConvLayoutElem.C_IN]
     output_channel_axis = 1
 
     @classmethod
@@ -95,7 +95,7 @@ class OVGroupConvolutionMetatype(OVOpMetatype):
     hw_config_names = [HWConfigOpName.CONVOLUTION]
     subtypes = [OVDepthwiseConvolutionMetatype]
     const_channel_axis = [0, 1]
-    const_layout = [ConvLayoutElem.GROUPS, ConvLayoutElem.C_OUT, ConvLayoutElem.C_IN]
+    const_layout = [OVConvLayoutElem.GROUPS, OVConvLayoutElem.C_OUT, OVConvLayoutElem.C_IN]
     output_channel_axis = 1
 
 
@@ -105,7 +105,7 @@ class OVGroupConvolutionBackpropDataMetatype(OVOpMetatype):
     op_names = ["GroupConvolutionBackpropData"]
     hw_config_names = [HWConfigOpName.CONVOLUTION]
     const_channel_axis = [0, 2]
-    const_layout = [ConvLayoutElem.GROUPS, ConvLayoutElem.C_IN, ConvLayoutElem.C_OUT]
+    const_layout = [OVConvLayoutElem.GROUPS, OVConvLayoutElem.C_IN, OVConvLayoutElem.C_OUT]
     output_channel_axis = 1
 
 

nncf/openvino/graph/node_utils.py

@@ -17,9 +17,8 @@
 
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.graph.graph import NNCFNode
-from nncf.common.graph.layer_attributes import ConvLayoutElem
-from nncf.common.graph.layer_attributes import LinearLayerAttributes
-from nncf.openvino.graph.layer_attributes import OVLayerAttributes
+from nncf.openvino.graph.layer_attributes import get_linear_weights_layout_from_node
+from nncf.openvino.graph.layout import OVConvLayoutElem
 from nncf.openvino.graph.metatypes.groups import OPERATIONS_WITH_BIAS
 from nncf.openvino.graph.metatypes.groups import OPERATIONS_WITH_WEIGHTS
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVAddMetatype
@@ -336,13 +335,9 @@ def get_matmul_channel_axes(node: ov.Node) -> List[int]:
     :param node: The target node.
     :return: List of channel axes for the MatMul operation.
     """
-    assert isinstance(node.layer_attributes, OVLayerAttributes)
-    layer_attributes = node.layer_attributes.get_backend_agnostic_attributes()
-    assert isinstance(layer_attributes, LinearLayerAttributes)
+    weights_layout = get_linear_weights_layout_from_node(node)
     return [
-        idx
-        for idx, elem in enumerate(layer_attributes.weights_layout)
-        if elem in [ConvLayoutElem.SPATIAL, ConvLayoutElem.C_OUT]
+        idx for idx, elem in enumerate(weights_layout) if elem in [OVConvLayoutElem.SPATIAL, OVConvLayoutElem.C_OUT]
     ]
 
 

nncf/openvino/quantization/weights_compression.py

@@ -15,10 +15,9 @@
 import openvino.runtime as ov
 from openvino.runtime import opset9 as opset
 
-from nncf.common.graph.layer_attributes import ConvLayoutElem
-from nncf.common.graph.layer_attributes import LinearLayerAttributes
 from nncf.common.graph.operator_metatypes import OperatorMetatype
-from nncf.openvino.graph.layer_attributes import get_weighted_layer_attributes
+from nncf.openvino.graph.layer_attributes import get_linear_weights_layout
+from nncf.openvino.graph.layout import OVConvLayoutElem
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVEmbeddingMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVMatMulMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import get_node_metatype
@@ -95,12 +94,10 @@ def _get_reduction_axes(
     """
     if metatype is OVMatMulMetatype:
         transpose = node.get_attributes()[f"transpose_{'a' if weight_port_id == 0 else 'b'}"]
-        constant_attributes = {weight_port_id: {"shape": weight_shape, "transpose": transpose}}
-        layer_attributes = get_weighted_layer_attributes(
-            ov_node=node, ov_metatype=OVMatMulMetatype, constant_attributes=constant_attributes
+        weights_layout = get_linear_weights_layout(
+            weights_shape=weight_shape, transpose=transpose, port_id=weight_port_id
         )
-        assert isinstance(layer_attributes, LinearLayerAttributes)
-        axes = tuple(idx for idx, elem in enumerate(layer_attributes.weights_layout) if elem == ConvLayoutElem.C_IN)
+        axes = tuple(idx for idx, elem in enumerate(weights_layout) if elem == OVConvLayoutElem.C_IN)
     elif metatype is OVEmbeddingMetatype:
         axes = (metatype.const_channel_axis[0] + 1) % 2
     else:

nncf/quantization/algorithms/channel_alignment/algorithm.py

@@ -18,7 +18,6 @@
 from nncf.common.factory import ModelTransformerFactory
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.graph.graph import NNCFNode
-from nncf.common.graph.layer_attributes import ConvLayoutElem
 from nncf.common.graph.patterns import GraphPattern
 from nncf.common.graph.transformations.commands import TargetPoint
 from nncf.common.graph.transformations.commands import TargetType
@@ -465,21 +464,7 @@ def __init__(
             bias = backend_entity.create_bias_tensor(conv_op, nncf_graph, 0)
         self.stated_bias = StatedTensor(bias)
         self._op = conv_op
-        weights_layout = conv_op.layer_attributes.get_backend_agnostic_attributes().weights_layout
-        if ConvLayoutElem.GROUPS in weights_layout:
-            # Using groups dim as output channels dim for ChannelAlignment algorithm
-            # TODO(dlyakhov) support group convolutions with groups number not in [1, out_channels]
-            self._dims = LayoutDescriptor(
-                weights_layout.index(ConvLayoutElem.GROUPS),
-                weights_layout.index(ConvLayoutElem.C_IN),
-                conv_op.metatype.output_channel_axis,
-            )
-        else:
-            self._dims = LayoutDescriptor(
-                weights_layout.index(ConvLayoutElem.C_OUT) if ConvLayoutElem.C_OUT in weights_layout else None,
-                weights_layout.index(ConvLayoutElem.C_IN),
-                conv_op.metatype.output_channel_axis,
-            )
+        self._dims = backend_entity.get_dims_descriptor(conv_op)
 
     @property
     def weight(self):