Wave: A Symbolic Python DSL And Compiler for High-Performance Machine Learning
Harsh Menon ⋅ Oleksandr Zinenko ⋅ Gaurav Verma ⋅ Stanley Winata ⋅ Ivan Butygin ⋅ Nithin Meganathan ⋅ Sanket Pandit ⋅ William Gallard Hatch ⋅ Surya Jasper ⋅ Megan Kuo ⋅ Sahil FAIZAL ⋅ Ashay Rane ⋅ Aurore De Spirlet ⋅ Martin P. Lücke
Abstract
Modern ML models demand ever-greater compute, prompting hardware vendors to add specialized matrix cores to their GPUs. While these units unlock high throughput, they impose intricate programming models and addressing schemes that are difficult to manage by hand. This paper introduces Wave, a Python-embedded DSL for kernel authoring that automates these complex address computations and lets authors focus on core computation. In experiments, it matches or surpasses the performance of state-of-the-art kernel DSLs and libraries.
Successful Page Load