IBM's New AI Chip Tool Makes Software Faster and Easier

This post is the middle chapter of three and explains how decisions made by torch-spyre’s Inductor become a program the Spyre runtime can launch. The contract between Inductor and the backend compiler is KTIR, an MLIR representation built on a data-parallel view of the accelerator, in which every compute tile can see every memory and a single logical tensor can live spread across several cores’ scratchpads.

A new open-source Dataflow Scheduler consumes KTIR and turns each tile’s work into a dataflow pipeline: it splits the kernel into programs, starts from a load to compute to store pipeline, routes every transfer through the memory hierarchy given in an architecture description, and expands that pipeline into a nested, double-buffered, multi-instance schedule. The scheduler’s output is represented using DataflowIR, which Spyre’s device-specific stages use to produce SpyreCode, a symbolic SpyreBinary plus a job plan that the Spyre runtime prepares once and launches many times, filling in tensor addresses and shapes at launch.