Code Migration with Formal Verification for Performance Improvement of Legacy Code

TensIR - Tenspiler's IR for Tensor Operations

While we can conceptually translate an input program in a source language directly into an output program in a target language via program synthesis, doing so would have a few drawbacks: 1) such a translator would be specific to the source and target language, and 2) the grammar of the target language (that determines the set of possible programs in the target language - also called as the search space) may allow programs of infinite depth (e.g., via loops or recursion), leading to unbounded search space. The solution adopted in the program synthesis community is to develop an intermediate language (IR) that captures the semantics of the target language while also restricting the size of the search space.

In Tenspiler, we have developed an IR, called TensIR, that captures commonly-used tensor operations (such as matrix multiplication, tensor-scalar operations, etc.) and types (e.g., vectors, scalars, tensors, etc.) from popular tensor DSLs. More importantly, the grammar of the TensIR is such that it restricts the search space of the possible TensIR programs while translating input programs into tensor DSLs. We show TensIR grammar below.

Design of Tenspiler

Figure below shows the overall design of Tenspiler. In essence, it is divided into three main components: synthesis, verification, and code generation. We cover all the three steps below.

Program Synthesis - Synthesizing TensIR Program

Now that we have discussed TensIR and its purpose, we now discuss how our program synthesis step leverages it. The goal of Tenspiler's program synthesis step is to search for a TensIR program that is semantically equivalent to the input program in a source language. Technically, the semantically-equivalent program in the target language is called program summary (PS) in the field of program synthesis. In addition to PS, the output of our synthesis step is also loop invariants (INV) for the input program. Conceptually, loop invariants are boolean expressions that capture relationships between input and output variables of a loop that hold true for any iteration of the loop. INV are required for our synthesis problem because the input programs had loops and loop invariants are required to prove the equivalence of such programs.

In Tenspiler, we formulate the search problem as a Syntax-guided Synthesis problem [5] (called SyGuS in short). The key idea in syntax-guided synthesis is that the syntax of the target language (TensIR in our case) guides the synthesis (i.e., search) of the target program. Our SyGuS problem is characterized by three parameters:

1. The specification: it describes the program property that the synthesized TensIR program should satisfy – intuitively, it is the condition that checks if the synthesized program is equivalent to the input program. In Tenspiler, we use verification conditions (VCs) as the specification for the synthesis phase as VCs provide full guarantees (up to a bound) on the equivalence of the input program and the synthesized program. Intuitively, VCs are logical expressions that encode the semantics of the input program.
2. The search space that describes the set of possible TensIR programs. We use TensIR grammar to define the search space. Conceptually, the search space imposes syntactic constraints on the structure of the output TensIR program.
3. The search algorithm that searches for the TensIR program. In Tenspiler, we use an enumerative algorithm, meaning it enumerates all possible programs up to a certain depth.

Verification

Given the limitations of the current SyGuS solvers, our synthesis step finds PS and INV by validating them against VC for a bounded set of program states only. In other words, our synthesis step does not guarantee full semantic equivalence between the input program and the synthesized TensIR program. To ensure complete semantic equivalence, we need a verification step. In Tenspiler, we leverage an existing SMT solver (CVC5) and ask it to find a PS and INV that violate VC (i.e., a counter example). Inability of the solver to find such an example indicates full semantic equivalence, otherwise we add the found counter example as a clause in VC and repeat the step.

Code generation

TensIR, as an IR, allows us to keep the program synthesis step tractable and also to verify the semantic equivalence between the input and the output programs. Given our goal of generating an output program in the target language though, we implement a rule-based code generator that accepts the TensIR program as an input and outputs the target program in the target languages. Tenspiler accepts input programs written in C/C++/Python and is able to generate output programs in six different AI-specific DSLs: Numpy, Tensorflow, PyTorch, Apple MLX, TPC-C, and Gemmini.

Tenspiler accepts input programs written in C/C++/Python and is able to generate output programs in six different AI-specific DSLs: Numpy, TensorFlow, PyTorch, Apple MLX, TPC-C, and Gemmini.

Verification

Given the limitations of the current SyGuS solvers, our synthesis step finds PS and INV by validating them against VC for a bounded set of program states only. In other words, our synthesis step does not guarantee full semantic equivalence between the input program and the synthesized TensIR program. To ensure complete semantic equivalence, we need a verification step. In Tenspiler, we leverage an existing SMT solver (CVC5) and ask it to find a PS and INV that violate VC (i.e., a counter example). Inability of the solver to find such an example indicates full semantic equivalence, otherwise we add the found counter example as a clause in VC and repeat the step.

Code Generation

TensIR, as an IR, allows us to keep the program synthesis step tractable and also to verify the semantic equivalence between the input and the output programs. Given our goal of generating an output program in the target language though, we implement a rule-based code generator that accepts the TensIR program as an input and outputs the target program in the target languages. Tenspiler accepts input programs written in C/C++/Python and is able to generate output programs in six different AI-specific DSLs: Numpy, Tensorflow, PyTorch, Apple MLX, TPC-C, and Gemmini.

Tenspiler accepts input programs written in C/C++/Python and is able to generate output programs in six different AI-specific DSLs: Numpy, TensorFlow, PyTorch, Apple MLX, TPC-C, and Gemmini.