---

Join Our Workgroup   GitHub Repo

Unlocking AI Potential: Navigating the Challenges and Opportunities of Diverse Hardware Accelerators

David Edelsohn, AI Alliance, IBM
Andrew Richards, UXL Foundation
Michael Wong, UXL Foundation, Khronos, ISOCPP, RISC-V
Thanks to AMD, Intel, Meta, and Modular

Introduction

The artificial intelligence (AI) landscape is experiencing an unprecedented surge in innovation, particularly in the realm of AI accelerator hardware. These accelerators, designed to significantly speed up AI computations, are the engines behind today’s AI breakthroughs, powering everything from natural language processing to computer vision tasks, generating songs or movies to driving cars. However, the rapid proliferation of AI hardware accelerators, each with its own unique architecture and capabilities, has given rise to a complex ecosystem. This complexity is further magnified by the diversity of AI frameworks, such as PyTorch, JAX, TensorFlow, PaddlePaddle, Pallas, and Triton, each offering distinct advantages and designed with specific use cases in mind.

The convergence of cutting-edge AI hardware and sophisticated software frameworks is where the true potential of AI is unleashed. It is at this intersection where groundbreaking advancements are made, enabling the development of more powerful, efficient, and intelligent AI systems. However, this junction also presents significant challenges for developers and researchers, who must navigate the intricate web of compatibility, performance optimization, and integration issues that arise when working with disparate hardware and software components.

This comprehensive review aims to demystify the complexities of the AI ecosystem, providing a clear and in-depth exploration of the interactions between various AI accelerators and frameworks. By examining the challenges posed by this rapidly evolving landscape and presenting strategies to effectively navigate its intricacies, this document serves as an invaluable resource for developers, researchers, and decision-makers alike. Whether you are a seasoned AI practitioner or a newcomer to the field, this guide will equip you with the knowledge and insights needed to harness the full potential of AI hardware and software, driving innovation and progress in this transformative domain.

Motivation

The increasing use of Artificial Intelligence, Deep Learning, Foundation Models, and Machine Learning throughout a wide range of tasks drives a wide range of requirements from both hardware and software. At one end of the spectrum, the training of large-scale models for natural language processing, self-driving vehicles, and generative AI demands immense computational resources. These tasks often require clusters of specialized hardware to process the vast number of calculations within a reasonable timeframe. On the other hand, the inference of these models, while less computationally intensive compared to training, can still pose challenges when deployed at scale, requiring a different mix of hardware and software optimizations to ensure efficient and responsive performance.

Moving along the spectrum, simpler AI models designed for applications such as fraud detection may have lower computational demands but often come with strict latency and throughput constraints. As we shift towards the edge of the network, AI deployments on self-driving cars, mobile devices, and smart appliances introduce additional requirements for power efficiency and real-time processing. To address these challenges, specialized AI chips with lower power consumption and dedicated AI processing capabilities have emerged as crucial components in edge computing scenarios.

At the extreme end of the spectrum, TinyML focuses on running AI models on devices with severely limited resources, such as wearables and sensor nodes. In these cases, specialized microcontrollers with built-in AI capabilities play a vital role in enabling basic AI tasks while maintaining minimal power consumption and form factor.

The challenge for vendors, software developers, and AI application developers lies in developing and utilizing common frameworks that can scale AI models across this diverse range of use cases, deployment environments, and hardware platforms. Adaptability is key for AI frameworks to effectively leverage the unique capabilities of each AI accelerator and hardware configuration, ultimately maximizing performance and efficiency for every AI application.

This guide aims to explore the strategies and solutions for navigating these complexities.

Introduction to AI Hardware Accelerators

AI accelerators have become pivotal in the modern AI landscape due to their specialized computational capabilities, designed to speed up artificial intelligence (AI) applications, including deep learning, machine learning, and data processing tasks. These accelerators perform numerically intensive computations at an unprecedented scale, allowing for rapid training and inference phases of AI models. The significance of AI accelerators lies in their ability to handle vast amounts of data and complex computations efficiently, reducing the time and energy consumption associated with traditional computing methods. This optimization unlocks new possibilities in AI development, enabling more sophisticated and accurate AI models, which are essential for advancing technologies in fields such as autonomous vehicles, healthcare, finance, entertainment, and natural language processing.

