diff --git a/qcom_arm64_backend/lib/OptimizeTritonShared/ARM64TilingPass.cpp b/qcom_arm64_backend/lib/OptimizeTritonShared/ARM64TilingPass.cpp index cbbc30f..7ef2029 100644 --- a/qcom_arm64_backend/lib/OptimizeTritonShared/ARM64TilingPass.cpp +++ b/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(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 decomposed = + linalg::downscaleSizeOneWindowedConvolution( + rewriter, cast(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(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(op)) + return tileAndDecomposeConv2D(cast(op)); + if (isa(op)) + return tileConv1D(cast(op)); if (isa(op)) return tileMmt4DOp(cast(op), vectorExt, numLoops); diff --git a/qcom_arm64_backend/lib/OptimizeTritonShared/ARM64VectorizationPass.cpp b/qcom_arm64_backend/lib/OptimizeTritonShared/ARM64VectorizationPass.cpp index 106dc3a..2f47911 100644 --- a/qcom_arm64_backend/lib/OptimizeTritonShared/ARM64VectorizationPass.cpp +++ b/qcom_arm64_backend/lib/OptimizeTritonShared/ARM64VectorizationPass.cpp @@ -101,7 +101,14 @@ static LogicalResult vectorizeLinalgOp(linalg::LinalgOp op, arm64VectorExtension SmallVector vecSizes(numLoops, 1); SmallVector scalableDims(numLoops, false); bool assumeDynamicDimsMatchVecSizes = false; - if (isa(op) && vectorExt == SME) { + if (isa(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(op) && vectorExt == SME) { scalableDims = {true, true, false}; vecSizes = {8, 8, 2}; } else if (isa(op) && vectorExt == SME) { diff --git a/qcom_arm64_backend/test/Transforms/arm64-tiling-pass.mlir b/qcom_arm64_backend/test/Transforms/arm64-tiling-pass.mlir index a33b227..dc613f6 100644 --- a/qcom_arm64_backend/test/Transforms/arm64-tiling-pass.mlir +++ b/qcom_arm64_backend/test/Transforms/arm64-tiling-pass.mlir @@ -174,3 +174,34 @@ func.func @test_no_tiling_zero_loops(%arg0: tensor, } -> tensor return %0 : tensor } + +// 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> +} diff --git a/qcom_arm64_backend/test/Transforms/arm64-vectorization-pass.mlir b/qcom_arm64_backend/test/Transforms/arm64-vectorization-pass.mlir index 302fc00..be42a4b 100644 --- a/qcom_arm64_backend/test/Transforms/arm64-vectorization-pass.mlir +++ b/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> +}