TASO: Memory & Latency Optimization for CNNs
- TASO is an ahead-of-time optimizer for CNN inference that uses an ILP formulation to jointly minimize latency and control memory usage on memory-constrained devices.
- It profiles 57 convolution primitive candidates and incorporates data layout conversions to accurately model inter-layer costs and ensure global optimization.
- TASO generates portable C code with a fixed schedule and exposes a latency-memory Pareto frontier, enabling deployment on embedded and mobile platforms under strict resource budgets.
TASO, in the sense of "Time and Space Optimization for Memory-Constrained DNN Inference," is an ahead-of-time optimizer for convolutional neural network inference on embedded and mobile platforms with tight RAM and latency budgets. It jointly selects, for each convolution layer, an implementation primitive and data layout, and it can also allocate a globally reusable workspace that upper-bounds per-layer temporary memory. The optimization is formulated as an integer linear program (ILP) driven by measured costs from the target device, so that TASO can minimize end-to-end latency under a memory budget, or minimize memory under a latency budget. Its output is a portable C implementation with fixed primitive and layout choices, together with a latency-memory Pareto frontier that makes deployment choices explicit (Wen et al., 2020).
1. Problem setting and design objective
TASO addresses a deployment regime in which CNN inference is dominated not only by arithmetic cost but also by transient memory and data movement. The motivating observation is that convolution layers dominate inference time, while many efficient convolution implementations require substantial temporary storage. In particular, lowering-based methods such as im2col can inflate temporary tensors by for kernels, which can push runtime memory into hundreds of MB; the paper reports that GoogleNet cannot run on devices with MB physical memory if implemented uniformly with im2col (Wen et al., 2020).
The system therefore treats latency and memory as coupled optimization targets rather than independent concerns. Its stated goal is to optimize end-to-end latency across the whole CNN while simultaneously controlling peak memory footprint, either in a whole-network-in-memory setting or in a layer-by-layer execution mode with a reusable workspace. This produces a family of Pareto-optimal operating points, allowing a deployment to be chosen under arbitrary system constraints without retraining, compression, or accuracy loss (Wen et al., 2020).
The underlying hardware assumptions are typical of edge deployment: limited DRAM capacity and bandwidth, small caches, and high energy cost for data movement. In this regime, temporary buffers and layout conversions can induce cache misses and memory pressure, so primitive selection cannot be done myopically on a per-layer basis. TASO’s optimization is network-level and explicitly accounts for inter-layer consequences.
2. Search space: primitives, layouts, and execution modes
TASO’s search space is defined over convolution-layer implementations. For each convolution layer , it considers a candidate set of implementation primitives drawn from a library of 57 candidates. A “primitive” denotes a concrete implementation algorithm together with its tuned variant and preferred layout. The candidate set spans direct convolution, lowering-based methods such as im2row or im2col plus GEMM, memory-efficient GEMM-based variants such as MEC and kern2row/kern2col, and Winograd. The paper discusses FFT convolution conceptually, but emphasizes that the short kernels typical of CNNs favor Winograd and GEMM-based methods in practice (Wen et al., 2020).
Data layout is part of the optimization because different primitives prefer different tensor organizations. TASO considers CHW (NCHW) and HWC (NHWC) orders, and each primitive carries its required input and output layout. When adjacent layers select incompatible layouts, TASO inserts explicit inter-layer layout conversions and charges both their latency and, if needed, their temporary memory against the optimization problem. This is essential because a locally fast primitive can become globally suboptimal once conversion overhead is included.
Two execution modes are central. In whole-network mode, TASO minimizes latency under a total memory budget, using the sum of selected per-layer temporary-memory costs. In layer-by-layer mode, it allocates a single global reusable workspace that upper-bounds all per-layer temporary memory. In that mode, TASO flips the roles of two main activation buffers between layers, streams weights, and uses additional small buffers for transforms and intermediates. This allows models to run even when the full network cannot be resident simultaneously (Wen et al., 2020).
The framework is not restricted to simple chains. Branching, fan-in, and fan-out topologies such as GoogleNet and ResNet are supported, and the pairwise conversion model extends along each edge. A plausible implication is that TASO’s main abstraction is the network dataflow graph plus a measured edge cost for layout mismatch, rather than any assumption of strictly sequential topology.
3. ILP formulation
TASO models a CNN with convolution layers. For each layer , the optimizer chooses exactly one primitive from . The principal binary decision variable is
where 0 means that primitive 1 is selected for layer 2. To linearize inter-layer conversion costs, TASO introduces
3
which encodes the adjacent choice of primitive 4 at layer 5 and primitive 6 at layer 7. In workspace mode, a scalar 8 denotes the maximum per-layer temporary memory usage across the network (Wen et al., 2020).
The measured parameters are:
- 9: latency of layer 0 with primitive 1,
- 2: temporary memory of layer 3 with primitive 4,
- 5: layout-conversion latency between adjacent layer choices,
- and optionally conversion memory 6.
The network-level latency objective under a whole-network memory budget 7 is
8
subject to
9
together with the primitive-assignment constraints and the standard linearization constraints linking 0 to adjacent 1 variables (Wen et al., 2020).
The dual formulation minimizes memory under a latency budget 2:
3
subject to total latency not exceeding 4.
For layer-by-layer execution, TASO instead minimizes peak workspace:
5
subject to the latency budget and
6
If conversions require extra temporary memory, TASO can also enforce
7
The paper notes that a weighted objective is also straightforward:
8
although the reported methodology primarily builds the frontier by sweeping hard memory or latency constraints rather than relying on a single weighted solve (Wen et al., 2020).
An optional explicit layout model with binaries 9 is possible, but in TASO layout is typically coupled to primitive selection, so the 0 variables suffice. This coupling is one reason the formulation remains directly tied to measured implementation behavior rather than to an abstract kernel taxonomy.
4. Measured cost model and ahead-of-time code generation
TASO does not rely on an analytical performance model. Instead, it profiles every candidate primitive for every convolution layer on the target hardware to obtain 1 and 2, and it profiles layout-conversion routines to populate the conversion matrix 3. In the reported implementation, this profiling is performed offline on an ODroid-XU4 with an Exynos-5422 SoC, using an ARM Cortex-A15 at 2.1 GHz with NEON (Wen et al., 2020).
This empirical modeling strategy is intended to capture microarchitectural effects that are difficult to represent analytically, including cache behavior, vectorization, instruction scheduling, and the performance characteristics of the platform’s BLAS or GEMM backend. The method is hardware-agnostic in the limited but practical sense that a new target only requires re-profiling and re-solving the ILP, rather than redesigning the optimization framework.
Code generation is ahead-of-time. TASO emits portable C code that sequences the selected primitives and inserts explicit layout-conversion calls where adjacent layers mismatch. In whole-network mode, all layer buffers are allocated statically. In layer-by-layer mode, the generated code allocates a global workspace, alternates two main activation buffers across layers, streams weights, and emits a fixed schedule for buffer flipping and auxiliary transform buffers. The paper characterizes the result as runnable on “any platform with a C compiler,” with performance portability derived from the measured cost model and the use of platform-efficient primitive implementations (Wen et al., 2020).
A common misconception is to regard TASO as a runtime autotuner. It is instead an ahead-of-time deployment optimizer: primitive identities, data layouts, conversion points, and workspace allocation are all fixed at compile time.
5. Pareto frontier and empirical results
The evaluation covers GoogleNet, AlexNet, VGG (VGG-D), ResNet, and SqueezeNet on ImageNet, measured on the ARM Cortex-A15 platform in CPU-only mode. The primitive library contains 57 candidates spanning direct convolution, im2row/im2col plus GEMM, memory-efficient GEMM-based variants, Winograd, and layout variants (Wen et al., 2020).
The central empirical claim is that global optimization matters most when memory is tight. Compared to a greedy memory-constrained heuristic baseline, TASO achieves up to 4 speedup on GoogleNet as memory budgets tighten. When memory is plentiful, the greedy method and TASO are similar; under tighter budgets, the greedy replacement heuristic misses globally optimal combinations of primitives and conversions that the ILP can identify (Wen et al., 2020).
Compared to a time-only solver identified as PBQP, TASO reduces memory requirement by 5 while sacrificing only about 6 of inference time when moving from the time-optimal solution to a nearby Pareto point. PBQP produces a single fastest-time configuration, whereas TASO exposes the full latency-memory frontier, including balanced points that are substantially more memory-efficient.
The comparison against uniform primitive choices clarifies why mixed per-layer selection is necessary. Direct-only execution gives the smallest memory footprint but is extremely slow; for GoogleNet, the paper reports a direct-only inference time of about 400 seconds in the scenario analyzed. Uniform im2col or im2row execution is much faster but often impractical because of memory usage, with GoogleNet exceeding 300 MB. Winograd-only execution is fast for small kernels but is not universally optimal. TASO’s mixed selection is faster than any single primitive choice at similar memory budgets and can match im2row speed at roughly half its memory (Wen et al., 2020).
In layer-by-layer mode, TASO also uncovers trade-offs invisible to time-only optimization. The paper reports configurations with nearly identical latency but significantly different workspace size, including an example where one point reduces approximately 10 MB relative to another, and configurations under tight memory of about 12 MB that differ by as much as 7 in speed.
The reported ILP solve times on an Intel Core i5 host are modest:
| Network | Decision variables | Solve time |
|---|---|---|
| AlexNet | 456 | 3.373 s |
| VGG | 741 | 5.259 s |
| SqueezeNet | 1710 | 8.835 s |
| GoogleNet | 3990 | 23.84 s |
These solve times are presented as a one-time ahead-of-time cost that is negligible relative to long-lived deployments (Wen et al., 2020).
6. Scope, limitations, and relation to adjacent systems
TASO is deliberately scoped. It targets static inference graphs and does not modify model semantics, retrain the model, quantize parameters, or prune the network. Its optimization focus is convolution layers, which dominate cost in the evaluated CNNs; extending the same methodology to pooling, fully connected layers, or normalization is described as feasible but as increasing modeling complexity (Wen et al., 2020).
Its main limitations arise from combinatorics and measurement fidelity. ILP solve time grows with the number of layers and candidate primitives, so larger search spaces may require pruning or decomposition. The quality of the solution depends on the accuracy of the measured cost model: firmware changes, BLAS versions, or CPU frequency scaling can shift the effective costs and necessitate re-profiling. Dynamic workloads with variable shapes or control flow are outside scope.
Relative to greedy primitive-selection heuristics, TASO’s distinctive property is global optimality under explicit constraints. Relative to PBQP-style time-only optimization, its contribution is the addition of memory modeling—either as total memory or as maximum reusable workspace—together with exposure of the full Pareto frontier rather than a single time-optimal point. Relative to systems such as TVM, Glow, Halide-based pipelines, Latte, Boda, and cuDNN heuristics, the paper positions TASO as an ahead-of-time joint time-memory optimizer with explicit modeling of primitive-dependent temporary memory and inter-layer layout conversions (Wen et al., 2020).
This suggests a precise characterization of TASO’s place in the optimization stack. It is neither a retraining method nor primarily a kernel generator. It is a deployment-time combinatorial optimizer that uses measured per-layer implementation costs to decide which already-available convolution primitive should run where, in which layout, and under what memory regime. Within that scope, its contribution is to make memory a first-class optimization variable rather than a post hoc constraint.