In the realm of AI accelerator hardware, key players include Graphics Processing Units (GPUs), Tensor Processing Units (TPUs), NPUs (Neural Processing Units), Field-Programmable Gate Arrays (FPGAs), Digital Signal Processors (DSPs), and Application-Specific Integrated Circuits (ASICs). GPUs, originally designed for rendering graphics, have been widely adopted for AI due to their high throughput and ability to handle parallel tasks. TPUs, developed by Google, are specifically tailored for TensorFlow operations, offering optimized performance for deep learning tasks. NPUs are specialized for neural network computations, often embedded in smartphones and IoT devices for on-device AI. FPGAs present a flexible architecture, allowing for customization to specific computational needs, making them suitable for prototyping and adaptive algorithms. ASICs are custom-designed for a particular use case, offering the highest efficiency and performance for specific AI tasks, but lack the versatility of GPUs and FPGAs.

AI accelerators play a crucial role in enhancing the performance of AI models by significantly reducing computation time and increasing efficiency. This improvement is critical for training complex models with billions of parameters, a common requirement in today’s AI challenges. By offloading heavy computational tasks to accelerators, developers can achieve higher throughput and lower latency in both training and inference phases. This enables more iterative experimentation, quicker model development, and the deployment of more advanced AI applications across various industries. Essentially, AI accelerators are the backbone that supports the rapid growth and scalability of AI technologies.

AI frameworks are the software foundations driving innovation in the AI field, providing developers with the tools and libraries needed to design, train, and deploy AI models. Leading AI frameworks include PyTorch, JAX, TensorFlow, Keras, Apache MXNet, Caffe, Chainer, Theano, Microsoft Cognitive Toolkit (CNTK), DL4J, PaddlePaddle, MindSpore, Pallas, and Triton. PyTorch, known for its dynamic computational graph and user-friendly interface, is favored for research and prototyping. JAX excels in high-performance numerical computing and machine learning research with its automatic differentiation and GPU/TPU acceleration. TensorFlow, developed by Google, offers a comprehensive ecosystem for developing and training ML models at scale, with strong support for TPUs and deployment. Keras provides a high-level API for building and training deep learning models and acts as an interface for TensorFlow. Apache MXNet offers flexibility and efficiency for deep learning. Caffe, developed by the Berkeley Vision and Learning Center, is known for its speed and modularity. Chainer is praised for its flexibility in enabling fast implementation of research ideas. Theano provides efficient definition, optimization, and evaluation of mathematical expressions. Microsoft Cognitive Toolkit (CNTK) supports commercial-grade distributed deep learning. DL4J, designed for business environments, supports distributed GPUs and CPUs. PaddlePaddle, developed by Baidu, emphasizes ease of use, scalability, and efficiency. MindSpore, by Huawei, offers comprehensive AI development and deployment capabilities. Pallas focuses on high-performance computing on GPUs, emphasizing efficiency and speed. Triton simplifies writing highly efficient GPU code, democratizing access to custom high-performance computations.

Each framework offers unique features and caters to different use cases. PyTorch is renowned for its flexibility and ease of use, making it ideal for academics and researchers focused on developing novel AI models. JAX’s ability to automatically differentiate through Python and NumPy code is particularly useful for scientists and researchers working on complex simulations and models. TensorFlow’s extensive community and robust tooling make it suitable for industrial applications requiring scalability and production readiness. Pallas, with its focus on maximizing GPU utilization, is well-suited for high-performance tasks that require extreme computational efficiency. Triton, offering an approachable way to write custom GPU kernels, appeals to developers needing to optimize specific operations beyond what is available in standard libraries.

