Tensor compilers, essential for generating efficient code for deep learning models across various applications, employ tensor graph rewrites as one of the key optimizations. These rewrites optimize tensor computational graphs with the expectation of preserving semantics for tensors of arbitrary rank and size. Despite this expectation, to the best of our knowledge, there does not exist a fully automated verification system to prove the soundness of these rewrites for tensors of arbitrary rank and size. Previous works, while successful in verifying rewrites with tensors of concrete rank, do not provide guarantees in the unbounded setting. To fill this gap, we introduce TensorRight, the first automatic verification system that can verify tensor graph rewrites for input tensors of arbitrary rank and size. We introduce a core language, TensorRight DSL, to represent rewrite rules using a novel axis definition, called aggregated-axis, which allows us to reason about an unbounded number of axes. We achieve unbounded verification by proving that there exists a bound on tensor ranks, under which bounded verification of all instances implies the correctness of the rewrite rule in the unbounded setting. We derive an algorithm to compute this rank using the denotational semantics of TensorRight DSL. TensorRight employs this algorithm to generate a finite number of bounded-verification proof obligations, which are then dispatched to an SMT solver using symbolic execution to automatically verify the correctness of the rewrite rules. We evaluate TensorRight's verification capabilities by implementing rewrite rules present in XLA's algebraic simplifier. The results demonstrate that TensorRight can prove the correctness of 115 out of 175 rules in their full generality, while the closest automatic, bounded-verification system can express only 18 of these rules.

翻译：张量编译器对于为深度学习模型在各种应用中生成高效代码至关重要，其采用张量图重写作为关键优化手段之一。这些重写旨在优化张量计算图，并期望在任意秩和尺寸的张量上保持语义不变。尽管存在这一期望，据我们所知，目前尚不存在完全自动化的验证系统来证明这些重写在任意秩和尺寸张量上的正确性。先前的研究虽成功验证了具体秩张量的重写，但无法在无界设定下提供保证。为填补这一空白，我们提出了TensorRight，首个能够验证任意秩和尺寸输入张量图重写的自动化验证系统。我们引入了一种核心语言——TensorRight DSL，通过一种称为聚合轴的新型轴定义来表示重写规则，这使我们能够推理无界数量的轴。我们通过证明存在一个张量秩的边界来实现无界验证，在该边界下对所有实例的有界验证意味着重写规则在无界设定下的正确性。我们利用TensorRight DSL的指称语义推导出一种计算该秩的算法。TensorRight运用该算法生成有限数量的有界验证证明义务，随后通过符号执行将这些义务分派给SMT求解器，以自动验证重写规则的正确性。我们通过实现XLA代数简化器中存在的重写规则来评估TensorRight的验证能力。结果表明，TensorRight能够证明175条规则中115条规则的完全通用正确性，而最接近的自动化有界验证系统仅能表达其中18条规则。