Graph Abstraction for ML Models

Graph abstraction is a key technique used in machine learning (ML) compilers to represent and reason about the structure and data flow of ML models. By abstracting the model into a graph representation, the compiler can perform various optimizations to improve performance and efficiency. This tutorial will cover the basics of graph abstraction, its key elements of Relax IR, and how it enables optimization in ML compilers.

What is Graph Abstraction?

Graph abstraction is the process of representing an ML model as a directed graph, where the nodes represent computational operations (e.g., matrix multiplication, convolution) and the edges represent the flow of data between these operations. This abstraction allows the compiler to analyze the dependencies and relationships between different parts of the model.

from tvm.script import relax as R

@R.function
def main(
    x: R.Tensor((1, 784), dtype="float32"),
    weight: R.Tensor((784, 256), dtype="float32"),
    bias: R.Tensor((256,), dtype="float32"),
) -> R.Tensor((1, 256), dtype="float32"):
    with R.dataflow():
        lv0 = R.matmul(x, weight)
        lv1 = R.add(lv0, bias)
        gv = R.nn.relu(lv1)
        R.output(gv)
    return gv

Key Features of Relax

Relax, the graph representation utilized in Apache TVM’s Unity strategy, facilitates end-to-end optimization of ML models through several crucial features:

• First-class symbolic shape: Relax employs symbolic shapes to represent tensor dimensions, enabling global tracking of dynamic shape relationships across tensor operators and function calls.

• Multi-level abstractions: Relax supports cross-level abstractions, from high-level neural network layers to low-level tensor operations, enabling optimizations that span different hierarchies within the model.

• Composable transformations: Relax offers a framework for composable transformations that can be selectively applied to different model components. This includes capabilities such as partial lowering and partial specialization, providing flexible customization and optimization options.

These features collectively empower Relax to offer a powerful and adaptable approach to ML model optimization within the Apache TVM ecosystem.