AI frameworks are rapidly evolving to leverage the diverse capabilities of various hardware accelerators, ensuring that the computational power of GPUs, TPUs, FPGAs, and ASICs can be fully harnessed for AI development. This adaptation involves the integration of specialized libraries and APIs that facilitate direct communication between the software and the underlying hardware, optimizing for performance and efficiency. For instance, TensorFlow and PyTorch have introduced support for TPUs and CUDA-enabled GPUs, allowing developers to more easily shift workloads to these accelerators for faster processing. Additionally, frameworks are incorporating features like automatic mixed precision (AMP) and graph optimization techniques to further enhance computational efficiency on accelerators. The development of hardware-agnostic interfaces, such as ONNX (Open Neural Network Exchange), also plays a crucial role in this adaptation, enabling models trained in one framework to be executed on different types of accelerators. By embracing these advancements, AI frameworks not only unlock the potential of existing hardware but also pave the way for the next generation of AI innovations, ensuring that developers can focus on creating cutting-edge models without being constrained by hardware compatibility issues.

The Complexity of the Ecosystem

The available paths through the AI software ecosystem from model to hardware are very fragmented and parochial. Models are evolving rapidly, frameworks are evolving rapidly, hardware is evolving rapidly, which makes it difficult to ensure interoperability among the complex matrix of options. New innovations from various quarters of the AI industry are addressing pieces of the problem and exposing an opportunity to stitch together a comprehensive solution.

The primary AI frameworks for large language models are PyTorch, originally developed by Meta, and TensorFlow and JAX, developed and sponsored by Google. The frameworks have developed their own infrastructure for deploying operators on hardware as well as third-party infrastructure, such as ONNX and TVM.

PyTorch: A Detailed Exploration

PyTorch provides two primary modes of operation: Eager Mode and Graph Mode.

• Eager Mode can be considered an interpreted mode in which each operator in the model is executed as it is encountered. This provides flexibility and ease of debugging.
• Graph Mode constructs a graph from the operators, permitting various forms of optimization before executing the transformed set of operators as a whole. PyTorch Graph Mode has evolved through multiple iterations of compilers, including TorchScript and the more recent Dynamo.

PyTorch also has evolved through multiple iterations to transmit operators to hardware. PyTorch Eager mode can invoke operator kernels directly, such as those written directly in NVIDIA CUDA, AMD HIP, or Khronos SYCL, all similar C++-based kernel languages. PyTorch Eager Mode invocation of CUDA C++ can be translated to AMD HIP C++ invocation of ROCm through use of HIPify tools. PyTorch can utilize NVIDIA CUDA kernels produced by the Python-based OpenAI Triton language, which also is gaining support to directly generate AMD ROCm code. PyTorch Eager Mode operators can be captured by CUDAGraph for optimization of invocation of kernels. PyTorch also is able to leverage oneDNN to access processor-specific kernels and operators, to complement the pervasive NVIDIA cuDNN kernels and the expanding AMD ZenDNN kernels. PyTorch continues to maintain a path to the OpenBLAS library to expand its breadth of CPU targets.

PyTorch Graph mode currently utilizes the Dynamo compiler, which has enhanced the compilation process with the concept of graph breaks to include assertions in the compiled graph that allow transparent fallback to Eager Mode when characteristics of the model dynamically change. Dynamo leverages the CPython Frame API to analyze the function at runtime and identify potential optimizations, including recognizing PyTorch operations. This allows for handling more dynamic Python features like conditionals, loops, and dynamic control flow, but might have slightly higher overhead compared to JAX for simple static functions. Dynamo feeds its optimized graph to the PyTorch Inductor compiler for hardware target specific transformation. OpenAI Triton primarily is intended as a high-level human-written language for AI model kernels; Inductor targets Triton as an intermediate representation to generate CUDA and ROCm code directly, in addition to the ability to invoke pre-existing kernels (either hand-written or parametrically generated).

Microsoft, Meta and others are cooperating on Triton-Shared – an ambitious effort to utilize the Triton Language for non-GPU hardware accelerators, such as Microsoft Maia and Meta MTIA accelerator processors. Triton-Shared utilizes the LLVM MLIR project, particularly targeting the LinAlg and MemRef dialects to transform and optimize the Triton kernels, or potentially multi-kernel sub-graphs, for NPU-like hardware accelerators. Triton-Shared ingests Triton kernels, which means that the IR already has been lowered and some of the context and semantic information lost relative to PyTorch Graphs, limiting some of the optimizations available in MLIR dialects.

Beyond PyTorch: The Broader Landscape

