TauLang: A Language and Compiler for Space-Time Programming on Statically Scheduled Accelerators
Xuxiao Yang, Junchi Wu, Feiyang Wu, Zhuoran Li, Shuo Tao, Dongyi Fan, Yiming Zhang, Keqin Luo, Xingrui Gao, Xuehai Qian, Youwei Zhuo
Abstract
Statically scheduled accelerators such as Groq's Tensor Streaming Processor (TSP) achieve low-latency large language model inference by stripping dynamic synchronization hardware from the chip. With no caches, arbiters, or hardware schedulers, the compiler must fix the exact slice and issue cycle of every operation, so when an operation runs is a condition for correctness. Existing kernel languages cannot target such machines because they rely on hardware to resolve cycle timing and instead encode schedules implicitly in code structure. We present TauLang, a kernel language and compiler that makes time a first-class citizen of the program. The programmer schedules operations, loops, and composed kernels with parameterized explicit times, so exploring schedules means changing parameter values rather than rewriting code. TauLang validates explicit times algebraically in TauIR, its timed intermediate representation. It then solves tensor placement and cycle scheduling jointly in an integer linear program, using explicit times to group connected operations into timing clusters that prune the search space. On a cycleaccurate model of a TSP-like machine, joint search achieves a 1.40× latency speedup across ten kernels over a space-first flow. A few explicit time pins keep latency within one cycle of the untimed schedules and scale joint search to 1,636 GEMM blocks, whereas untimed search finds no schedule at 32.
Create a lesson
Related papers
Verified Implementation of the Next-gen Borrow Checking Algorithm for Rust via Abstract Interpretation
Jinhua Wu, Yuting Wang, Zhemin Chen
Demystifying Information Flow Downgrading using the Real--Ideal Paradigm
Suraaj Kanniwadi, Andrew C. Myers
Linking as Evaluation: Correct Separate Compilation via First-Class Environments
Jam Kabeer Ali Khan, Jinhao Tan, Bruno C. d. S. Oliveira
Properly defining Erlang lists
Albert Schimpf, Stefan Wehr, Annette Bieniusa
Proving the Wrong Program: Miscompilations as a Threat to zkVM Guarantees
Hui Xuan Tan, Stefanos Chaliasos, Alastair F. Donaldson
Handling JAX
Oghenevwogaga Ebresafe, Daniel Zheng, Gordon Plotkin et al.