PaperlayerView exact revision on arXiv

Source-linked AI summary

Ansor: Generating High-Performance Tensor Programs for Deep Learning

Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, Joseph E. Gonzalez, Ion Stoica

arXiv:2006.06762v5cs.LGcs.NEcs.PFcs.PLstat.ML

TL;DR

Ansor targets the difficulty of finding efficient tensor programs across operators and hardware, where manual libraries demand engineering effort and existing searches restrict coverage. It automatically explores a hierarchical search space, fine-tunes complete programs with evolutionary search and a learned cost model, and schedules optimization across subgraphs. The reported maximum gains are 3.8× on Intel CPUs, 2.6× on ARM CPUs, and 1.7× on NVIDIA GPUs relative to state-of-the-art approaches.

  • Problem

    Manual kernel libraries require platform- and operator-specific engineering, while existing search methods restrict optimization coverage or exploration.

  • Method

    Ansor automatically constructs a hierarchical search space, samples complete programs, fine-tunes them with evolutionary search and a learned cost model, and prioritizes important subgraphs.

  • Results

    Ansor improves DNN execution performance by up to 3.8× on Intel CPU, 2.6× on ARM CPU, and 1.7× on NVIDIA GPU relative to state-of-the-art approaches.

  • Takeaways & Limitations

    Ansor finds high-performance programs outside existing search spaces and outperforms manual libraries and search-based frameworks across diverse neural networks and hardware platforms.

  • Takeaways & Limitations

    Ansor requires static, known shapes and supports dense operators but cannot optimize dynamic-shape graphs or sparse operators without redesigning its search space.

Abstract

from arXiv · show

High-performance tensor programs are crucial to guarantee efficient execution of deep neural networks. However, obtaining performant tensor programs for different operators on various hardware platforms is notoriously challenging. Currently, deep learning systems rely on vendor-provided kernel libraries or various search strategies to get performant tensor programs. These approaches either require significant engineering effort to develop platform-specific optimization code or fall short of finding high-performance programs due to restricted search space and ineffective exploration strategy. We present Ansor, a tensor program generation framework for deep learning applications. Compared with existing search strategies, Ansor explores many more optimization combinations by sampling programs from a hierarchical representation of the search space. Ansor then fine-tunes the sampled programs with evolutionary search and a learned cost model to identify the best programs. Ansor can find high-performance programs that are outside the search space of existing state-of-the-art approaches. In addition, Ansor utilizes a task scheduler to simultaneously optimize multiple subgraphs in deep neural networks. We show that Ansor improves the execution performance of deep neural networks relative to the state-of-the-art on the Intel CPU, ARM CPU, and NVIDIA GPU by up to $3.8\times$, $2.6\times$, and $1.7\times$, respectively.

1 Introduction

Ansor addresses the difficulty of generating high-performance tensor programs by expanding search coverage, improving exploration, and prioritizing important DNN subgraphs. Its evaluation reports gains over existing libraries and search-based systems across hardware platforms.

  • Vendor kernel libraries require substantial manual tuning for each hardware platform and operator.
  • Existing search approaches miss optimization combinations because manual templates and aggressive pruning restrict their search spaces.
  • Ansor automatically constructs a large hierarchical search space, samples complete programs, and fine-tunes them with evolutionary search and a learned cost model.
  • Ansor dynamically prioritizes subgraphs that are more likely to improve end-to-end DNN performance.
  • 3.8×, 2.6×, and 1.7× are Ansor's reported maximum execution-performance improvements on Intel CPU, ARM CPU, and NVIDIA GPU, respectively.

2 Background

Tensor-program generation requires navigating a large optimization space, while existing template-guided and sequential approaches constrain coverage or struggle to evaluate incomplete programs. Ansor instead separates high-level structures from low-level choices and tunes complete programs.

  • High-performance tensor programs require searching a complicated combinatorial space of tile structures, vectorization, parallelization, and related optimizations.
  • Template-guided search: Template-guided search depends on manually written templates and searches tunable parameters for specific shapes and hardware targets.
  • Sequential construction based search: Sequential construction search makes fixed-order decisions, but incomplete-program performance is difficult to estimate accurately with a cost model trained on complete programs.
  • Sequential construction based search: 50% pairwise accuracy and 0% top-k recall represent random guessing before programs become complete; both curves improve as completion increases.
  • Ansor’s hierarchical approach: Ansor automatically builds a hierarchical search space, samples complete programs, and fine-tunes them without estimating incomplete-program performance.