IREE-Turbine is an effort to ingest PyTorch FX/Dynamic graphs into the IREE pipeline through the Torch-MLIR dialect, leveraging the entire IREE optimization infrastructure described in more detail in a later section. The Torch-MLIR dialect serves as a bridge, translating PyTorch models into an intermediate representation compatible with the MLIR ecosystem. Torch-MLIR continues to expand its support for a widening variety of models, including improvements for the important feature of dynamic tensor shapes. The translation facilitates the application of various optimizations and transformations inherent to MLIR, ensuring efficient execution across different hardware backends. This integrated approach streamlines the deployment of PyTorch models, enhancing their portability and efficiency across a wide range of AI accelerators, thereby accelerating the development and deployment of AI applications.

PyTorch can also utilize Google’s XLA infrastructure to leverage its optimization framework and to target XLA-based accelerators, such as Google TPU. XLA will be discussed in the TensorFlow and JAX section.

Other projects have tied PyTorch Graphs into MLIR, either directly or via ONNX. The lack of a formal specification for MLIR dialects and the instability of the MLIR and LLVM APIs complicate the ability to utilize MLIR directly. Triton and IREE offer the potential of a stable API interface for AI/ML compiler passes utilizing MLIR.

While MLIR provides a path to leverage PyTorch Graph Mode on diverse hardware, PyTorch Eager Mode remains widely used and more strongly tied to GPUs and DNN libraries.

TensorFlow, JAX and Pallas: Building for Diverse Hardware

TensorFlow and JAX have created an extensive infrastructure in support of diverse hardware. Google has developed the OpenXLA compiler ecosystem to support diverse accelerator hardware.

XLA Compiler

The XLA compiler optimizes linear algebra computations for CPUs, GPUs and ML Accelerators. JAX utilizes a tracer to convert Python functions to its own, internal representation, which is sent to XLA for compilation. This requires functions to be mostly pure (no side effects) and have static shapes for optimal results. It compiles the function once for specific input shapes and types. Subsequent calls with the same shapes reuse the compiled version, leading to speedups. However, changing input shapes triggers recompilation. Similar to PyTorch, the TensorFlow/JAX design utilizes XLA to optimize and tie together tensor operations in conjunction with calls to hand-optimized tensor kernel libraries for critical performance operations.

Kernel Optimization

TensorFlow and JAX adopted Eigen as an early design choice to optimize kernels for a wide breadth of CPU targets. They have added oneAPI oneDNN as a path to invoke optimized kernels and operators for x86_64 CPUs AVX, ARM CPUs SVE, and Intel GPUs, to complement support for NVIDIA cuDNN and AMD ZenDNN kernels.

Pallas extension and Custom Kernel Development

JAX has added the Pallas extension to write custom kernels for GPUs and TPUs. Pallas utilizes a code generation path through Mosaic for TPUs and through Triton for GPUs.This enables developers to write high-performance custom kernels without needing to delve into the intricacies of each hardware platform.

The combination of TensorFlow, JAX, and Pallas creates a powerful ecosystem for accelerating machine learning workloads across diverse hardware. By leveraging the optimization capabilities of XLA, the performance benefits of hand-optimized kernel libraries, and the flexibility of custom kernel development through Pallas, developers can harness the full potential of their hardware accelerators while maintaining a high level of productivity and portability.

PaddlePaddle: Designed for Ultra-Large-Scale AI

PaddlePaddle has developed a high performance software stack that flexibly targets a wide variety of workloads. Its design has been optimized for ultra large scale neural network training.

Compiler Infrastructure

PaddlePaddle utilizes its own compiler infrastructure, which includes importing the model either in native Paddle format or X2Paddle conversion from TensorFlow, Caffe, or ONNX. High-level optimizations are applied in the HLIR representation that is close to the original model. HLIR captures the high-level structure and semantics of the model, including the computation graph and operators. It is used to perform high-level optimizations, such as operator fusion and algebraic simplifications, before translating the model into Paddle IR (PIR). PIR abstracts the details of the model, making it easier to apply intermediate-level optimizations, such as inlining, dead code elimination, and loop transformations. PIR serves as a bridge between the high-level model definitions and the lower-level optimizations performed by Compiler Infrastructure for Neural Networks (CINN), which compiles it into executable code for specific hardware platforms. CINN applies hardware-specific optimizations, manages memory allocation, and ensures efficient data movement to maximize performance.

