Introduction to TVM Workflow

This article explores how to deploy neural networks on new hardware using AI compilers, covering the entire process from algorithm design to actual execution. We'll use TVM (Tensor Virtual Machine) as our example, examining its workflow:

• Model Import: TVM can import models from TensorFlow, PyTorch, ONNX, and other frameworks.
• Conversion to Relay: Relay is TVM's intermediate representation. Imported models are converted to Relay format (with Relax as the latest version).
• Conversion to TE (Tensor Expression): After graph-level optimizations in Relay, the model undergoes subgraph partitioning and conversion to TE representation.
• Auto-tuning: Using AutoTVM or AutoScheduler to search for optimal scheduling that leverages hardware characteristics for better performance. This generates the best configuration, with tuning logs selecting optimal scheduling for each subgraph.
• Conversion to TIR (Tensor Intermediate Representation): TVM's low-level intermediate representation specific to hardware platforms. TIR undergoes low-level optimizations before being lowered to the target compiler for the hardware platform. This is the final code generation stage producing optimized models for production deployment.
• Compilation to Machine Code: TVM supports various hardware backends including LLVM for general microprocessor architectures (X86, ARM), specialized compilers like NVCC for NVIDIA GPUs, and embedded/specialized platforms through TVM's code generation framework.

Algorithm Layer

Operator Optimization

When deploying networks, we must first consider the feasibility on target hardware. If a hardware platform doesn't support certain operators, we need to replace them with similar supported ones. For example, in YOLOv5, the default activation function is SiLU, which involves natural exponent and division operations—unfriendly for FPGA hardware. YOLOv5 also provides alternatives like ReLU6, which has a simpler computation method with minimal accuracy difference, making it a viable replacement.

Besides activation function replacement, another common optimization is convolution improvement. To reduce the computational overhead of standard convolutions, researchers have developed various compact convolutions:

• Channel Squeeze: SqueezeNet uses 1×1 filters with fewer channels than input channels in each network block, then expands using multiple 1×1 and 3×3 kernels. This reduces computational complexity while maintaining accuracy. SqueezeNext further reduces parameters using separable convolutions.
• Depthwise Separable Convolution: Composed of Depthwise Convolution and Pointwise Convolution. Depthwise Convolution applies a separate kernel to each input channel, while Pointwise Convolution is essentially 1×1 convolution. This approach offers high computational efficiency, as seen in Xception and MobileNet.
• Linear Bottleneck Layers: Bottleneck structures map high-dimensional spaces to low-dimensional ones, reducing channel count. MobileNet V2 introduced linear bottleneck structures (replacing ReLU with linear activation) to mitigate information loss from ReLU activations.
• Group Convolution: Input channels are divided into groups, with each group convolving independently. This reduces MAC operations compared to standard convolutions.

Quantization

Network lightweighting methods like quantization can compress networks at the algorithm level. PyTorch offers two quantization APIs: dynamic graph (eager) quantization and static graph (fx) quantization.

Before torch.fx (PyTorch 1.8), quantization was performed in eager mode. Here's a typical code example for quantization-aware training of ResNet18:

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import dataset_loader

# Model loading
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torchvision.models.resnet18(pretrained=False)
model_weights_path = "./weights/resnet18-f37072fd.pth"
model.load_state_dict(torch.load(model_weights_path), strict=False)

