Flow-Graph-Aware Tiling and Rescheduling for Memory-Efficient On-Device Inference

Abstract

With the increasing popularity of artificial intelligence (AI) applications, deep neural networks (DNNs) are in demand for on-device serving in various real-life fields. Running DNN inference on a resource-constrained edge device requires aggressive memory optimization. While several recent tiling-based techniques reduce peak memory usage by partitioning large tensors into micro-tensors, they are specialized for the MCU environment and do not provide scalability to various edge platforms. Moreover, they greedily search for the target of tiling without considering the memory flow across the model while partitioning. In this paper, we propose OKO, a compiler-based optimization technique that minimizes peak memory usage by considering both tiling and the corresponding operation rescheduling. OKO estimates the memory savings from the tiling method based on the lifetime and dependencies of the tensor, and then algorithmically selects the optimal tiling strategy. It further maximizes the reuse of memory spaces by efficiently reordering operations and immediately releasing unnecessary tensors. Evaluations on various edge devices show that OKO achieves effective memory savings of up to 80% and an average of 59% with no loss of accuracy and negligible overhead, supporting memory-efficient inference across a broad range of target devices.