Compiler Backend

Paddle utilizes Eigen for optimized tensor operations. It also can utilize oneDNN as a backend, alongside GPUs (CUDA, ROCm) and XPU, and optionally OpenBLAS for BLAS operations on CPUs. For specialized hardware, PaddlePaddle offers support for XPU, Baidu’s custom-designed accelerator for deep learning. XPU provides high performance and energy efficiency for AI workloads, and PaddlePaddle’s integration with XPU allows users to seamlessly target this accelerator for their models.

The Key Advantages are:

• Scalability: Designed to handle ultra-large-scale models and distributed training across multiple nodes.
• Flexibility: Supports a wide range of hardware platforms and offers both static and dynamic graph execution modes.
• Performance: Leverages a sophisticated compiler infrastructure and optimized libraries for high performance.

XLA: Accelerated Linear Algebra for AI

Accelerated Linear Algebra (XLA) is a domain-specific compiler for linear algebra optimization originally designed for TensorFlow, but expanded to target a diverse set of frameworks (TensorFlow, JAX, PyTorch) and a diverse set of hardware. XLA utilizes MLIR for some components where appropriate, such as StableHLO, and utilizes LLVM for hardware device support, but utilizes a unique HLO IR and passes infrastructure instead of MLIR. XLA was developed and designed before MLIR was formed, and it influenced the development of MLIR. XLA is a pluggable compiler framework that can utilize IREE. It initially focused on linear algebra optimization.

XLA operates at the full-graph level and performs transformations accordingly. It also considers the topology of the target (for multi-devices environments) and performs all sharding/partitioning necessary, and optimizes the cross-device communication scheduling to overlap computations. It is parameterized and customized for a given execution environment (platform and runtime).

XLA restricts the IR to “functional” programs with no aliasing and mutations (side effects), unlike PyTorch Dynamo and Inductor. Many AI models, including ones with dynamic tensor shapes, may not run optimally in XLA. Special care needs to be taken to avoid graph breaks and graph recompilations.

IREE: MLIR-Based Framework for End-to-End AI Optimization

Intermediate Representation Execution Environment (IREE) is an MLIR-based framework to optimize machine learning models. It is an effort to develop an end-to-end AI compiler infrastructure completely in MLIR, ideally with all tensor operations generated from and by the compiler pipeline. IREE targets the entire range of ML hardware environments from data centers to mobile and edge deployments.

Model Support and Code Generation

IREE can utilize models expressed in PyTorch, TensorFlow, LiteRT (formerly TensorFlow Lite), JAX, and ONNX, and can generate code for CPUs, NVIDIA CUDA, AMD ROCm, and SPIR-V for a growing list of hardware accelerators. PyTorch connects to Torch-MLIR, and TensorFlow and JAX interface with the StableHLO IR into MHLO graph optimizer, progressively lowering to MLIR LinAlg dialect and eventually the IREE execution environment utilizing LLVM to generate CPU and GPU code.

MLIR Dialects and Optimization Frameworks

IREE utilizes the rich set of MLIR dialects, including Affine, Arith, LinAlg, MemRef, MHLO, StableHLO, SCF, Tensor, and TOSA to construct optimization passes. The dialects can be utilized in optimization passes to apply semantic-specific optimizations that permit complex and aggressive transformations relevant to dialect-specific representations of the objects and model.

IREE leverages these dialects to perform a wide range of optimizations, including:

• Operator Fusion: Combining multiple operations into a single kernel to reduce overhead and improve performance.
• Tiling: Decomposing computations into smaller tiles to improve data locality and cache utilization.
• Loop Unrolling: Duplicating loop bodies to reduce loop overhead and expose parallelism.
• Hardware-Specific Optimizations: Leveraging specific instructions and capabilities of the target hardware for maximum performance.

IREE’s focus on a unified MLIR-based infrastructure and its comprehensive approach to optimization make it a promising tool for accelerating AI workloads across diverse hardware platforms.

Edge AI Ecosystem

ExecuTorch: Streamlining AI Deployment on Mobile and Edge Devices

