Operator Fusion

The operator fusion supported by the training tool can be divided into two main categories: 1. absorb BN. 2. fuse Add, ReLU(6).

Absorb BN

The purpose of absorbing BN is to reduce the computing workload of the model. Since BN is a linear transformation process, the parameters of BN can be absorbed into the parameters of Conv when both BN and Conv occur together, thus eliminating the computation of BN from the deployed model.

The computation of absorption proceeds as follows:

By absorbing BN, Conv2d + BN2d can be simplified to Conv2d:

Fuse Add、ReLU(6)

Different from the CUDA Kernel Fusion that fuses the CUDA Kernel to increase the computation speed, the fusion supported by the training tool is more on the quantization level.

The BPU hardware provides optimization for common model basic structures, and when computing Conv -> Add -> ReLU operator combinations, it allows the data transfer between operators to retain a high degree of accuracy, improving the overall numerical accuracy of the model. Therefore, when quantizing the model, we can consider Conv -> Add -> ReLU as a whole.

The operator fusion, in addition to allowing the intermediate results to retain the high precision state, also eliminates the need to convert the intermediate results to a low precision representation, and therefore the execution speed will be faster compared to no fusion.

Since operator fusion improves both model accuracy and model speed, it should generally be done for all fusable parts. In graph mode, the fusion process is performed automatically by the tool, and you do not need to pay much attention to it.

Implementation Principle

Thanks to FX's advantage of obtaining computational graphs, the training tool can automate the analysis of the model's computational graphs, match the fusible parts according to the predefined fusion pattern, and implement the fusion operation by submodule substitution.

Conv+BN Fusion

import torch
from torch import nn
from torch.quantization import DeQuantStub

from horizon_plugin_pytorch import qint8, set_march
from horizon_plugin_pytorch.quantization import QuantStub, get_qconfig, prepare
from horizon_plugin_pytorch.quantization.qconfig_setter import *


class ModelForFusion(torch.nn.Module):
    def __init__(self):
        super(ModelForFusion, self).__init__()
        self.quant = QuantStub()
        self.conv = nn.Conv2d(3, 3, 3)
        self.bn = nn.BatchNorm2d(3)
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.conv(x)
        x = self.bn(x)
        x = self.dequant(x)

        return x


float_model = ModelForFusion()

set_march("nash-m")
fused_model = prepare(
    float_model,
    torch.rand(1, 3, 32, 32),
    QconfigSetter(
        get_qconfig(), [ModuleNameTemplate({"": qint8}), ConvDtypeTemplate()]
    ),
)

print(fused_model)
"""
GraphModuleImpl(
  (quant): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (conv): Conv2d(
    3, 3, kernel_size=(3, 3), stride=(1, 1)
    (weight_fake_quant): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_channel_symmetric, ch_axis=0, scale=tensor([1., 1., 1.]), zero_point=tensor([0, 0, 0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (bn): Identity()
  (dequant): DeQuantStub()
)
"""

As you can see, after performing the operator fusion operation on the model, BN is fused into Conv, and the original submodule is replaced with Identity.

Conv+Add/ReLU Fusion

import torch
from torch import nn
from torch.quantization import DeQuantStub

from horizon_plugin_pytorch import qint8, set_march
from horizon_plugin_pytorch.quantization import QuantStub, get_qconfig, prepare
from horizon_plugin_pytorch.quantization.qconfig_setter import *


class ModelForFusion(torch.nn.Module):
    def __init__(self):
        super(ModelForFusion, self).__init__()
        self.quantx = QuantStub()
        self.quanty = QuantStub()
        self.conv = nn.Conv2d(3, 3, 3)
        self.bn = nn.BatchNorm2d(3)
        self.relu = nn.ReLU()
        self.dequant = DeQuantStub()

    def forward(self, x, y):
        x = self.quantx(x)
        y = self.quanty(y)
        x = self.conv(x)
        x = self.bn(x)
        x = x + y
        x = self.relu(x)
        x = self.dequant(x)
        return x


float_model = ModelForFusion()

set_march("nash-m")
fused_model = prepare(
    float_model,
    (torch.rand(1, 3, 32, 32), torch.rand(1, 3, 30, 30)),
    QconfigSetter(
        get_qconfig(), [ModuleNameTemplate({"": qint8}), ConvDtypeTemplate()]
    ),
)

print(fused_model)
"""
GraphModuleImpl(
  (quantx): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (quanty): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (conv): Conv2d(
    3, 3, kernel_size=(3, 3), stride=(1, 1)
    (weight_fake_quant): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_channel_symmetric, ch_axis=0, scale=tensor([1., 1., 1.]), zero_point=tensor([0, 0, 0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (bn): Identity()
  (relu): ReLU()
  (dequant): DeQuantStub()
  (_generated_add_0): FloatFunctional()
)
"""

The tool automatically replaces the plus sign in x = x + y within the model with a Module named _generated_add_0, to support operations related to operator fusion and quantization.

Each layer out qconfig:
+------------------+----------------------------------------------------------------------+--------------------------+-----------------+-------------------+-----------------------------------------+
| Module Name      | Module Type                                                          | Input dtype              | out dtype       | ch_axis           | observer                                |
|------------------+----------------------------------------------------------------------+--------------------------+-----------------+-------------------+-----------------------------------------|
| quantx           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub                        | [torch.float32]          | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| quanty           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub                        | [torch.float32]          | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| conv             | horizon_plugin_pytorch.nn.qat.conv2d.Conv2d                          | ['qint8']                | [torch.float32] | activation = None |                                         |
| _generated_add_0 | horizon_plugin_pytorch.nn.qat.functional_modules.FloatFunctional.add | [torch.float32, 'qint8'] | [torch.float32] | activation = None |                                         |
| relu             | torch.nn.modules.activation.ReLU                                     | [torch.float32]          | [torch.float32] | qconfig = None    |                                         |
| dequant          | horizon_plugin_pytorch.nn.qat.stubs.DeQuantStub                      | [torch.float32]          | [torch.float32] | qconfig = None    |                                         |
+------------------+----------------------------------------------------------------------+--------------------------+-----------------+-------------------+-----------------------------------------+

