NineToothed: A Triton-Based High-Level Domain-Specific Language for Machine Learning

TL;DR

NineToothed introduces a high-level, serial programming language for machine learning that simplifies kernel development by automatically transforming serial code into efficient parallel code, maintaining performance.

Contribution

It presents NineToothed, a DSL with tensor-oriented metaprogramming and an automatic code transformation, reducing complexity in deep learning kernel development.

Findings

Simplifies compute kernel development process

Achieves performance comparable to Triton

Enables serial programming for deep learning workloads

Abstract

The emergence of deep learning domain-specific languages (DSLs) has substantially reduced the obstacles in developing high-performance, cross-platform compute kernels. However, current DSLs, such as Triton, still demand that developers possess expertise in parallel programming and expose them to many low-level details. This requirement complicates the development process and adds to the difficulty of maintaining compute kernels. Consequently, developing a new programming model that supports serial programming for deep learning workloads is crucial. This paper introduces NineToothed, a domain-specific language that offers serial semantics for machine learning programming. Through the automatic transformation of serial code into parallel code, NineToothed significantly streamlines the development process while causing minimal performance degradation. NineToothed encompasses (1) a language with tensor-oriented metaprogramming (TOM) that adopts the arrange-and-apply paradigm, enabling the expression of tiled computations without the need to manage low-level details and (2) a code generator for generating high-performance parallel code. Our evaluation results indicate that NineToothed can greatly simplify compute kernel development while maintaining performance comparable to that of Triton.