Lower Conv2D to direct convolution

qcom_arm64_backend/lib/OptimizeTritonShared/ARM64TilingPass.cpp

@@ -98,6 +98,55 @@ static LogicalResult tileMmt4DOp(linalg::Mmt4DOp op, arm64VectorExtension vector
   return success();
 }
 
+// Tile height dimensions, OH and KH, then decompose NCHW/FCHW convolution to
+// NCW/FCW.
+static LogicalResult tileAndDecomposeConv2D(linalg::Conv2DNchwFchwOp op) {
+  IRRewriter rewriter(op.getContext());
+  rewriter.setInsertionPoint(op);
+
+  OpFoldResult zero = rewriter.getIndexAttr(0);
+  OpFoldResult one = rewriter.getIndexAttr(1);
+  scf::SCFTilingOptions options;
+  // Loop order: N, F, OH, OW, C, KH, KW.
+  options.setTileSizes({zero, zero, one, zero, zero, one, zero});
+
+  auto tiled = scf::tileUsingSCF(
+      rewriter, cast<TilingInterface>(op.getOperation()), options);
+  if (failed(tiled) || tiled->tiledOps.size() != 1)
+    return failure();
+  rewriter.replaceOp(op, tiled->replacements);
+
+  Operation *tiledOp = tiled->tiledOps.front();
+  rewriter.setInsertionPoint(tiledOp);
+  // The tiled convolution has unit OH and KH dimensions.
+  FailureOr<linalg::LinalgOp> decomposed =
+      linalg::downscaleSizeOneWindowedConvolution(
+          rewriter, cast<linalg::LinalgOp>(tiledOp));
+  if (failed(decomposed))
+    return failure();
+  return success();
+}
+
+// Tile F and OW, and split C and KW, before vectorization.
+static LogicalResult tileConv1D(linalg::Conv1DNcwFcwOp op) {
+  IRRewriter rewriter(op.getContext());
+  rewriter.setInsertionPoint(op);
+
+  OpFoldResult zero = rewriter.getIndexAttr(0);
+  OpFoldResult one = rewriter.getIndexAttr(1);
+  scf::SCFTilingOptions options;
+  // Loop order: N, F, OW, C, KW.
+  options.setTileSizes({zero, rewriter.getIndexAttr(4),
+                        rewriter.getIndexAttr(4), one, one});
+
+  auto tiled = scf::tileUsingSCF(
+      rewriter, cast<TilingInterface>(op.getOperation()), options);
+  if (failed(tiled))
+    return failure();
+  rewriter.replaceOp(op, tiled->replacements);
+  return success();
+}
+
 static LogicalResult tileLinalgOp(linalg::LinalgOp op, arm64VectorExtension vectorExt) {
   // Early exit.
   unsigned numLoops = op.getNumLoops();
@@ -108,6 +157,10 @@ static LogicalResult tileLinalgOp(linalg::LinalgOp op, arm64VectorExtension vect
   }
 
   // Specialized tiling for certain ops
+  if (isa<linalg::Conv2DNchwFchwOp>(op))
+    return tileAndDecomposeConv2D(cast<linalg::Conv2DNchwFchwOp>(op));
+  if (isa<linalg::Conv1DNcwFcwOp>(op))
+    return tileConv1D(cast<linalg::Conv1DNcwFcwOp>(op));
   if (isa<linalg::Mmt4DOp>(op))
     return tileMmt4DOp(cast<linalg::Mmt4DOp>(op), vectorExt, numLoops);
 

qcom_arm64_backend/lib/OptimizeTritonShared/ARM64VectorizationPass.cpp

@@ -101,7 +101,14 @@ static LogicalResult vectorizeLinalgOp(linalg::LinalgOp op, arm64VectorExtension
   SmallVector<int64_t> vecSizes(numLoops, 1);
   SmallVector<bool> scalableDims(numLoops, false);
   bool assumeDynamicDimsMatchVecSizes = false;
-  if (isa<linalg::MatmulOp>(op) && vectorExt == SME) {
+  if (isa<linalg::Conv1DNcwFcwOp>(op)) {
+    // The preceding Conv2D schedule decomposes one output/filter-height tile
+    // to NCW/FCW. Vectorize F and OW as a 4x4 outer product while retaining
+    // C and KW as scalar reduction tiles.
+    if (numLoops != 5)
+      return failure();
+    vecSizes = {1, 4, 4, 1, 1};
+  } else if (isa<linalg::MatmulOp>(op) && vectorExt == SME) {
     scalableDims = {true, true, false};
     vecSizes = {8, 8, 2};
   } else if (isa<linalg::Mmt4DOp>(op) && vectorExt == SME) {

qcom_arm64_backend/test/Transforms/arm64-tiling-pass.mlir

@@ -174,3 +174,34 @@ func.func @test_no_tiling_zero_loops(%arg0: tensor<f32>,
   } -> tensor<f32>
   return %0 : tensor<f32>
 }
+
+// CHECK-LABEL: func.func @test_conv1d_outer_product_tiling
+// CHECK:       scf.for {{.+}} step {{.+}}4
+// CHECK:       scf.for {{.+}} step {{.+}}4
+// CHECK:       scf.for {{.+}} step {{.+}}1
+// CHECK:       scf.for {{.+}} step {{.+}}1
+// CHECK:       linalg.conv_1d_ncw_fcw
+func.func @test_conv1d_outer_product_tiling(
+    %input: tensor<1x3x6xf32>, %filter: tensor<4x3x3xf32>,
+    %output: tensor<1x4x4xf32>) -> tensor<1x4x4xf32> {
+  %0 = linalg.conv_1d_ncw_fcw
+    ins(%input, %filter : tensor<1x3x6xf32>, tensor<4x3x3xf32>)
+    outs(%output : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
+  return %0 : tensor<1x4x4xf32>
+}
+
+// CHECK-LABEL: func.func @test_conv2d_decomposes_to_conv1d
+// CHECK:       scf.for
+// CHECK:       scf.for
+// CHECK-NOT:   linalg.conv_2d_nchw_fchw
+// CHECK:       linalg.conv_1d_ncw_fcw
+// CHECK:       return
+func.func @test_conv2d_decomposes_to_conv1d(
+    %input: tensor<1x3x5x6xf32>,
+    %filter: tensor<4x3x3x3xf32>,
+    %output: tensor<1x4x3x4xf32>) -> tensor<1x4x3x4xf32> {
+  %0 = linalg.conv_2d_nchw_fchw
+    ins(%input, %filter : tensor<1x3x5x6xf32>, tensor<4x3x3x3xf32>)
+    outs(%output : tensor<1x4x3x4xf32>) -> tensor<1x4x3x4xf32>
+  return %0 : tensor<1x4x3x4xf32>
+}

qcom_arm64_backend/test/Transforms/arm64-vectorization-pass.mlir

@@ -1,5 +1,6 @@
 // RUN: arm64-mlir-opt %s -triton-arm64-linalg-vectorization="vector-ext=neon" | FileCheck %s
 // RUN: arm64-mlir-opt %s -triton-arm64-linalg-vectorization | FileCheck %s --check-prefix=CHECK-SME
+// RUN: arm64-mlir-opt %s -triton-arm64-linalg-tiling="vector-ext=neon" -triton-arm64-linalg-vectorization="vector-ext=neon" | FileCheck %s --check-prefix=CHECK-CONV
 
 // CHECK-LABEL: func.func @test_elementwise_static
 // CHECK-NOT:   linalg.map
@@ -192,3 +193,20 @@ func.func @test_library_call_skipped_generic(%arg0: tensor<8xf32>,
   } -> tensor<8xf32>
   return %0 : tensor<8xf32>
 }
+
+// CHECK-CONV-LABEL: func.func @test_conv1d_outer_product
+// CHECK-CONV:       scf.for
+// CHECK-CONV:       vector.transfer_read {{.+}} : tensor<1x1x4xf32>, vector<1x1x4xf32>
+// CHECK-CONV:       vector.transfer_read {{.+}} : tensor<4x1x1xf32>, vector<4x1x1xf32>
+// CHECK-CONV:       vector.transfer_read {{.+}} : tensor<1x4x4xf32>, vector<1x4x4xf32>
+// CHECK-CONV:       vector.contract
+// CHECK-CONV:       vector.transfer_write {{.+}} : vector<1x4x4xf32>, tensor<1x4x4xf32>
+// CHECK-CONV-NOT:   linalg.conv_1d_ncw_fcw
+func.func @test_conv1d_outer_product(%input: tensor<1x3x6xf32>,
+                                     %filter: tensor<4x3x3xf32>,
+                                     %output: tensor<1x4x4xf32>) -> tensor<1x4x4xf32> {
+  %0 = linalg.conv_1d_ncw_fcw
+    ins(%input, %filter : tensor<1x3x6xf32>, tensor<4x3x3xf32>)
+    outs(%output : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
+  return %0 : tensor<1x4x4xf32>
+}