3 Design Overview

Ansor combines automatic hierarchical program sampling, performance fine-tuning, and resource allocation across subgraphs. Its system overview connects DNN partitioning, program generation, hardware measurement, and scheduler updates.

  • Ansor partitions input DNNs into smaller subgraphs and generates tensor programs for them.
  • Program sampler: The program sampler uses sketches for high-level structures and annotations for billions of low-level choices such as tile size, parallelization, and unrolling.
  • Performance tuner: The performance tuner iteratively fine-tunes sampled programs with evolutionary search and a learned cost model.
  • Gray arrows represent subgraph extraction and optimized-program generation, while green arrows represent profiling feedback that updates system components.
  • Task scheduler: The task scheduler uses gradient descent to allocate resources to subgraphs likely to improve end-to-end DNN performance.

4 Program Sampling

Ansor automatically generates sketches from computation DAGs and samples complete programs from a hierarchical search space. Its derivation rules cover diverse tensor-program structures while leaving low-level choices for efficient sampling.

  • Existing search spaces are limited by manual template enumeration and aggressive pruning of incomplete programs.These limitations prevent comprehensive coverage of effective optimization combinations.
  • The hierarchical representation separates high-level sketches from billions of low-level annotations such as tile sizes, parallelism, and unrolling.This separation enables flexible enumeration of structures and efficient sampling of implementation details.
  • The sampler accepts partitioned subgraphs represented as mathematical expressions, naive programs, or computational DAGs, then produces annotated complete programs.Generated sketches are built iteratively in topological order, with compute-intensive nodes tiled and simple element-wise nodes inlineable.
  • Ansor recursively applies flexible derivation rules to automatically enumerate sketches from a computation DAG.The derivation begins from a naive program and recursively produces terminal sketches through a queue of intermediate states.
  • Static analysis identifies predicates for inlining, data reuse, fusible consumers, and additional reduction parallelism when applying derivation rules.The analysis parses read/write patterns in mathematical expressions automatically.
  • Derivation rules support multi-level tiling, fusion, caching, reduction factorization, and user-defined extensions for specialized algorithms or hardware.Ansor can incorporate new rules for cases such as Winograd convolution and TensorCore-oriented tile structures.

5 Performance Fine-tuning

Ansor fine-tunes randomly sampled tensor programs through iterative evolutionary search, hardware measurement, and learned cost-model retraining. Tensor-program-specific mutations and node-based crossover expand optimization choices while preserving valid programs.

  • Randomly sampled programs provide broad search-space coverage, but their randomly chosen optimizations do not guarantee high quality.Performance fine-tuning addresses this quality gap.
  • Each fine-tuning iteration selects promising programs, measures them on hardware, and retrains the learned cost model with profiling data.The model becomes more accurate for the target hardware as measured data accumulates.
  • The learned cost model predicts program performance quickly, enabling comparison of tens of thousands of candidates before hardware measurement.It uses program features and aggregates predictions for innermost non-loop statements into a program score.
  • Evolutionary search selects programs using predicted fitness, then applies tensor-program-specific mutations and crossover to generate new candidates.The initial population combines random samples with high-quality programs from previous measurements.
  • Node-based crossover combines rewriting histories by selecting a parent per DAG node and verifying the merged program’s validity.Node-level granularity reduces dependency violations compared with arbitrarily combining rewriting steps.
  • Ansor’s evolutionary operations handle complicated dependencies and out-of-order program modifications beyond fixed grid-like parameter spaces.This distinguishes them from parameter-grid search in TVM and FlexTensor and sequential unfolding rules in Halide.

6 Task Scheduler

