Value Function Based Performance Optimization of Deep Learning Workloads
This work provides a method to rapidly optimize the performance of deep learning workloads, which is crucial for developers and researchers deploying machine learning models.
The paper addresses the challenge of finding efficient schedules for deep neural network implementations by modeling the scheduling problem as a sequence of optimization choices. They developed a technique to predict the expected performance of partial schedules, which allowed them to greedily identify efficient schedules that improved throughput by 2.6x over Halide and 1.5x over TVM. Their method is also significantly faster, completing in seconds instead of hours.
As machine learning techniques become ubiquitous, the efficiency of neural network implementations is becoming correspondingly paramount. Frameworks, such as Halide and TVM, separate out the algorithmic representation of the network from the schedule that determines its implementation. Finding good schedules, however, remains extremely challenging. We model this scheduling problem as a sequence of optimization choices, and present a new technique to accurately predict the expected performance of a partial schedule. By leveraging these predictions we can make these optimization decisions greedily and rapidly identify an efficient schedule. This enables us to find schedules that improve the throughput of deep neural networks by 2.6x over Halide and 1.5x over TVM. Moreover, our technique is two to three orders of magnitude faster than that of these tools, and completes in seconds instead of hours.