Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Hardware-Aware Reformulation of Convolutions for Efficient Execution on Specialized AI Hardware: A Case Study on NVIDIA Tensor Cores

Ganesh Bikshandi

TL;DR

This work introduces a hardware-aware, post-training reformulation for CNNs that targets NVIDIA Tensor Core alignment by trading width for channels through width folding. The core idea reshapes inputs and expands filters so that convolution computations become block-diagonal and inherently compatible with hardware tiling, all while preserving exact input–output semantics without retraining or zero-padding. The authors formalize the transformation, prove correctness, and show how to realize it as an MLIR IR pass with extensions to GEMM via 1×1 convolutions. Empirical results on NVIDIA A100 indicate substantial speedups over baseline cuDNN/TensorRT implementations, validating the potential of semantic tuning to improve hardware utilization. This semantic-tuning framework lays groundwork for broader, compiler-driven, hardware-aware CNN optimizations across accelerators.

Abstract

Convolutional Neural Networks (CNNs) are central to modern AI, but their performance is often limited by hardware constraints. NVIDIA Tensor Cores, for instance, require input channels to be multiples of 8 and sometimes 512 for efficient execution. {\em oneDNN} framework for CPU imposes such a requirement for the blocked format. Traditional approaches address such alignment issue using zero-padding, which can be inefficient. In this work, we present a first-step, hardware-aware reformulation of CNN computations using rewrite rules, restructuring the underlying math to satisfy hardware alignment entirely {\bf post-training} without modifying network weights. While our current implementation focuses on a single transformation for Tensor Cores, this approach is generalizable, laying the foundation to explore additional transformations for CPU and accelerators. This study represents an initial step toward {\em semantic tuning}, a systematic, hardware-aware optimization strategy for efficient deployment of CNN models on specialized AI hardware.

Hardware-Aware Reformulation of Convolutions for Efficient Execution on Specialized AI Hardware: A Case Study on NVIDIA Tensor Cores

TL;DR

This work introduces a hardware-aware, post-training reformulation for CNNs that targets NVIDIA Tensor Core alignment by trading width for channels through width folding. The core idea reshapes inputs and expands filters so that convolution computations become block-diagonal and inherently compatible with hardware tiling, all while preserving exact input–output semantics without retraining or zero-padding. The authors formalize the transformation, prove correctness, and show how to realize it as an MLIR IR pass with extensions to GEMM via 1×1 convolutions. Empirical results on NVIDIA A100 indicate substantial speedups over baseline cuDNN/TensorRT implementations, validating the potential of semantic tuning to improve hardware utilization. This semantic-tuning framework lays groundwork for broader, compiler-driven, hardware-aware CNN optimizations across accelerators.

Abstract

Convolutional Neural Networks (CNNs) are central to modern AI, but their performance is often limited by hardware constraints. NVIDIA Tensor Cores, for instance, require input channels to be multiples of 8 and sometimes 512 for efficient execution. {\em oneDNN} framework for CPU imposes such a requirement for the blocked format. Traditional approaches address such alignment issue using zero-padding, which can be inefficient. In this work, we present a first-step, hardware-aware reformulation of CNN computations using rewrite rules, restructuring the underlying math to satisfy hardware alignment entirely {\bf post-training} without modifying network weights. While our current implementation focuses on a single transformation for Tensor Cores, this approach is generalizable, laying the foundation to explore additional transformations for CPU and accelerators. This study represents an initial step toward {\em semantic tuning}, a systematic, hardware-aware optimization strategy for efficient deployment of CNN models on specialized AI hardware.
Paper Structure (32 sections, 35 equations, 3 figures, 1 table, 1 algorithm)

This paper contains 32 sections, 35 equations, 3 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Width Folding Transformation
  3. Input Width Folding
  4. Filter Construction
  5. Bias Construction
  6. Resulting Convolution
  7. Width Folding: Mathematical Perspective
  8. Correctness of Width Folding Transformation
  9. Width Folding Transformation.
  10. Filter Expansion.
  11. Transformed Convolution.
  12. Conclusion.
  13. N-D Convolutions
  14. Width-Folding Transformation as an IR Transformation
  15. Motivation
  16. ...and 17 more sections

Figures (3)

  • Figure 1: Hardware-aware reformulation of a CNN via width-to-channel folding. The rewrite preserves exact semantics while satisfying Tensor Core alignment constraints without zero padding or retraining.
  • Figure 2: Semantic-preserving CNN reformulation via width folding. The input width is partitioned into $F=8$ interleaved slices and stacked along the channel dimension, increasing the effective number of input channels without altering spatial height. The original $K \times 1$ convolution kernel is replicated along the main diagonal of the expanded filter matrix, ensuring independent convolution of each folded slice. This post-training transformation preserves exact convolution semantics while aligning channel dimensions with accelerator constraints.
  • Figure : Width-Folding Transformation for Convolution