Tensorize: Fast Synthesis of Tensor Programs from Legacy Code using Symbolic Tracing, Sketching and Solving

Tensor domain specific languages (DSLs) achieve substantial performance due to high-level compiler optimization and hardware acceleration. However, to achieve such performance for existing applications requires the programmer to manual rewrite their legacy code in evolving Tensor DSLs. Prior efforts to automate this translation face significant scalability issues which greatly reduces their applicability to real-world code. This paper presents Tensorize, a novel MLIR-based compiler approach to automatically lift legacy code to high level Tensor DSLs using program synthesis. Tensorizeuses a symbolic trace of the legacy program as a specification and automatically selects sketches from the target Tensor DSLs to drive the program synthesis. It uses an algebraic solver to rapidly simplify the specification, resulting in a fast, automatic approach that is correct by design. We evaluate Tensorizeon several legacy code benchmarks and compare against state-of-the-art techniques. Tensorizeis able to lift more code than prior schemes, is an order of magnitude faster in synthesis time, and guarantees correctness by construction.