Mirage is the first multi-level superoptimizer for tensor programs. A key idea in Mirage is muGraphs, a hierarchical graph representation that specifies tensor programs across multiple levels of the GPU compute hierarchy (kernel, thread block, and thread levels). muGraphs enable Mirage to discover novel optimizations that combine algebraic transformations, schedule transformations, and generation of new custom kernels.
To navigate the large search space, Mirage introduces a pruning technique based on abstract expressions that significantly reduces the search space and provides a certain optimality guarantee. To ensure that the optimized muGraph is equivalent to the input program, Mirage introduces a probabilistic equivalence verification procedure with strong theoretical guarantees.
Our evaluation shows that Mirage significantly outperforms existing approaches even for DNNs that are widely used and heavily optimized.