Using MLIR Transform to Design Sliced Convolution Algorithm

TL;DR

This paper introduces SConvTransform, an MLIR extension for optimizing 2D convolutions through declarative tiling and packing strategies guided by shape analysis, achieving significant performance on ARM and Intel architectures.

Contribution

It presents a novel MLIR Transform dialect extension, SConvTransform, that automates convolution optimization via a parametric, analysis-guided tiling and packing pipeline.

Findings

Achieves up to 60% of peak performance on ARM SME.

Achieves up to 67% of peak performance on Intel AVX512.

Demonstrates effectiveness of static shape analysis combined with structured tiling.

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.