Ansor’s task scheduler allocates tuning resources across DNN subgraphs according to end-to-end objectives rather than optimizing every subgraph equally. It combines objective functions, gradient-based prioritization, task similarity, and exploration.

  • Tuning every DNN subgraph equally can waste time when a subgraph is not a bottleneck or has little room for improvement.Ansor therefore allocates different amounts of tuning time to different subgraphs.
  • The scheduler iteratively selects a task, generates promising programs for its subgraph, measures them on hardware, and treats each iteration as one time unit.A single DNN may require dozens of tasks, such as 29 for ResNet-50.
  • Users can define objectives for reducing latency, meeting latency requirements, minimizing tuning time, or optimizing multiple DNNs.Custom objective functions are also supported.
  • For a single DNN, end-to-end latency is approximated by a weighted sum of subgraph latencies, with weights equal to task appearance counts.The formulation minimizes f(g1(t),g2(t),...,gn(t)) over the allocation vector.
  • The scheduler uses gradient descent to estimate which task receives the next resource unit, incorporating allocation history and task similarity.The estimate also uses task FLOP counts, achievable throughput, and tunable prediction weights.
  • After round-robin warm-up, Ansor repeatedly updates the allocation vector and uses ε-greedy selection to retain random exploration.For single-DNN latency, it can leave a task when additional iterations stop reducing its latency.

7 Evaluation

Ansor is evaluated against search frameworks and manual libraries across operators, subgraphs, and complete DNNs on Intel, NVIDIA, and ARM hardware. It achieves strong performance and search efficiency, with gains attributed to broad search coverage, fine-tuning, and task scheduling.

  • Single Operator Benchmark: Ansor performs best or equally best across all operator and batch-size settings, outperforming existing search frameworks by 1.1−22.5×.Its gains arise from a larger search space and more effective exploration, including optimizations unavailable to competing frameworks.
  • Ablation Study: Removing fine-tuning or restricting the search space significantly decreases final performance, while random sampling alone cannot beat AutoTVM within the given time budget.The limited-space configuration performs worst because the best programs are excluded from that space.
  • Subgraph Benchmark: Ansor outperforms manual libraries and other search frameworks on subgraphs by 1.1−14.2× across both evaluated platforms.FlexTensor performs well for single operators but is less advantageous on subgraphs because it lacks operator-fusion support.
  • End-to-End Network Benchmark: Ansor performs best or equally best across all end-to-end network cases, improving over the best alternative by up to 3.8× on Intel, 2.6× on ARM, and 1.7× on NVIDIA.It matches or outperforms AutoTVM in all cases, with 1.0−21.8× speedup.
  • End-to-End Network Benchmark: 40% improvement comes from Ansor’s layout rewrite for ResNet-50 on the Intel CPU.Ansor does not perform global layout search but rewrites weight-tensor layouts and packs them into more tiling levels.
  • Network Search: The task scheduler prioritizes performance bottlenecks, such as subgraphs containing 3x3 convolution, improving both search efficiency and final performance.Ansor allocates time across subgraphs instead of assigning each a fixed, equal measurement budget.
  • Search Time: Ansor matches AutoTVM with an order of magnitude less search time on the Intel CPU network benchmark.Savings come from task scheduling, fine-tuning, broad optimization coverage, low search overhead, and efficient measurement.
  • Cost Model: The learned cost model achieves 0.079 RMSE, 0.958 R2 correlation, 0.851 pairwise comparison accuracy, and 0.624 recall@30 on the test set.Measured-throughput points cluster near the diagonal, indicating accurate predictions.

8 Related Work

Related work includes scheduling-language systems, polyhedral compilers, graph-level optimization, and general auto-tuning frameworks. These approaches differ in how they represent programs, define search spaces, and explore optimization choices.

  • Scheduling Languages: Halide and TVM provide scheduling languages for manual optimization and automatic search, while AutoTVM and FlexTensor use template-guided exploration.Halide’s latest evaluated auto-scheduler combines beam search with a learned cost model.
  • Polyhedral Compilation: Polyhedral compilers formulate optimization as affine loop transformations, often using integer linear programming to reduce data-reuse distances.Tiramisu and TensorComprehensions apply polyhedral compilation to deep learning workloads.
  • Graph-Level Optimization: Graph-level optimization treats operators as graph units and applies transformations such as layout optimization, operator fusion, constant folding, and auto-batching.These optimizations operate at graph level without changing operators’ internal implementations.
  • Search-Based Compilation: General auto-tuning systems explore spaces using random search, multiarmed bandits, or domain-specific tuning, but may depend on user-specified search spaces.Ansor instead generates tensor programs with loop nests and automatically constructs its search space.

9 Limitations and Future work

