LoopTune: Optimizing Tensor Computations with Reinforcement Learning

TL;DR

LoopTune employs reinforcement learning to optimize tensor computations on CPUs, achieving significant speedups over existing auto-tuners and compiler frameworks with rapid tuning times.

Contribution

It introduces a novel RL-based compiler framework that optimizes tensor traversal and hardware-specific code generation for deep learning models.

Findings

LoopTune speeds up LoopNest by 3.2x.

Generates code 10x faster than TVM.

Tunes code in seconds.

Abstract

Advanced compiler technology is crucial for enabling machine learning applications to run on novel hardware, but traditional compilers fail to deliver performance, popular auto-tuners have long search times and expert-optimized libraries introduce unsustainable costs. To address this, we developed LoopTune, a deep reinforcement learning compiler that optimizes tensor computations in deep learning models for the CPU. LoopTune optimizes tensor traversal order while using the ultra-fast lightweight code generator LoopNest to perform hardware-specific optimizations. With a novel graph-based representation and action space, LoopTune speeds up LoopNest by 3.2x, generating an order of magnitude faster code than TVM, 2.8x faster than MetaSchedule, and 1.08x faster than AutoTVM, consistently performing at the level of the hand-tuned library Numpy. Moreover, LoopTune tunes code in order of seconds.