使用 MLIR Transform 设计分片卷积算法
计算机视觉与模式识别
2025-11-25 v1 机器学习
性能
摘要
本文提出了 SConvTransform，这是一种 Transform 方言扩展，提供了用于优化 2D 卷积的运算。其主要运算 SConvOp 通过完全声明性转换管道将 Linalg 卷积降低为分块和打包的 generic 运算。该过程由一种卷积分片分析驱动，根据输入和滤波器形状以及目标架构参数确定块大小和数据布局策略。SConvOp 通过分割不规则区域并必要时调整仿射映射来处理边缘情况。所有打包和分块运算均源自参数化仿射方程组，这使得转换可重用且可分析。虽然功能正确性是本工作的首要目标，但实验评估表明 SConvTransform 的有效性，在不同目标架构上实现了相当好的性能。未来工作将致力于优化性能并移植到其他目标设备。在标准卷积配置下，生成的代码在 ARM SME 上可达峰值性能的 60%，在 Intel AVX512 上可达 67%。这些结果验证了将静态形状分析与结构化分块和打包策略结合到 MLIR Transform 方言中的好处。此外，SConvTransform 的模块化设计促进了与未来扩展的集成，使通过 MLIR 的可扩展编译基础设施继续优化卷积工作负载成为可能。
引用
@article{arxiv.2511.18222,
title = {Using MLIR Transform to Design Sliced Convolution Algorithm},
author = {Victor Ferrari and Marcio Pereira and Lucas Alvarenga and Gustavo Leite and Guido Araujo},
journal= {arXiv preprint arXiv:2511.18222},
year = {2025}
}