Ansor has several scope and optimization limitations, including inability to handle dynamic shapes or sparse operators and incomplete exploitation of platform-specific instructions.

  • Scope limitations: Ansor requires static, known tensor shapes and cannot optimize computational graphs with dynamic shapes.Dynamic-shape support is identified as future work.
  • Scope limitations: Ansor supports dense operators but not sparse operators such as SpMM used in sparse and graph neural networks.Supporting sparse operators would require redesigning the search space, although much of Ansor could be reused.
  • Hardware-specific optimization: Ansor performs high-level program optimization but relies on LLVM and NVCC for platform-dependent optimizations.This limits its handling of special instructions for mixed-precision and low-precision operators.
  • Hardware-specific optimization: Ansor does not effectively utilize Intel VNNI, NVIDIA Tensor Core, and ARM Dot instructions through current off-the-shelf code generators.These instruction-specific limitations affect mixed-precision and low-precision operators.

10 Conclusion

Ansor is an automated tensor-program search framework that explores broad optimization spaces and prioritizes performance bottlenecks. It finds programs beyond existing approaches and improves performance across neural networks and hardware platforms.

  • Conclusion: Ansor efficiently explores a large search space and prioritizes performance bottlenecks to find programs outside existing approaches.The framework targets automated generation of high-performance tensor programs for deep neural networks.
  • Conclusion: Up to 3.8× performance improvement is reported over existing manual libraries and search-based frameworks across diverse neural networks and hardware platforms.The reported result is a maximum improvement rather than a uniform gain across all evaluations.
  • Conclusion: Ansor is integrated into the Apache TVM open-source project.The paper presents this integration as part of the framework’s deployment context.

A Gradient Approximation for the Task Scheduler

The task scheduler’s gradient approximation assumes task cost depends primarily on allocated resources, predicts unavailable values, and uses program features for cost modeling.

  • Gradient approximation: The gradient approximation assumes that a task’s best cost depends only on the resource units allocated to it.The paper acknowledges that shared cost-model training data can violate this assumption.
  • Gradient approximation: Different resource allocations can produce different training data and therefore different shared cost models.This is identified as a reason the approximation may not hold exactly.
  • Gradient approximation: The scheduler predicts the unknown future task cost using an optimistic estimate and similarity among structurally similar subgraphs.The prediction combines these factors through a weighted approximation controlled by parameters described in the surrounding derivation.
  • Cost-model features: The cost model uses categorical and numerical features extracted from innermost non-loop statements in complete tensor programs.The same feature set is used for CPU and GPU programs, with 164 features per statement.
  • Cost-model features: Features describe arithmetic, vectorization, unrolling, parallelization, GPU thread binding, memory access, reuse, allocation, and loop structure.Memory-related features include accessed bytes, unique bytes, cache lines, reuse, and arithmetic intensity.
  • Cost-model features: The feature extraction also includes reuse distance, reuse counts, stride, accessed-bytes ratios, allocation properties, and outer-loop characteristics.These features complement operation and scheduling descriptors in the statement representation.

C Shape Configurations in the Evaluation

The evaluation covers diverse tensor operators and neural-network components by encoding each workload with shape parameters specific to its operator structure.

  • Convolution configurations: C1D, C2D, and C3D configurations encode convolution dimensions, channels, kernel size, stride, and padding.C3D additionally includes depth alongside height and width.
  • Convolution configurations: GRD and DIL configurations extend convolution shapes with groups or dilation, respectively.Both retain spatial, channel, kernel, stride, and padding parameters.
  • Convolution configurations: DEP configurations specify height, width, channel, kernel size, stride, and padding for depthwise convolution.The channel parameter replaces separate input and output channel fields in this format.
  • Specialized convolutions: T2D and CAP configurations represent transposed and capsule 2D convolutions, with CAP adding capsule size.Their formats include spatial dimensions, channels, kernel size, stride, and padding.
  • Convolution configurations: ConvLayer uses a convolution-layer format containing height, width, input and output channels, kernel size, stride, and padding.It provides a named layer-level configuration alongside the operator-specific formats.
  • Non-convolution configurations: TBS represents transposed, batch-matrix-multiplication, and softmax operations in multi-head attention using sequence length, head count, and hidden dimension.This configuration broadens the evaluation beyond convolutional operators.
Loading 2006.06762v5…