Using MLIR Transform to Design Sliced Convolution Algorithm

2025-11-22

AI Summary

To address the performance optimization challenge of 2D convolutions in MLIR, this paper proposes SConvTransform—a declarative optimization framework built upon the Transform dialect. The method introduces *convolutional slicing analysis*, which integrates static shape inference with parametric affine modeling to automatically derive optimal tiling strategies and data layouts—including packing—enabling end-to-end, structured lowering of Linalg convolutions to hardware-friendly low-level operations. Designed as a modular, reusable, and edge-case-aware framework, SConvTransform significantly improves maintainability and extensibility. Evaluation on ARM Scalable Matrix Extension (SME) and Intel AVX-512 achieves 60% and 67% of peak theoretical performance, respectively. These results validate the efficacy of static-analysis-driven, structured tiling and demonstrate strong cross-architecture portability.

Technology Category

• Machine Learning: Hardware-aware ML
• Search and Optimization: Non-convex Optimization
• Computer Vision: Learning & Optimization for CV

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Security and Privacy: Security and privacy of machine learning and AI applications

Abstract

This paper proposes SConvTransform, a Transform dialect extension that provides operations for optimizing 2D convolutions in MLIR. Its main operation, SConvOp, lowers Linalg convolutions into tiled and packed generic operations through a fully declarative transformation pipeline. The process is guided by a Convolution Slicing Analysis that determines tile sizes and data layout strategies based on input and filter shapes, as well as target architecture parameters. SConvOp handles edge cases by splitting irregular regions and adjusting affine maps where needed. All packing and tiling operations are derived from a parametric set of affine equations, enabling reusable and analyzable transformations. Although functional correctness was the primary goal of this work, the experimental evaluation demonstrates the effectiveness of SConvTransform, achieving good enough performance across different target architectures. Future work will focus on optimizing performance and porting to other target devices. When applied to standard convolution configurations, the generated code achieves up to 60% of peak performance on ARM SME and 67% on Intel AVX512. These results validate the benefit of combining static shape analysis with structured tiling and packing strategies within the MLIR Transform dialect. Furthermore, the modular design of SConvTransform facilitates integration with future extensions, enabling continued optimization of convolution workloads through MLIR's extensible compilation infrastructure.

Problem

• Optimizing 2D convolutions through MLIR Transform dialect extensions
• Lowering convolutions into tiled operations using declarative transformations
• Determining optimal tile sizes through convolution slicing analysis

Innovation

• Extends MLIR Transform dialect for convolution optimization
• Lowers convolutions into tiled operations via declarative pipeline
• Uses parametric affine equations for packing and tiling