# Specify fused layers
fused_layers = [
    ['model.conv1', 'model.bn1', 'model.relu'],
    ['model.layer1.0.conv1', 'model.layer1.0.bn1', 'model.layer1.0.relu'],
    ['model.layer1.0.conv2', 'model.layer1.0.bn2'],
    ['model.layer1.1.conv1', 'model.layer1.1.bn1', 'model.layer1.1.relu'],
    ['model.layer1.1.conv2', 'model.layer1.1.bn2'],
    ['model.layer2.0.conv1', 'model.layer2.0.bn1', 'model.layer2.0.relu'],
    ['model.layer2.0.conv2', 'model.layer2.0.bn2'],
    ['model.layer2.1.conv1', 'model.layer2.1.bn1', 'model.layer2.1.relu'],
    ['model.layer2.1.conv2', 'model.layer2.1.bn2'],
    ['model.layer3.0.conv1', 'model.layer3.0.bn1', 'model.layer3.0.relu'],
    ['model.layer3.0.conv2', 'model.layer3.0.bn2'],
    ['model.layer3.1.conv1', 'model.layer3.1.bn1', 'model.layer3.1.relu'],
    ['model.layer3.1.conv2', 'model.layer3.1.bn2'],
    ['model.layer4.0.conv1', 'model.layer4.0.bn1', 'model.layer4.0.relu'],
    ['model.layer4.0.conv2', 'model.layer4.0.bn2'],
    ['model.layer4.1.conv1', 'model.layer4.1.bn1', 'model.layer4.1.relu'],
    ['model.layer4.1.conv2', 'model.layer4.1.bn2'],
]

# Insert quantization and dequantization nodes
class QuantizedResNet18(nn.Module):
    def __init__(self, model_fp32):
        super(QuantizedResNet18, self).__init__()
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()
        self.model_fp32 = model_fp32

    def forward(self, x):
        x = self.quant(x)
        x = self.model_fp32(x)
        x = self.dequant(x)
        return x

quantized_model = QuantizedResNet18(model)
quantized_model.eval()

# Layer fusion
quantized_model = torch.ao.quantization.fuse_modules(quantized_model, fused_layers)
quantized_model.train()

# Quantization configuration
quantization_config = torch.quantization.QConfig(
    activation=torch.quantization.MinMaxObserver.with_args(dtype=torch.quint16),
    weight=torch.quantization.MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_affine)
)
quantized_model.qconfig = quantization_config
quantized_model = torch.ao.quantization.prepare_qat(quantized_model, inplace=True)
quantized_model.to(device)

# Hyperparameters
batch_size = 64
learning_rate = 0.0001
epochs = 10

