# CoSA: Scheduling by Constrained Optimization for Spatial Accelerators

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

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 71 references that are in the map.
- First posted on arXiv 5 May 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [You Only Look Once: Unified, Real-Time Object Detection](https://mlsanity.com/paper/1506.02640/) (2015)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Aggregated Residual Transformations for Deep Neural Networks](https://mlsanity.com/paper/1611.05431/) (2016)
- [WaveNet: A Generative Model for Raw Audio](https://mlsanity.com/paper/1609.03499/) (2016)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [SCNN: An Accelerator for Compressed-sparse Convolutional Neural Networks](https://mlsanity.com/paper/1708.04485/) (2017)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [The Architectural Implications of Facebook's DNN-based Personalized Recommendation](https://mlsanity.com/paper/1906.03109/) (2019)
- [Explaining How a Deep Neural Network Trained with End-to-End Learning Steers a Car](https://mlsanity.com/paper/1704.07911/) (2017)
- [TVM: An Automated End-to-End Optimizing Compiler for Deep Learning](https://mlsanity.com/paper/1802.04799/) (2018)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018)
- [Interstellar: Using Halide's Scheduling Language to Analyze DNN Accelerators](https://mlsanity.com/paper/1809.04070/) (2018)
- [Mind Mappings: Enabling Efficient Algorithm-Accelerator Mapping Space Search](https://mlsanity.com/paper/2103.01489/) (2021)
- [Communication-Optimal Tilings for Projective Nested Loops with Arbitrary Bounds](https://mlsanity.com/paper/2003.00119/) (2020)

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

- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks](https://mlsanity.com/paper/2107.06419/) (2021)
- [Sparseloop: An Analytical Approach To Sparse Tensor Accelerator Modeling](https://mlsanity.com/paper/2205.05826/) (2022)
- [Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators](https://mlsanity.com/paper/2505.18574/) (2025)
- [Deeploy: Enabling Energy-Efficient Deployment of Small Language Models On Heterogeneous Microcontrollers](https://mlsanity.com/paper/2408.04413/) (2024)
- [Performance Modeling and Workload Analysis of Distributed Large Language Model Training and Inference](https://mlsanity.com/paper/2407.14645/) (2024)
- [Enabling Unstructured Sparse Acceleration on Structured Sparse Accelerators](https://mlsanity.com/paper/2403.07953/) (2024)
- [SCAR: Scheduling Multi-Model AI Workloads on Heterogeneous Multi-Chiplet Module Accelerators](https://mlsanity.com/paper/2405.00790/) (2024)
- [DOSA: Differentiable Model-Based One-Loop Search for DNN Accelerators](https://mlsanity.com/paper/2509.10702/) (2025)
- [MTP: A Meaning-Typed Language Abstraction for AI-Integrated Programming](https://mlsanity.com/paper/2405.08965/) (2024)
- [eIQ Neutron: Redefining Edge-AI Inference with Integrated NPU and Compiler Innovations](https://mlsanity.com/paper/2509.14388/) (2025)
- [DNNFuser: Generative Pre-Trained Transformer as a Generalized Mapper for Layer Fusion in DNN Accelerators](https://mlsanity.com/paper/2201.11218/) (2022)
- [Demystifying Map Space Exploration for NPUs](https://mlsanity.com/paper/2210.03731/) (2022)
- [SALSA: Simulated Annealing based Loop-Ordering Scheduler for DNN Accelerators](https://mlsanity.com/paper/2304.12931/) (2023)
- [Tackling the Matrix Multiplication Micro-kernel Generation with Exo](https://mlsanity.com/paper/2310.17408/) (2023)
- [Weight Block Sparsity: Training, Compilation, and AI Engine Accelerators](https://mlsanity.com/paper/2407.09453/) (2024)
- [Obsidian: Cooperative State-Space Exploration for Performant Inference on Secure ML Accelerators](https://mlsanity.com/paper/2409.02817/) (2024)
- [Blockbuster, Part 1: Block-level AI Operator Fusion](https://mlsanity.com/paper/2505.07829/) (2025)
- [KAPLA: Pragmatic Representation and Fast Solving of Scalable NN Accelerator Dataflow](https://mlsanity.com/paper/2306.15676/) (2023)
- [CaMDN: Enhancing Cache Efficiency for Multi-tenant DNNs on Integrated NPUs](https://mlsanity.com/paper/2505.06625/) (2025)
- [A High-Level Compiler Integration Approach for Deep Learning Accelerators Supporting Abstraction and Optimization](https://mlsanity.com/paper/2507.04828/) (2025)
- [Mozart: A Chiplet Ecosystem-Accelerator Codesign Framework for Composable Bespoke Application Specific Integrated Circuits](https://mlsanity.com/paper/2510.08873/) (2025)

## Abstract (arXiv)

Recent advances in Deep Neural Networks (DNNs) have led to active development of specialized DNN accelerators, many of which feature a large number of processing elements laid out spatially, together with a multi-level memory hierarchy and flexible interconnect. While DNN accelerators can take advantage of data reuse and achieve high peak throughput, they also expose a large number of runtime parameters to the programmers who need to explicitly manage how computation is scheduled both spatially and temporally. In fact, different scheduling choices can lead to wide variations in performance and efficiency, motivating the need for a fast and efficient search strategy to navigate the vast scheduling space. To address this challenge, we present CoSA, a constrained-optimization-based approach for scheduling DNN accelerators. As opposed to existing approaches that either rely on designers' heuristics or iterative methods to navigate the search space, CoSA expresses scheduling decisions as a constrained-optimization problem that can be deterministically solved using mathematical optimization techniques. Specifically, CoSA leverages the regularities in DNN operators and hardware to formulate the DNN scheduling space into a mixed-integer programming (MIP) problem with algorithmic and architectural constraints, which can be solved to automatically generate a highly efficient schedule in one shot. We demonstrate that CoSA-generated schedules significantly outperform state-of-the-art approaches by a geometric mean of up to 2.5x across a wide range of DNN networks while improving the time-to-solution by 90x.