DOCSEARCHbyAlgolia
Submit Docs

tract

repository·main·Indexed 24 days ago

https://github.com/sonos/tract

A neural-network inference engine developed by Sonos for loading, optimizing, and running models such as ONNX, NNEF, and TensorFlow. It supports a wide range of platforms, including embedded ARM CPUs, GPUs, and WebAssembly. Features include a programmatic interface for causal LLMs with speculative decoding, a CLI tool for model translation, and support for dynamic shapes and symbols.

Tokens
45.2K
Snippets
96
Records
268
Agent score
85%

What's inside tract

  1. Overview of tract-linalg

    main
    The tract-linalg crate provides low-level, architecture-dependent optimizations used by tract-core. Despite its name, it is not a general-purpose linear algebra library, but rather a collection of highly optimized kernels for specific operations like matrix multiplication and activation functions.

  2. Overview of x86_64 int8 GEMM kernels

    main
    The linalg package provides a family of int8 (i32-accumulator) matrix-multiply (GEMM) kernels for x86_64 architectures. These kernels form a throughput cascade ranging from portable AVX2 emulation to Intel AMX, with AVX-512-VNNI in between. The optimal kernel is selected at runtime based on CPUID and an einsum kernel scorer.

  3. Understand Symbolic Shapes and TDim

    main

    In tract, dimensions that depend on runtime inputs (like batch size or sequence length) are represented as symbolic expressions using the TDim algebraic data type rather than concrete usize values. This allows the optimizer to perform algebraic reasoning (e.g., proving S >= 0) even when exact values are unknown.

    TDim (found in tract_data::dim) is an enum containing:

    • Val(i64): A known integer.
    • Sym(Symbol): A named symbolic atom (e.g., B, S).
    • Arithmetic: Add, Mul, MulInt(i64, Box<TDim>), and Div(Box<TDim>, u64).
    • Reductions: Min, Max, and Broadcast (dimension-wise broadcast rule).
    • Comparisons: Ge and Eq (evaluate to 0 or 1).

  5. Understand the Tract Graph structure

    main

    Tract represents Neural Networks as a Directed Acyclic Graph (DAG) using Graph and BaseNode.

    • Graph: Contains a list of BaseNode objects, along with inputs and outputs defined by OutletId.
    • BaseNode: Represents an operator (op). It contains a list of inputs (inlets) and outputs (outlets).
    • OutletId: A unique identifier for a wire, consisting of a node index and a slot (the output number of the Op).
    • Outlet: Represents a connection point. An outlet can be connected to multiple inlets, but an inlet can only be set by one unique outlet.

    To run a network, tract executes nodes in an order where a node is only run once all its input values are known. The process follows:

    1. Set values for Source nodes (provided by the caller).
    2. Compute outputs for unvisited nodes whose inputs are all known.
    3. Extract model outputs from the wires specified in Graph.outputs.
    pub struct Graph<F, O> {
    pub nodes: Vec<BaseNode<F, O>>,
    pub inputs: Vec<OutletId>,
    pub outputs: Vec<OutletId>,
    /* [...] */
}

pub struct BaseNode<F, O> {
    pub inputs: Vec<OutletId>,
    pub op: O,
    pub outputs: Vec<Outlet<F>>,
}

pub struct OutletId {
    pub node: usize,
    pub slot: usize,
}