train_loader, test_loader, test_size = dataset_loader.load_data(batch_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(quantized_model.parameters(), lr=learning_rate)

# Training
best_accuracy = 0.0
save_path_fp32 = './weights/resNet18_fp32.pth'
for epoch in range(epochs):
    quantized_model.train()
    running_loss = 0.0
    for step, data in enumerate(train_loader, start=0):
        images, labels = data
        optimizer.zero_grad()
        outputs = quantized_model(images.to(device))
        loss = criterion(outputs, labels.to(device))
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        progress = (step + 1) / len(train_loader)
        progress_bar = "*" * int(progress * 50)
        remaining = "." * int((1 - progress) * 50)
        print(f"\rTraining loss: {int(progress * 100)}%[{progress_bar}{remaining}]{loss:.4f}", end="")

    quantized_model.eval()
    accuracy = 0.0
    with torch.no_grad():
        for val_data in test_loader:
            val_images, val_labels = val_data
            outputs = quantized_model(val_images.to(device))
            predictions = torch.max(outputs, dim=1)[1]
            accuracy += (predictions == val_labels.to(device)).sum().item()
        val_accuracy = accuracy / test_size
        if val_accuracy > best_accuracy:
            best_accuracy = val_accuracy
            torch.save(quantized_model.state_dict(), save_path_fp32)
        print(f'[Epoch {epoch + 1}] Training Loss: {running_loss / step:.3f}  Test Accuracy: {val_accuracy:.3f}')

save_path_int8 = './weights/resNet18_INT16.pth'
quantized_model.eval()

# Convert to int8 model for saving
int8_model = torch.quantization.convert(quantized_model, inplace=True)
torch.save(int8_model.state_dict(), save_path_int8)

This process requires several manual model modifications:

1. Insert quant and dequant nodes for explicit float32 to int8 conversion. For layers that cannot be quantized or should run in float, these nodes must be wrapped around them.
2. In eager mode quantization, all operators must be defined in __init__ using torch.nn forms, not torch.nn.functional, even for parameterless operators like ReLU.
3. Layer fusion (operator fusion) may be needed during inference, requiring manual specification of fusion patterns.

FX mode quantization traces the model first, obtaining the entire graph information before quantization. This eliminates the need for manual layer fusion specification and quant/dequant node insertion. However, FX mode has limitations like lack of support for runtime dynamic control flow.

Here's a typical FX mode quantization-aware training code:

from torch.quantization.quantize_fx import prepare_qat_fx, convert_fx
import torch.quantization.observer as observer
import copy

# Convert model to QAT version
def create_qat_model(model):
    qconfig_dict = {
        "": torch.ao.quantization.get_default_qat_qconfig('qnnpack'),
        "object_type": [
            (PositionalEmbedding, None),
            (torch.nn.Softmax, softmax_qconfig),
        ],
    }
    model_to_quantize = copy.deepcopy(model)
    prepared_model = prepare_qat_fx(model_to_quantize, qconfig_dict)
    return prepared_model

qat_model = create_qat_model(float_model)

# Model training
for i in range(epochs):
    train(qat_model, train_data_loader)
    ...

# Model evaluation
quantized_model = convert_fx(qat_model)
evaluate(quantized_model, eval_data_loader)
...

# Save quantized model
# As TorchScript
trace_model = torch.jit.trace(quantized_model, (input1, input2, input3))
torch.jit.save(trace_model, 'quantized_model.pt')

# As ONNX
torch.onnx.export(quantized_model, (input1, input2, input3), "quantized_model.onnx",
                 input_names=['input1', 'input2', 'input3'], output_names=['result'],
                 opset_version=16, export_params=True)

Key considerations for FX quantization:

• When using prepare_qat_fx, unsupported functions may cause errors. These can be resolved by annotating functions with @torch.fx.wrap.
• If loss or quantization parameters become NaN during training, try different qconfig or pre-train the model before quantization-aware training.
• If accuracy remains unchanged during quantization-aware training, ensure the optimizer is reinitialized for the quantized model.
• If convergence is slow, pre-train the model until it converges before switching to quantization-aware training.

Compilation Layer

Frontend Quantized Model Parsing

TVM uses QNN Dialect to parse quantized models. QNN is a framework developed for TVM that supports importing pre-quantized models with the following characteristics:

• QNN is a high-level IR at the computation graph level, adding new operators without graph or operator-level optimizations.
• QNN quantized operators are at a higher level than any graph-level operators.
• New QNN optimization passes can be defined to transform graphs for specific hardware backends.

TVM maintains a convert_map for each framework's Converter, mapping framework operators to Relay operators. For example, "quantized::linear" maps to _linear(), and "quantized::conv2d" maps to _quantized_conv2d().

The _quantized_conv2d() implementation demonstrates the import process:

def _quantized_conv2d(with_relu=False):
    def _impl(inputs, _):
        # inputs[0]: input tensor
        # inputs[1]: (weight, scale, zero_point, bias)
        # inputs[2-5]: stride, padding, dilation, groups
        # inputs[6]: output_scale
        # inputs[7]: output_zero_point
        # inputs[8]: input_scale
        # inputs[9]: input_zero_point
        
        conv_params = inputs[1]
        weight = conv_params[0]
        weight_scale = conv_params[1]
        weight_zero_point = conv_params[2]
        bias = conv_params[3]

        if len(conv_params) > 4:
            strides = conv_params[4]
            padding = conv_params[5]
            dilation = conv_params[6]
            groups = conv_params[7]
            output_scale = _expr.const(inputs[2])
            output_zero_point = _expr.const(inputs[3])
            input_scale = _expr.const(inputs[4])
            input_zero_point = _expr.const(inputs[5])
        else:
            strides = inputs[2]
            padding = inputs[3]
            dilation = inputs[4]
            groups = inputs[5]
            output_scale = _expr.const(inputs[6])
            output_zero_point = _expr.const(inputs[7])
            input_scale = _expr.const(inputs[8])
            input_zero_point = _expr.const(inputs[9])

        weight_shape = infer_shape(weight)
        kernel_size = (weight_shape[2], weight_shape[3])
        out_channels = weight_shape[0]

        if padding[0] != 0 or padding[1] != 0:
            pad_value = _get_scalar(input_zero_point)
            inp = _op.nn.pad(
                inputs[0],
                pad_width=((0, 0), (0, 0),
                           (padding[0], padding[0]), (padding[1], padding[1])),
                pad_value=float(pad_value),
            )
        else:
            inp = inputs[0]

        conv_out = relay.qnn.op.conv2d(
            inp,
            weight,
            input_zero_point,
            weight_zero_point,
            input_scale,
            weight_scale,
            kernel_size=kernel_size,
            dilation=dilation,
            strides=strides,
            padding=(0, 0),
            groups=groups,
            channels=out_channels,
        )

        return _do_bias_and_requantize(
            conv_out, bias, input_scale, weight_scale, output_scale, output_zero_point, with_relu
        )

    return _impl

Key functions include add_input_quant_params_to_op_inputs for extracting quantization parameters and get_quant_param_for_input for maintaining output parameter indices.

Quantized convolution can be decomposed into existing Relay operations based on tensor quantization principles:

A = scale_a * (QA - zp_A)

C(m, n) = Σ(k) (A(m, k) * W(n, k))

Σ(k) ([QA(m, k) - zp_a] * [QW(n, k) - zp_w])

# Expanded form:
Σ(k) QA(m, k) * QW(n, k)  # Term1
- Σ(k) zp_w * QA(m, k)    # Term2
- Σ(k) zp_a * QW(n, k)    # Term3
- Σ(k) * zp_a * zp_w      # Term4

Other operators follow similar quantization patterns:

• Multiplication: Quantized multiplication involves integer multiplication followed by requantization to handle scale differences.
• Requantization: Converts floating-point scale factors to integer operations using mantissa and exponent decomposition.
• Addition: Requires requantization before addition since scale factors cannot be factored out.
• Activation Functions: Complex activations use approximation or lookup tables; simple ones like ReLU6 can be implemented directly.
• Passive Quantization Operators: Operators like padding, ReLU, and max pooling don't change scale and require specific handling.

Graph Optimization

TVM applies various optimization passes during compilation. Key passes include:

Backend Codegen Integration

TVM's BYOC (Bring Your Own Codegen) framwork allows hardware vendors to integrate their code generation tools. The workflow includes:

1. Model loading and conversion to unified IR
2. Hardware-agnostic graph optimizations
3. Graph partitioning into host and accelerator sections

Graph partitioning can be based on patterns or costs. Pattern-based grouping uses sequence matching to identify operator sequences that can be mapped to hardware instructions. Cost-based partitioning considers resource constraints and decides which operators to offload.

After partitioning, hardware-specific optimizations are applied to accelerator regions, followed by code generation in various formats (JSON, C code, or custom representations).

Simulation Layer

Operator Simulation

When developing for new hardware, operator correctness must be verified through multiple steps:

1. Computational Correctness: Implement the operator using TVM DSL and compare results with PyTorch.
2. Compiler Simulator: Generate hardware-specific instructions and simulate on CPU to verify instruction correctness.
3. Hardware Simulation: Test instructions on hardware simulators (CSIM/RTL) using values from the compiler simulator.
4. Hardware Testing: Perform actual hardware validation after simulation passes.

Network Simulation

After individual operators are validated, network-level testing is performed using subnetwork benchmarks that capture common patterns. Full network testing then compares results with CPU execution or evaluates task-specific metrics like accuracy.