Mobile and edge devices have special requirements for AI deployment, including diverse hardware, critical power requirements, low or no internet connectivity, and realtime processing constraints. ExecuTorch is an emerging framework designed to facilitate the deployment of AI models on mobile and edge devices, complementing the capabilities of existing platforms like PyTorch. As part of the broader push towards enabling efficient on-device AI, ExecuTorch focuses on providing a streamlined runtime environment that optimizes PyTorch models for execution on constrained hardware. This includes support for various hardware accelerators and integration with platform-specific APIs to ensure optimal performance. ExecuTorch aims to leverage the flexibility and ease of use of PyTorch while introducing optimizations that reduce the computational footprint and power consumption of AI models. This makes it an attractive choice for developers looking to deploy sophisticated AI applications in scenarios where computational resources are limited, and real-time processing is crucial.

Model Transformation

ExecuTorch exports a PyTorch model graph in ATen dialect in which it can be transformed and optimized in architecture-agnostic ways. The ATen dialect is lowered to the Edge dialect that is aware of the parameterized constraints of the target device for additional specialization. The Edge dialect then is transformed to the Backend dialect that leverages delegates for specific hardware, allowing Core ML on iOS, QNN on Qualcomm, or TOSA on Arm to rewrite the graph. The resulting graph can be further prepared for the runtime environment through memory layout and usage planning, selectively linking actively-used kernels, and optimally serializing and packing the program for efficient loading and execution by the runtime. For optimal memory planning, tensor mutations need to be expunged, as required by XLA.

Target Backend and Dialects

ExecuTorch is designed to target a number of mobile hardware accelerators, including CPUs via XNNPACK, GPUs via Vulkan, Apple Neural Engine via Apple CoreML, and DSPs, with flexibility for additional targets.

Model Optimization

ExecuTorch applies a variety of optimizations to improve model performance on mobile and edge devices:

• Quantization: Reduces the precision of model weights and activations to lower memory usage and computational requirements.
• Pruning: Removes redundant connections or neurons in the model to reduce its size and complexity.
• Fusion: Combines multiple operations into a single kernel to reduce overhead and improve performance.
• Memory Planning: Optimizes memory layout and usage to reduce memory footprint and improve cache utilization. To ensure optimal memory planning, tensor mutations are eliminated as required by XLA (Accelerated Linear Algebra). This approach guarantees predictable and efficient memory usage, which is critical for mobile and edge deployments where resources are limited.
• Kernel Selection: Selectively links actively used kernels to minimize the runtime’s size.

By addressing the specific needs of mobile and edge AI deployment, ExecuTorch provides a robust framework that combines the strengths of PyTorch with targeted optimizations for constrained environments. This ensures that AI models can be efficiently deployed and executed on a wide range of devices, enabling real-time, on-device AI applications.

LiteRT: Enabling On-Device Machine Learning

Lite Runtime (LiteRT), formerly TensorFlow Lite (TFLite), is a lightweight, open-source deep learning framework designed specifically for mobile and edge devices. It enables developers to deploy machine learning models on resource-constrained environments such as smartphones, embedded systems, and IoT devices. LiteRT provides a suite of tools and APIs for optimizing and converting TensorFlow models into a format that is efficient for mobile and edge hardware. Its interpreter is optimized for low latency and high performance, supporting hardware acceleration on various platforms, including Android Neural Networks API (NNAPI), GPU, and Hexagon DSP. LiteRT’s ability to perform on-device inference ensures real-time data processing and reduced dependency on cloud services, making it ideal for applications requiring low latency, privacy, and offline capabilities.

Use cases

LiteRT has been widely adopted across various industries and applications:

• Mobile Apps: LiteRT powers on-device AI features in mobile apps, such as real-time image recognition, language translation, and voice assistants.
• Embedded Systems: LiteRT enables AI on devices like microcontrollers and IoT devices, allowing for intelligent edge computing in areas like smart homes, wearables, and industrial automation.
• Healthcare: LiteRT has been used to develop medical image analysis tools, disease prediction models, and personalized health monitoring applications.
• Finance: LiteRT can be leveraged for fraud detection, risk assessment, and personalized financial recommendations.
• Autonomous Systems: LiteRT is used in applications like object detection and classification for autonomous vehicles and drones.

Paddle Lite: Lightweight Inference for Edge Devices