TPU-MLIR: A Compiler For TPU Using MLIR

• URL: http://arxiv.org/abs/2210.15016v1
• Date: Sun, 23 Oct 2022 10:45:54 GMT
• Title: TPU-MLIR: A Compiler For TPU Using MLIR
• Abstract summary: TPU-MLIR deploys pre-trained neural network (NN) models to a custom ASIC called a Processing Unit (TPU) A NN model is translated to the TOP dialect and then lowered to the TPU dialect for different TPUs according to the chip's configuration. We demonstrate how to use the MLIR pass pipeline to organize and perform optimization on TPU to generate machine code.
• Score: 2.6519283973116963
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Multi-level intermediate representations (MLIR) show great promise for reducing the cost of building domain-specific compilers by providing a reusable and extensible compiler infrastructure. This work presents TPU-MLIR, an end-to-end compiler based on MLIR that deploys pre-trained neural network (NN) models to a custom ASIC called a Tensor Processing Unit (TPU). TPU-MLIR defines two new dialects to implement its functionality: 1. a Tensor operation (TOP) dialect that encodes the deep learning graph semantics and independent of the deep learning framework and 2. a TPU kernel dialect to provide a standard kernel computation on TPU. A NN model is translated to the TOP dialect and then lowered to the TPU dialect for different TPUs according to the chip's configuration. We demonstrate how to use the MLIR pass pipeline to organize and perform optimization on TPU to generate machine code. The paper also presents a verification procedure to ensure the correctness of each transform stage.

Related papers

• An MLIR-Based Compilation Method for Large Language Models [9.14762452723213]
  This paper presents an MLIR based compilation method for large language models.<n>A model is first represented as TopOp, then lowered layer by layer to TpuOp, and finally a deployable binary is generated.<n>The method has been implemented in the TPU-MLIR compiler and the LLM-TPU deployment project.
  arXiv  Detail & Related papers  (2026-07-17T11:24:45Z)
• $\mathcal{S}^2$IT: Stepwise Syntax Integration Tuning for Large Language Models in Aspect Sentiment Quad Prediction [51.165432266846096]
  syntactic structure information is underutilized in the generative paradigm of large language models (LLMs)<n>S2IT progressively integrates syntactic structure knowledge into LLMs through a multi-step tuning process.<n> Experiments demonstrate that S2IT significantly improves state-of-the-art performance across multiple datasets.
  arXiv  Detail & Related papers  (2026-04-25T13:24:29Z)
• Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU [3.9009842917937534]
  We present Ragged Paged Attention (RPA), a high-performance and flexible attention kernel for Google's TPUs.<n>RPA achieves up to 86% memory bandwidth utilization (MBU) in decode and 73% model FLOPs utilization (MFU) in prefill.
  arXiv  Detail & Related papers  (2026-04-16T18:30:13Z)
• TorchLean: Formalizing Neural Networks in Lean [71.68907600404513]
  We introduce TorchLean, a framework that treats learned models as first-class mathematical objects with a single, precise semantics shared by execution and verification.<n>We validate TorchLean end-to-end on certified robustness, physics-informed residual bounds for PINNs, and Lyapunov-style neural controller verification.
  arXiv  Detail & Related papers  (2026-02-26T05:11:44Z)
• Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models [96.0074341403456]
  Inference-time compute has re-emerged as a practical way to improve LLM reasoning.<n>Most test-time scaling (TTS) algorithms rely on autoregressive decoding.<n>We propose Prism, an efficient TTS framework for dLLMs.
  arXiv  Detail & Related papers  (2026-02-02T09:14:51Z)
• NNTile: a machine learning framework capable of training extremely large GPT language models on a single node [83.9328245724548]
  NNTile is based on a StarPU library, which implements task-based parallelism and schedules all provided tasks onto all available processing units. It means that a particular operation, necessary to train a large neural network, can be performed on any of the CPU cores or GPU devices.
  arXiv  Detail & Related papers  (2025-04-17T16:22:32Z)
• Disambiguation of Chinese Polyphones in an End-to-End Framework with Semantic Features Extracted by Pre-trained BERT [81.99600765234285]
  We propose an end-to-end framework to predict the pronunciation of a polyphonic character. The proposed method consists of a pre-trained bidirectional encoder representations from Transformers (BERT) model and a neural network (NN) based classifier.
  arXiv  Detail & Related papers  (2025-01-02T06:51:52Z)
• Training Neural Networks as Recognizers of Formal Languages [87.06906286950438]
  Formal language theory pertains specifically to recognizers. It is common to instead use proxy tasks that are similar in only an informal sense. We correct this mismatch by training and evaluating neural networks directly as binary classifiers of strings.
  arXiv  Detail & Related papers  (2024-11-11T16:33:25Z)
