Hands-on Practical: Inspecting Relay IR

Apache TVM offers a practical platform for inspecting high-level Intermediate Representations (IRs) and understanding the role of Static Single Assignment (SSA) form in a production environment. TVM utilizes Relay, a high-level, functional intermediate representation designed specifically for machine learning. Unlike a traditional computational graph which might only represent a Directed Acyclic Graph (DAG) of operations, Relay is a full programming language that supports control flow, recursion, and complex data structures, although most deep learning models utilize a static subset of these features.

In this section, we will ingest a model from a standard framework, convert it into Relay IR, and inspect the resulting text format. This process reveals how the compiler captures shape information, data types, and operator semantics before any hardware-specific optimization occurs.

Ingesting a Model into TVM

To observe the IR, we first need a source model. We will define a simple convolutional block in PyTorch, which serves as a representative example of a typical deep learning workload. This block includes a convolution, a batch normalization, and a ReLU activation.

import torch
import torch.nn as nn
import tvm
from tvm import relay

# Define a standard convolution block
class ConvBlock(nn.Module):
    def __init__(self):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(16)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

# Instantiate and trace the model
model = ConvBlock().eval()
input_shape = (1, 3, 224, 224)
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data)

The tracing step converts the dynamic PyTorch execution into a static graph representation (TorchScript). TVM requires this static definition to import the graph structure accurately.

Next, we invoke the TVM frontend to translate this TorchScript graph into a Relay IRModule. The IRModule is the central container in TVM that holds the functions and type definitions for the program.

# map input name to shape
shape_list = [("input_0", input_shape)]

# Import the graph to Relay
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)

print(mod)

Analyzing the Textual Representation

When you print the mod object, TVM outputs the text format of the Relay IR. The output will resemble the following structure. This text is not just a debug string; it is syntactically valid code in the Relay language.

def @main(%input_0: Tensor[(1, 3, 224, 224), float32]) {
  %0 = nn.conv2d(%input_0, meta[relay.Constant][0], padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]);
  %1 = nn.batch_norm(%0, meta[relay.Constant][1], meta[relay.Constant][2], meta[relay.Constant][3], meta[relay.Constant][4]);
  %2 = %1.0;
  nn.relu(%2)
}

This output explicitly demonstrates several compiler engineering principles discussed in earlier sections:

1. Function Definition: The entire graph is encapsulated in a main function. This aligns with the functional programming approach where a model is simply a function transforming inputs to outputs.
2. Strong Typing: Look at %input_0. It is annotated with a concrete shape (1, 3, 224, 224) and a data type float32. Unlike Python, Relay strictly enforces these types during compilation. If a dimension mismatch occurs between the convolution and the input, the compiler catches it here, long before code generation.
3. SSA Form: The variables %0, %1, and %2 represent intermediate tensors. Each variable is assigned exactly once. This structure simplifies dataflow analysis, allowing the compiler to easily track where data is produced and consumed.
4. Attributes: Operations like nn.conv2d carry attributes such as padding, channels, and kernel_size directly in the call. These are essential for the lowering phase to select the correct implementation.

Visualizing the Dataflow

While textual IR is precise, visualizing the dependencies as a graph helps in understanding the flow of data, especially for complex topologies. The IRModule describes a structure where data flows from the input parameters through a series of transformations.

The following diagram illustrates the structure of the Relay program we just generated. Notice how the weights (represented as Constants) feed into the operators alongside the data path.

The diagram depicts the dependency graph derived from the Relay IR. The blue node represents the input tensor, yellow nodes represent learned parameters (weights/biases), and green nodes represent computational operators.

One specific detail in the IR that often confuses new compiler engineers is the TupleGetItem node (shown as %2 = %1.0 in the text and the purple node in the diagram). The batch_norm operator in Relay returns a tuple containing three elements: the normalized tensor, the moving mean, and the moving variance. Since we are only interested in the normalized data for the forward pass, the %1.0 instruction extracts the zeroth element of that tuple. This explicit handling of multiple return values is a feature of Relay's expression-based design.

Manipulating the IR with Passes

Inspecting the IR is passive; the power of a compiler lies in transforming it. In TVM, transformations are applied via "Passes". A pass takes an IRModule and returns a new, optimized IRModule.

To see this in action, we can apply a simple optimization pass: Constant Folding. If our graph contained mathematical operations on constant values (e.g., 3 + 4), the compiler should compute the result (7) at compile time rather than runtime. While our current ConvBlock is mostly dynamic, applying a pass reveals how the infrastructure modifies the graph.

# Apply a transformation pass
seq = tvm.transform.Sequential([
    relay.transform.SimplifyInference(),
    relay.transform.FoldConstant(),
    relay.transform.DeadCodeElimination()
])

# The optimization occurs within a PassContext
with tvm.transform.PassContext(opt_level=3):
    opt_mod = seq(mod)

print("--- Optimized IR ---")
print(opt_mod)

Running this code might change the structure of the batch_norm. In inference mode, Batch Normalization can often be folded into the preceding convolution by updating the convolution's weights and bias. The SimplifyInference pass attempts to remove the batch_norm operator by replacing it with simpler arithmetic or merging it. When you inspect opt_mod, you may notice the nn.batch_norm call has disappeared, replaced by explicit broadcast, multiply, and add operations, or fused entirely into the conv2d if the backend supports it.

This hands-on inspection confirms that the IR is the ground truth for the compiler. It is the contract between the high-level framework and the low-level code generator. Understanding how to read and visualize this IR is a prerequisite for writing custom compiler passes or debugging performance regressions in deep learning models.