# UNIT: Unifying Tensorized Instruction Compilation

arXiv: https://arxiv.org/abs/2101.08458 · HTML: https://mlsanity.com/paper/2101.08458/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 48 references that are in the map.
- First posted on arXiv 21 Jan 2021.

## Builds on (most cited first, 10 of 10)

- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [Glow: Graph Lowering Compiler Techniques for Neural Networks](https://mlsanity.com/paper/1805.00907/) (2018)
- [TVM: An Automated End-to-End Optimizing Compiler for Deep Learning](https://mlsanity.com/paper/1802.04799/) (2018)
- [Optimizing CNN Model Inference on CPUs](https://mlsanity.com/paper/1809.02697/) (2018)
- [Relay: A New IR for Machine Learning Frameworks](https://mlsanity.com/paper/1810.00952/) (2018)
- [Efficient Execution of Quantized Deep Learning Models: A Compiler Approach](https://mlsanity.com/paper/2006.10226/) (2020)
- [Automatic Kernel Generation for Volta Tensor Cores](https://mlsanity.com/paper/2006.12645/) (2020)
- [PolyDL: Polyhedral Optimizations for Creation of High Performance DL primitives](https://mlsanity.com/paper/2006.02230/) (2020)

## Built on by (most cited first, 10 of 10)

- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [ArchGym: An Open-Source Gymnasium for Machine Learning Assisted Architecture Design](https://mlsanity.com/paper/2306.08888/) (2023)
- [Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs](https://mlsanity.com/paper/2510.08726/) (2025)
- [PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch](https://mlsanity.com/paper/2503.19779/) (2025)
- [REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Serving](https://mlsanity.com/paper/2506.01374/) (2025)
- [Joint Program and Layout Transformations to enable Convolutional Operators on Specialized Hardware based on Constraint Programming](https://mlsanity.com/paper/2104.04731/) (2021)
- [Automating Energy-Efficient GPU Kernel Generation: A Fast Search-Based Compilation Approach](https://mlsanity.com/paper/2411.18873/) (2024)
- [ALT: Boosting Deep Learning Performance by Breaking the Wall between Graph and Operator Level Optimizations](https://mlsanity.com/paper/2210.12415/) (2022)
- [HIERA: Workload-Aware Planning Across Implementation Spaces for GPU Kernel Optimization](https://mlsanity.com/paper/2608.21157/) (2026)

## Abstract (arXiv)

Because of the increasing demand for computation in DNN, researchers develope both hardware and software mechanisms to reduce the compute and memory burden. A widely adopted approach is to use mixed precision data types. However, it is hard to leverage mixed precision without hardware support because of the overhead of data casting. Hardware vendors offer tensorized instructions for mixed-precision tensor operations, like Intel VNNI, Tensor Core, and ARM-DOT. These instructions involve a computing idiom that reduces multiple low precision elements into one high precision element. The lack of compilation techniques for this makes it hard to utilize these instructions: Using vendor-provided libraries for computationally-intensive kernels is inflexible and prevents further optimizations, and manually writing hardware intrinsics is error-prone and difficult for programmers. Some prior works address this problem by creating compilers for each instruction. This requires excessive effort when it comes to many tensorized instructions. In this work, we develop a compiler framework to unify the compilation for these instructions -- a unified semantics abstraction eases the integration of new instructions, and reuses the analysis and transformations. Tensorized instructions from different platforms can be compiled via UNIT with moderate effort for favorable performance. Given a tensorized instruction and a tensor operation, UNIT automatically detects the applicability, transforms the loop organization of the operation,and rewrites the loop body to leverage the tensorized instruction. According to our evaluation, UNIT can target various mainstream hardware platforms. The generated end-to-end inference model achieves 1.3x speedup over Intel oneDNN on an x86 CPU, 1.75x speedup over Nvidia cuDNN on an NvidiaGPU, and 1.13x speedup over a carefully tuned TVM solution for ARM DOT on an ARM CPU.