TensorRight: Automated Verification of Tensor Graph Rewrites

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Arora, Jai, Lu, Sirui, Jain, Devansh, Xu, Tianfan, Houshmand, Farzin, Phothilimthana, Phitchaya Mangpo, Lesani, Mohsen, Narayanan, Praveen, Murthy, Karthik Srinivasa, Bodik, Rastislav, Sabne, Amit, Mendis, Charith
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866917098306928640
author Arora, Jai
Lu, Sirui
Jain, Devansh
Xu, Tianfan
Houshmand, Farzin
Phothilimthana, Phitchaya Mangpo
Lesani, Mohsen
Narayanan, Praveen
Murthy, Karthik Srinivasa
Bodik, Rastislav
Sabne, Amit
Mendis, Charith
author_facet Arora, Jai
Lu, Sirui
Jain, Devansh
Xu, Tianfan
Houshmand, Farzin
Phothilimthana, Phitchaya Mangpo
Lesani, Mohsen
Narayanan, Praveen
Murthy, Karthik Srinivasa
Bodik, Rastislav
Sabne, Amit
Mendis, Charith
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2511_17838
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle TensorRight: Automated Verification of Tensor Graph Rewrites
Arora, Jai
Lu, Sirui
Jain, Devansh
Xu, Tianfan
Houshmand, Farzin
Phothilimthana, Phitchaya Mangpo
Lesani, Mohsen
Narayanan, Praveen
Murthy, Karthik Srinivasa
Bodik, Rastislav
Sabne, Amit
Mendis, Charith
Programming Languages
F.3.1; F.3.2; D.3.4
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.
title TensorRight: Automated Verification of Tensor Graph Rewrites
topic Programming Languages
F.3.1; F.3.2; D.3.4
url https://arxiv.org/abs/2511.17838