Graph Rewriting

Graph Rewriting in the context of TensorRT-LLM is a process that involves manipulating the structure of a neural network at a lower level, specifically at the ILayer/INetworkDefinition level in TensorRT.

This technique is particularly useful in optimizing and transforming neural network models for efficient execution, especially when using NVIDIA's TensorRT, a high-performance deep learning inference optimizer and runtime library.

When to Use Graph Rewriting?

Graph Rewriting is used in situations where fine-grained control and manipulation of the network at the layer level are required, particularly after the network has been defined. It is different from Module Rewriting, which operates at a higher level (before the network is converted into TensorRT's graph format). Graph Rewriting is useful:

1. When only ILayer/INetworkDefinition is available: If you're working directly with these lower-level constructs.

2. For complex manipulations: That cannot be efficiently or feasibly done at the Module level, such as layer fusion, or when dealing with nested control flow.

Key Concepts in Graph Rewriting

1. Tensor-Related Methods: These methods allow manipulation of tensors, including getting parent layers of a tensor, consumer layers, and replacing tensors in consumer layers.

2. FLayerInfo and FLayerInfoMemo: These are used to store and retrieve high-level information about layers, especially useful for layers implemented as TensorRT plugins. They maintain the original input and attribute information of layers.

3. Pattern and Pattern Manager:

   • PatternRewriter: For defining a rewriting pattern that alters the network.

   • PatternAnalyzer: For defining an analysis pattern that collects information from the network.

   • RewritePatternManager: Manages multiple rewriting patterns.

   • AnalysisPatternManager: Manages multiple analysis patterns.

4. @record_signature Decorator: Used to record the high-level signature (FLayerInfo) for functionals, which is crucial for Graph Rewriting when analyzing or rewriting certain functionals.

5. Workflow for Defining a Graph Rewriting Pattern: Typically involves defining a class that inherits from PatternRewriter or PatternAnalyzer and implementing the match and rewrite (or analyze) methods to specify how the network should be transformed or analyzed.

Practical Examples of Graph Rewriting

1. Layer Replacement: Replacing one type of layer with another (e.g., replacing a sum layer with a subtract layer) while maintaining the connections in the network.

2. Layer Fusion: Combining multiple layers into a single, more efficient layer, often used to reduce computational overhead.

3. Optimization for Specific Hardware: Tailoring the network structure to leverage specific features of hardware accelerators, like GPUs, for more efficient execution.

4. Enabling Advanced Features: Modifying layers to enable advanced features like dynamic batch sizes, mixed precision, or remove-padding modes in specific layers.

Importance in Neural Network Optimization

Graph Rewriting is a crucial step in optimizing neural networks for deployment, particularly in scenarios where high throughput and low latency are critical, such as real-time applications or edge computing. By transforming and optimizing the network at a granular level, it's possible to achieve significant improvements in performance on specific hardware architectures.

In summary, Graph Rewriting in TensorRT-LLM offers a powerful toolset for deep manipulation and optimization of neural networks, allowing for fine-grained control and customization to achieve high-performance inference.