• Hardware-Assisted Virtualization of Neural Processing Units for Cloud Platforms [3.93151902626205]
  We present Neu10, a holistic NPU virtualization framework. Neu10 consists of (1) a flexible NPU abstraction called vNPU, which enables fine-grained virtualization of the heterogeneous compute units in a physical NPU (pNPU); (2) a vNPU resource allocator that enables pay-as-you-go computing model and flexible vNPU-to-pNPU mappings for improved resource utilization and cost-effectiveness; and (3) an ISA extension of modern NPU architecture for facilitating fine-grained tensor operator scheduling for multiple vNPUs.
  arXiv  Detail & Related papers  (2024-08-07T21:45:01Z)
• Flex-TPU: A Flexible TPU with Runtime Reconfigurable Dataflow Architecture [0.0]
  The work herein consists of developing a reconfigurable dataflow TPU, called the Flex-TPU, which can dynamically change the dataflow per layer during run-time. The results show that our Flex-TPU design achieves a significant performance increase of up to 2.75x compared to conventional TPU, with only minor area and power overheads.
  arXiv  Detail & Related papers  (2024-07-11T17:33:38Z)
• On the Representational Capacity of Neural Language Models with Chain-of-Thought Reasoning [87.73401758641089]
  Chain-of-thought (CoT) reasoning has improved the performance of modern language models (LMs) We show that LMs can represent the same family of distributions over strings as probabilistic Turing machines.
  arXiv  Detail & Related papers  (2024-06-20T10:59:02Z)
• Exploration of TPUs for AI Applications [0.0]
  Processing Units (TPUs) are specialized hardware accelerators for deep learning developed by Google. This paper aims to explore TPUs in cloud and edge computing focusing on its applications in AI.
  arXiv  Detail & Related papers  (2023-09-16T07:58:05Z)
• Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
  This work introduces a framework to develop efficient, portable Deep Learning and High Performance Computing kernels. We decompose the kernel development in two steps: 1) Expressing the computational core using Processing Primitives (TPPs) and 2) Expressing the logical loops around TPPs in a high-level, declarative fashion. We demonstrate the efficacy of our approach using standalone kernels and end-to-end workloads that outperform state-of-the-art implementations on diverse CPU platforms.
  arXiv  Detail & Related papers  (2023-04-25T05:04:44Z)
• Heterogeneous Integration of In-Memory Analog Computing Architectures with Tensor Processing Units [0.0]
  This paper introduces a novel, heterogeneous, mixed-signal, and mixed-precision architecture that integrates an IMAC unit with an edge TPU to enhance mobile CNN performance. We propose a unified learning algorithm that incorporates mixed-precision training techniques to mitigate potential accuracy drops when deploying models on the TPU-IMAC architecture.
  arXiv  Detail & Related papers  (2023-04-18T19:44:56Z)
• Tensor Slicing and Optimization for Multicore NPUs [2.670309629218727]
  This paper proposes a compiler optimization pass for Multicore NPUs, called Slicing Optimization (TSO) TSO identifies the best tensor slicing that minimizes execution time for a set of CNN models. Results show that TSO is capable of identifying the best tensor slicing that minimizes execution time for a set of CNN models.
  arXiv  Detail & Related papers  (2023-04-06T12:03:03Z)
• An Adaptive Device-Edge Co-Inference Framework Based on Soft Actor-Critic [72.35307086274912]
  High-dimension parameter model and large-scale mathematical calculation restrict execution efficiency, especially for Internet of Things (IoT) devices. We propose a new Deep Reinforcement Learning (DRL)-Soft Actor Critic for discrete (SAC-d), which generates the emphexit point, emphexit point, and emphcompressing bits by soft policy iterations. Based on the latency and accuracy aware reward design, such an computation can well adapt to the complex environment like dynamic wireless channel and arbitrary processing, and is capable of supporting the 5G URL
  arXiv  Detail & Related papers  (2022-01-09T09:31:50Z)
• Compiling ONNX Neural Network Models Using MLIR [51.903932262028235]
  We present a preliminary report on our onnx-mlir compiler, which generates code for the inference of deep neural network models. Onnx-mlir relies on the Multi-Level Intermediate Representation (MLIR) infrastructure recently integrated in the LLVM project.
  arXiv  Detail & Related papers  (2020-08-19T05:28:08Z)
• Synthetic Datasets for Neural Program Synthesis [66.20924952964117]
  We propose a new methodology for controlling and evaluating the bias of synthetic data distributions over both programs and specifications. We demonstrate, using the Karel DSL and a small Calculator DSL, that training deep networks on these distributions leads to improved cross-distribution generalization performance.
  arXiv  Detail & Related papers  (2019-12-27T21:28:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.