From the model check results, it can be seen that during the quantization phase, the float32 data type is used between conv, add, and relu operators to simulate the "intermediate high precision" after fusion. In the subsequent model deployment phase, multiple computations will be merged to achieve actual fusion.

Version Migration Note

The aforementioned document demonstrates the fusion process when configuring quantization using QconfigSetter. Prior to this, the tool offered an alternative configuration method via TemplateQconfigSetter.

Under the old configuration method, the fusion result is:

GraphModuleImpl(
  (quantx): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (quanty): QuantStub(
    (activation_post_process): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
  (conv): Identity()
  (bn): Identity()
  (relu): Identity()
  (dequant): DeQuantStub()
  (_generated_add_0): ConvAddReLU2d(
    3, 3, kernel_size=(3, 3), stride=(1, 1)
    (weight_fake_quant): FakeQuantize(
      dtype=qint8, fake_quant_enabled=True, observer_enabled=True, qscheme=torch.per_channel_symmetric, ch_axis=0, scale=tensor([1., 1., 1.]), zero_point=tensor([0, 0, 0])
      (activation_post_process): MinMaxObserver(averaging_constant=0.01)
    )
  )
)

Each layer out qconfig:
+------------------+----------------------------------------------------+--------------------+-----------------+-------------------+-----------------------------------------+
| Module Name      | Module Type                                        | Input dtype        | out dtype       | ch_axis           | observer                                |
|------------------+----------------------------------------------------+--------------------+-----------------+-------------------+-----------------------------------------|
| quantx           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub      | [torch.float32]    | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| quanty           | horizon_plugin_pytorch.nn.qat.stubs.QuantStub      | [torch.float32]    | ['qint8']       | -1                | MinMaxObserver(averaging_constant=0.01) |
| _generated_add_0 | horizon_plugin_pytorch.nn.qat.conv2d.ConvAddReLU2d | ['qint8', 'qint8'] | [torch.float32] | activation = None |                                         |
| dequant          | horizon_plugin_pytorch.nn.qat.stubs.DeQuantStub    | [torch.float32]    | [torch.float32] | qconfig = None    |                                         |
+------------------+----------------------------------------------------+--------------------+-----------------+-------------------+-----------------------------------------+

As can be seen, a group of Module operators is fused into a single ConvAddReLU2d, which replaces _generated_add_0 in the original model.

Therefore, when migrating quantization configuration methods, discrepancies may appear in the model print outputs and model check results. Please do not mechanically modifying qconfig to align the model check result.

Operators that can be Fused

The currently supported combinations of fusable operators are shown in the following function definitions:

import operator
import torch
from torch import nn
from horizon_plugin_pytorch import nn as horizon_nn


def register_fusion_patterns():
    convs = (
        nn.Conv2d,
        nn.ConvTranspose2d,
        nn.Conv3d,
        nn.Linear,
    )
    bns = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
    adds = (
        nn.quantized.FloatFunctional.add,
        horizon_nn.quantized.FloatFunctional.add,
        torch.add,
        operator.add,  # the plus sign used in the code
    )
    relus = (nn.ReLU, nn.ReLU6, nn.functional.relu, nn.functional.relu6)

    for conv in convs:
        for bn in bns:
            for add in adds:
                for relu in relus:
                    # conv bn
                    register_fusion_pattern((bn, conv))(ConvBNAddReLUFusion)

                    # conv relu
                    register_fusion_pattern((relu, conv))(ConvBNAddReLUFusion)

                    # conv add
                    register_fusion_pattern((add, conv, MatchAllNode))(
                        ConvBNAddReLUFusion
                    )  # the output of conv is used as the first input to add
                    register_fusion_pattern((add, MatchAllNode, conv))(
                        ConvBNAddedReLUFusion
                    )  # The output of conv is used as the second input to add

                    # conv bn relu
                    register_fusion_pattern((relu, (bn, conv)))(
                        ConvBNAddReLUFusion
                    )

                    # conv bn add
                    register_fusion_pattern((add, (bn, conv), MatchAllNode))(
                        ConvBNAddReLUFusion
                    )
                    register_fusion_pattern((add, MatchAllNode, (bn, conv)))(
                        ConvBNAddedReLUFusion
                    )

                    # conv add relu
                    register_fusion_pattern((relu, (add, conv, MatchAllNode)))(
                        ConvBNAddReLUFusion
                    )
                    register_fusion_pattern((relu, (add, MatchAllNode, conv)))(
                        ConvBNAddedReLUFusion
                    )

                    # conv bn add relu
                    register_fusion_pattern(
                        (relu, (add, (bn, conv), MatchAllNode))
                    )(ConvBNAddReLUFusion)
                    register_fusion_pattern(
                        (relu, (add, MatchAllNode, (bn, conv)))
                    )(ConvBNAddedReLUFusion)

#Operator Fusion

#Absorb BN

#Fuse Add、ReLU(6)

#Implementation Principle

#Conv+BN Fusion

#Conv+Add/ReLU Fusion

#Version Migration Note

#Operators that can be Fused