Skip to content

Commit 9e081a7

Browse files
committed
[mlir][linalg] Implement Winograd Conv2D.
This patch implements the Winograd Conv2D algorithm. It supports several configurations of Winograd Conv2D, including F(2, 3), F(4, 3) and F(2, 5). These configurations show that the implementation can support different kernel size (3 and 5) and different output size (2 and 4). Besides symetric kernel size 3x3 and 5x5, this patch also supports 1x3, 3x1, 1x5, and 5x1 kernels. The implementation is based on the paper, Fast Algorithm for Convolutional Neural Networks. (https://arxiv.org/abs/1509.09308)
1 parent 419e7b8 commit 9e081a7

File tree

5 files changed

+1416
-0
lines changed

5 files changed

+1416
-0
lines changed

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1692,6 +1692,9 @@ void populateTransposeMatmulPatterns(RewritePatternSet &patterns,
16921692
void populateBlockPackMatmulPatterns(RewritePatternSet &patterns,
16931693
const ControlBlockPackMatmulFn &controlFn);
16941694

1695+
/// Patterns to apply Winograd Conv2D algorithm.
1696+
void populateWinogradConv2DPatterns(RewritePatternSet &patterns);
1697+
16951698
} // namespace linalg
16961699
} // namespace mlir
16971700

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
3838
Transforms.cpp
3939
TransposeConv2D.cpp
4040
Vectorization.cpp
41+
WinogradConv2D.cpp
4142

4243
ADDITIONAL_HEADER_DIRS
4344
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg

0 commit comments

Comments
 (0)