Graph Serialization

The core component is the directed acyclic graph (DAG) defining the model's architecture. This includes:

• Nodes/Operators: Represent mathematical operations (e.g., Conv2D, MatMul, ReLU).
• Edges/Tensors: Represent the multi-dimensional data arrays flowing between operators.
• Attributes: Static parameters for each operator (e.g., kernel size, stride, padding). This structural definition is abstract and hardware-agnostic, describing what to compute, not how.

Serialization embeds the learned parameters of the model. This is the state acquired during training.

• Weights: The core matrices and filters for layers like convolutions and linear transformations.
• Biases: The additive constants applied in many layers. These are typically stored as raw binary blobs of floating-point or quantized integer data, often compressed. Their exact memory layout must be preserved for deterministic execution.

This component provides the semantic context for the raw tensor data.

• Data Types: Precision of each tensor (e.g., FP32, FP16, INT8).
• Tensor Shapes: Static or dynamic dimensions (e.g., [batch, channels, height, width]).
• Tensor Names: Human-readable identifiers for inputs, outputs, and intermediates, crucial for runtime binding. This metadata enables the runtime to allocate correct memory buffers and validate graph execution.

Serialization formats include provenance information to ensure compatibility and traceability.

• Producer Information: The framework and version that created the graph (e.g., PyTorch 2.1, TensorFlow 2.13).
• Format Version: The specific version of the serialization protocol itself.
• Model Signature: A formal definition of the model's inputs and outputs, including expected names and types. This is critical for deployment toolchains.

Advanced serialization may include hardware-aware directives to guide the deployment runtime.

• Fused Operator Patterns: Annotations indicating groups of nodes that are candidates for kernel fusion.
• Preferred Data Layouts: Hints for memory format (e.g., NCHW vs. NHWC) to avoid costly runtime transformations.
• Quantization Parameters: Scales and zero-points for quantized tensors. These hints bridge the gap between the abstract graph and efficient hardware execution.

Major ML frameworks provide their own serialization formats, often combining graph structure and weights into a single archive.

• PyTorch TorchScript (.pt/.pth): Serializes a PyTorch model's code and state via Python's pickle module, often combined with TorchScript's intermediate representation for graph capture.
• TensorFlow SavedModel: A directory containing a protobuf file (saved_model.pb) storing the MetaGraphDef, and a subdirectory (variables/) storing checkpointed weights. It's the standard for serving.
• JAX Flax Checkpoints: Often use msgpack or TensorFlow's tf.train.Checkpoint format via orbax-checkpoint to serialize the model state dictionary (parameters) separately from the functional model code.

These formats are typically the starting point before conversion to more deployment-oriented formats like ONNX or TFLite.

Serialization for NPU targets involves unique requirements beyond simple persistence.

• Hardware-Specific Annotations: The serialized graph may need to embed vendor-specific attributes (e.g., #tpu, #npu.core_assignment) to guide later compilation stages.
• Quantization Information: Must preserve quantization parameters (scale/zero-point) for INT8/INT4 weights and activations, often stored as metadata alongside tensors.
• Static Graph Requirements: Most NPU compilers require a static graph—all tensor shapes must be inferable at serialization/compile time. Dynamic shapes complicate this process.
• Separated Weights: For very large models, weights may be serialized in a separate, indexable file (like safetensors) to allow memory-mapped loading, reducing RAM pressure on the host before NPU transfer.

The serialized artifact is the input to the vendor's proprietary compiler toolchain (e.g., NVIDIA's TensorRT, Google's Edge TPU Compiler, Qualcomm's AI Engine Direct).

An Intermediate Representation is a compiler's internal, abstract data structure used to represent a program between its high-level source form and low-level machine code. For neural networks, the IR captures the computational graph, enabling machine-independent analysis and transformations.

• Purpose: Serves as a common language for multiple optimization passes.
• Examples: LLVM IR, MLIR dialects, ONNX graph representation.
• Relation to Serialization: The optimized IR is the primary input to the serialization process, which encodes it into a persistent byte stream format like ONNX or TensorFlow SavedModel.

Graph canonicalization is a compiler transformation that rewrites a computational graph into a standard, simplified form. It eliminates syntactic variations (e.g., different operator names for the same function) to make subsequent analysis and optimization passes more effective and predictable.

• Key Benefit: Ensures deterministic optimization outcomes.
• Common Actions: Standardizing operator names, flattening nested structures, removing identity operations.
• Pipeline Order: Typically occurs early in the compilation pipeline, before major optimizations and serialization.

Static shape inference is a compiler analysis that determines the dimensions (shape) of all tensors in a computational graph at compile time. This is crucial for memory planning and enabling certain optimizations before execution.

• Enables: Precise static memory allocation, kernel selection, and validation.
• Contrast with Dynamic Shapes: Graphs with fully static shapes are often more efficiently optimized and serialized.
• Serialization Role: Inferred shape metadata is a key component serialized alongside the graph structure, informing the runtime about buffer sizes.

Memory planning is a compiler optimization pass that allocates memory buffers for tensors in a computational graph. It aims to minimize peak memory usage through techniques like buffer reuse and in-place operations.

• Objective: Reduce device memory footprint for deployment on constrained hardware (e.g., NPUs, mobile).
• Techniques: Liveness analysis, in-place optimization, shared memory allocation.
• Serialization Output: The memory plan may be explicitly encoded in the serialized format or left for the runtime to recompute using the serialized graph metadata.

Graph lowering is the process of transforming a high-level, abstract computational graph representation into a lower-level, more hardware-specific representation. This occurs through a series of legalization and conversion passes.

• Purpose: Bridge the gap between framework-level operators and hardware-specific kernels or primitives.
• Multi-Stage Process: May progress from a framework graph (e.g., PyTorch) to a vendor-specific graph (e.g., for a specific NPU).
• Final Stage: The fully lowered graph, targeting a specific execution backend, is what is ultimately serialized for deployment.

Ahead-Of-Time compilation is a strategy where source code or an intermediate representation is fully compiled into a native machine code binary before execution. This eliminates runtime compilation overhead, providing predictable performance.

• Contrast with JIT: AOT trades runtime flexibility for startup speed and determinism.
• NPU Context: Essential for embedded and edge NPU deployment where runtime compilers are unavailable or too costly.
• Serialization Link: The final output of an AOT compilation pipeline for an NPU is often a serialized graph and the compiled kernel binaries, packaged together into a deployable artifact.