Tadashi: Enabling AI-Based Automated Code Generation With Guaranteed Correctness

2024-10-04

arXiv.org

AI Summary

Existing black-box neural-network-based AI code generation methods lack formal guarantees for the correctness and legality of loop scheduling transformations. Method: We propose the first end-to-end, polyhedral-driven framework supporting formal correctness guarantees, tightly integrating machine learning optimization with compilation theory. Our approach constructs a verifiable program transformation space, introduces differentiable modeling of legality constraints, and jointly generates training data annotated with machine-checkable proofs. Contribution/Results: All generated scheduling transformations are formally proven to preserve semantic equivalence under the polyhedral model. The framework incurs negligible runtime overhead and demonstrates strong generalization across diverse hardware platforms—including CPUs, GPUs, and FPGAs—as well as canonical loop optimization scenarios. An open-source implementation is publicly available.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Hardware-aware ML
• Computer Vision: Learning & Optimization for CV

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Frameworks and domain-specific languages for auto-generating code have traditionally depended on human experts to implement rigorous methods ensuring the legality of code transformations. Recently, machine learning (ML) has gained traction for generating code optimized for specific hardware targets. However, ML approaches-particularly black-box neural networks-offer no guarantees on the correctness or legality of the transformations they produce. To address this gap, we introduce Tadashi, an end-to-end system that leverages the polyhedral model to support researchers in curating datasets critical for ML-based code generation. Tadashi provides an end-to-end system capable of applying, verifying, and evaluating candidate transformations on polyhedral schedules with both reliability and practicality. We formally prove that Tadashi guarantees the legality of generated transformations, demonstrate its low runtime overhead, and showcase its broad applicability. Tadashi available at https://github.com/vatai/tadashi/.

Problem

• Ensuring correctness in AI-based automated code generation
• Guaranteeing legality of transformations in ML-generated code
• Providing reliable polyhedral model for ML dataset curation

Innovation

• Leverages polyhedral model for ML-based code generation
• Guarantees legality of generated code transformations
• Provides end-to-end system with low runtime overhead