Technical Summary: Enabling Memory-efficient Im2win Convolution with Multi-precision Support on GPU CUDA and Tensor Cores

Problem Statement

Convolution operations constitute the primary computational bottleneck in Deep Neural Networks (DNNs), often accounting for 50–90% of execution time. Existing GPU convolution methods face significant trade-offs:

• Direct Convolution: Offers zero memory overhead but suffers from poor cache utilization.
• FFT-based Convolution: Provides speedups for large kernels but introduces latency for small filters.
• Winograd Convolution: Reduces computation for small fixed kernels but exhibits numerical instability at larger sizes.
• Im2col-based GEMM: Leverages highly optimized libraries like cuBLAS but incurs high memory overhead due to data transformation and creates irregular matrices that yield suboptimal performance.
• Implicit GEMM (cuDNN): While efficient, it often relies on mini-batch strategies that may not fully exploit hardware potential across all kernel sizes and precision levels.

The core challenge is to develop a unified convolution framework that minimizes memory overhead, maximizes cache utilization, supports various kernel sizes, and efficiently leverages both full-precision CUDA cores and half-precision Tensor Cores.

Methodology

The paper extends the im2win (image to window) paradigm, a memory-efficient convolution method that transforms input tensors to enable contiguous memory access and data reuse. The authors implement this on modern GPU architectures (specifically NVIDIA Ampere) with multi-precision support:

1. Data Transformation: Unlike im2col, which creates a highly redundant matrix by duplicating overlapping window elements, im2win arranges sliding window elements consecutively across different channels. This transformation reduces memory usage by eliminating redundant storage of overlapping elements, particularly when the stride is smaller than the filter height.
2. Multi-Precision Implementation:
   • CUDA Cores: Implements full-precision (FP32) kernels.
   • Tensor Cores: Implements half-precision (FP16) kernels using the WMMA (Warp Matrix Multiply-Accumulate) API. The design utilizes the NHWC layout to enhance memory access contiguity for convolutional windows.
3. Kernel Optimizations: To maximize hardware utilization on Tensor Cores, the authors introduce several specific optimizations:
   • Index Precomputation: Offsets are precomputed on the host and stored in constant memory to avoid redundant runtime calculations.
   • Asynchronous Data Movement: Utilizes PTX instructions to separate data movement from computation, hiding latency by overlapping data transfers with processing.
   • Zig-Zag Access: Rearranges data access order based on thread block parity to mitigate shared memory bank conflicts.
   • Double Buffering: Employs two shared memory buffers and register sets to create a producer-consumer pipeline, allowing the next data block to be fetched while the current block is being computed.
4. Execution Strategy: The implementation processes all batches simultaneously (full-batch strategy) to achieve higher parallelism, contrasting with the single-batch or mini-batch strategies often used by cuBLAS and cuDNN.

Key Contributions

The paper makes three primary contributions:

1. Multi-Precision Extension: Extends the im2win paradigm to support full-precision execution on CUDA cores and half-precision execution on Tensor Cores.
2. Optimized Kernel Design: Augments im2win with a suite of optimizations (index precomputation, Zig-Zag access, asynchronous data movement, and double buffering) specifically tailored for high-performance execution on modern GPU architectures.
3. Comprehensive Evaluation: Conducts a thorough experimental evaluation across twelve diverse CNN benchmarks, comparing the optimized im2win against PyTorch's im2col (cuBLAS) and multiple variants of cuDNN (IPG, FT, WN) on both CUDA and Tensor Cores. An ablation study is also performed to quantify the impact of individual optimization techniques.

Experimental Results

The evaluation was conducted on an NVIDIA GeForce RTX 3090 (Ampere architecture) using 12 distinct convolutional benchmarks (cv1–cv12) covering various kernel sizes and input dimensions.

• Performance (TFLOPS):
  • Tensor Cores: The im2win FP16 implementation achieved up to 2.7× higher TFLOPS than its own CUDA core (FP32) implementation. It demonstrated a 1.4× improvement over cuDNN's Tensor Core implementations and a 6.4× improvement over im2col-based convolution using cuBLAS.
  • CUDA Cores: im2win (FP32) outperformed cuDNN's implicit GEMM and Winograd variants in specific benchmarks, achieving up to 2.5× higher TFLOPS against cuDNN_WN.
• Memory Efficiency:
  • im2win significantly reduced memory consumption. On Tensor Cores, it used approximately 53% of the memory required by cuDNN and 35% of that required by im2col-cuBLAS.
  • On CUDA cores, it consistently used less memory than cuDNN variants across most benchmarks.
• Ablation Study:
  • Double Buffering was identified as the most critical optimization, providing the largest performance gains by effectively hiding data transfer latency.
  • Asynchronous Data Movement provided the second most significant improvement.
  • Zig-Zag Access yielded the most modest gains, as the other optimizations had already reduced the probability of bank conflicts.

Significance

The paper establishes im2win as a unified, high-performance convolution framework for modern GPU architectures. By addressing the high memory overhead of im2col and the poor locality of direct convolution, while simultaneously leveraging the mixed-precision capabilities of Tensor Cores, the proposed method offers a scalable solution for deep learning workloads. The results suggest that im2win can serve as a robust alternative to existing state-of-the-art libraries, particularly in scenarios where memory bandwidth and utilization of Tensor Cores are critical constraints. The work demonstrates that careful algorithmic design, combined with hardware-specific optimizations like asynchronous data movement and double buffering, can yield substantial performance and efficiency gains over established standards like cuDNN and cuBLAS.