I’m working on a tensor computing language/compiler called i with a simple explicit scheduling model (loop splitting, loop ordering, input “staging”). These mechanisms alone are enough to express complex algorithms like FlashAttention, generating target code with techniques like loop fusion, minimized intermediate allocations, and “online” reductions.
Right now there is a runtime and compiler targeting C, written in dependency-free Rust, and a minimal Python frontend. The project is very much proof-of-concept stage so not yet fast. Working on a CUDA backend now.
The goal is to enable automatic discovery of FlashAttention-style optimizations which is not feasible with current compilers.
Very open to feedback/discussion from anybody interested in or knowledgeable about tensor compilers!
I made i, a tensor computation language with a simple explicit scheduling model. The only scheduling concepts in i are loop splits, loop ordering, and input producer staging (where one intermediate value is computed inside the loop nest of another).
Those three concepts are enough to write numerically stable online blockwise FlashAttention. Loop tiling, loop fusion, storage folding, and (critically) online reduction rewriting fall out as _predictable_ consequences of the lowering.
The hope is that this scheduling model makes it easier for people and search algorithms to find performant schedules.
Current status: working proof-of-concept. Compiler/runtime in Rust, C backend, small Python front-end, zero dependencies. Not fast (yet).
George just said today during a Q+A session that they do not intend to do learned control because traditional control is already greatly exceeds human ability.
Right now there is a runtime and compiler targeting C, written in dependency-free Rust, and a minimal Python frontend. The project is very much proof-of-concept stage so not yet fast. Working on a CUDA backend now.
The goal is to enable automatic discovery of FlashAttention-style optimizations which is not feasible with current compilers.
Very open to feedback/discussion from anybody interested in or knowledgeable about tensor compilers!
repo: https://github.com/ilang-dev/i