pub struct Outlet<F: Fact + Hash> {
    pub fact: F,
    pub successors: Vec<InletId>,
}

  6. Understand the MatMul kernel cost model

    main

    Tract uses a LinearCostModel to select the optimal microkernel for MatMatMul operations on a specific CPU. The model predicts runtime for a shape (m, k, n) using the formula:

    time ≈ a · padded_work + b · n_tiles + c + restream · a_restream

    Where:

    • a: Inverse steady-state throughput.
    • b: Per-tile setup cost.
    • c: Fixed call overhead.
    • restream: A shared coefficient capturing the cost of re-streaming weights (A) between layers, which is captured via cold-cache probes.
    • padded_work, n_tiles, and a_restream are derived from the input shape and kernel parameters (mr, nr).

    Target-specific coefficient tables are stored as generated Rust files (e.g., linalg/src/arm64/cortex_a53_linear.rs) and are dispatched based on the CPU type detected at runtime.

  7. Use the api/rs crate for client applications

    main
    When building applications, examples, or language bindings, you should use the api/rs crate. This is the stable public API surface. Do not use internal crates such as core, nnef, or onnx directly, as they do not provide a stable API.

  8. Dumping and loading models using OPL

    main

    OPL is a tract-specific format based on NNEF. Operators can be dumped to OPL and loaded from it. While some operators use standard NNEF forms, many ONNX and TensorFlow operators require extensions within OPL.

    Each OPL module (such as nnef, pulse-opl, or onnx-opl) maintains a Registry that contains both OPL loaders and OPL dumpers.

    • Dumping: Uses a from_tract mapping that connects a Rust TypeId (for a TypedOp) to a FromTract function. This function modifies an IntoAst object to store the operator's representation, potentially adding NNEF fragments to the document.
    • Loading: Uses the registry to find the appropriate loader for the operator defined in the OPL file.
    pub struct Registry {
    pub id: String,
    pub fragments: HashMap<String, FragmentDef>,
    pub primitives: HashMap<String, (Vec<ast::Parameter>, ToTract)>,
    pub unit_element_wise_ops: Vec<(String, Box<dyn ElementWiseMiniOp>)>,
    pub element_wise_ops: Vec<(String, TypeId, FromTract, Vec<ast::Parameter>, ToTract)>,
    pub binary_ops: Vec<(String, Box<dyn BinMiniOp>)>,
    pub from_tract: HashMap<TypeId, FromTract>,
}

/// Function type for converting a tract model to OPL
pub type ToTract = fn(&mut ModelBuilder, &ResolvedInvocation) -> TractResult<TVec<OutletId>>;

/// Function type for loading an operator from OPL into a tract model
pub type FromTract = fn(&mut IntoAst, node: &TypedNode) -> TractResult<Option<Arc<RValue>>>;

  9. Benchmark AMX and VNNI Kernels

    main

    To measure performance, use taskset to pin a core for stability. Run the following benchmarks to compare AMX, VNNI, and AVX2 implementations:

    1. int8 (AVX2 vs VNNI 8x8 vs AMX 8x8 vs AMX 16x16):

    taskset -c 2 cargo bench -p tract-linalg --bench amx_i32

    2. f32 via bf16 (FMA 16x6 vs AVX-512 16x12 vs AMX-BF16 16x16):

    taskset -c 2 cargo bench -p tract-linalg --bench amx_f32

    3. VNNI 16x16 isolation (AVX2 vs VNNI 8x8 vs VNNI 16x16):

    taskset -c 2 cargo bench -p tract-linalg --bench vnni_i32

    Benchmark reports are generated as HTML files in target/criterion/.

    taskset -c 2 cargo bench -p tract-linalg --bench amx_i32

  10. Load a model and specify input shapes

    main

    To load a model and set its input shapes (equivalent to using InferenceFact and with_input_fact in the API), use the -i flag. The format for shapes is batch,channels,height,width (or similar depending on the model) followed by the data type.

    Example for an ONNX model with input shape 1,3,224,224 and f32 type:

    tract mobilenetv2-7.onnx -i 1,3,224,224,f32 dump
    tract mobilenetv2-7.onnx -i 1,3,224,224,f32 dump

  11. Reproduce AMX benchmarks

    main

    To reproduce or extend the AMX benchmark results, you must use an AMX-capable host (such as Intel Sapphire Rapids, Emerald Rapids, Granite Rapids Xeon, or Xeon Max) and follow the instructions in linalg/AMX_BENCH_RUNBOOK.md.

    Benchmark Environment Requirements:

    • ISA: amx_tile, amx_int8, amx_bf16 + AVX-512-VNNI.
    • Kernel: $\ge$ 5.16 (e.g., 6.18.5).
    • Method: Use cargo bench with default criterion sampling, pinned to a specific core (e.g., taskset -c 2) to ensure stability.