From 0c6c63d4393bdada02b801e4515f0b919fd5367a Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Fri, 15 Nov 2024 20:18:47 +0000 Subject: [PATCH] [mlir] Rename `GeneralizeOuterUnitDims{Un}PackOpPattern`s Renames: * `GeneralizeOuterUnitDimsPackOpPattern`, * `GeneralizeOuterUnitDimsUnPackOpPattern`, as * `DecomposeOuterUnitDimsPackOpPattern`, * `OuterUnitDimsUnPackOpPattern`, respectively. The new name better describes the underlying transformation. **NOTE** Depends on #116373 (please only review the top commit) --- .../Linalg/TransformOps/LinalgTransformOps.td | 9 ++++--- .../Dialect/Linalg/Transforms/Transforms.h | 6 ++--- .../TransformOps/LinalgTransformOps.cpp | 4 ++-- .../Dialect/Linalg/Transforms/Transforms.cpp | 8 +++---- ...e.mlir => decompose-tensor-pack-tile.mlir} | 5 +++- ...r-pack.mlir => decompose-tensor-pack.mlir} | 4 +++- ...mlir => decompose-tensor-unpack-tile.mlir} | 2 +- ...pack.mlir => decompose-tensor-unpack.mlir} | 2 +- ...neralize-pack.mlir => decompose-pack.mlir} | 4 ++-- .../Linalg/CPU/pack-dynamic-inner-tile.mlir | 2 +- .../Dialect/Linalg/TestLinalgTransforms.cpp | 24 +++++++++---------- 11 files changed, 37 insertions(+), 33 deletions(-) rename mlir/test/Dialect/Linalg/{generalize-tensor-pack-tile.mlir => decompose-tensor-pack-tile.mlir} (94%) rename mlir/test/Dialect/Linalg/{generalize-tensor-pack.mlir => decompose-tensor-pack.mlir} (98%) rename mlir/test/Dialect/Linalg/{generalize-tensor-unpack-tile.mlir => decompose-tensor-unpack-tile.mlir} (98%) rename mlir/test/Dialect/Linalg/{generalize-tensor-unpack.mlir => decompose-tensor-unpack.mlir} (99%) rename mlir/test/Dialect/Linalg/td/{generalize-pack.mlir => decompose-pack.mlir} (70%) diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td index 42057d8d0c910..e3084530bd11b 100644 --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td @@ -41,13 +41,12 @@ def ApplyEraseUnnecessaryInputsPatternsOp : Op]> { let description = [{ - Collect patterns to generalize tensor.pack and tensor.unpack (i.e. to - decompose it into e.g. tensor::PadOp, linalg::transposeOp etc). Requires - all outer dims to be unit. + Collect patterns to decompose tensor.pack and tensor.unpack into e.g. + tensor::PadOp, linalg::transposeOp Ops. Requires all outer dims to be unit. }]; let assemblyFormat = "attr-dict"; diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h index 0b55a76f88433..51967f83fee37 100644 --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h @@ -1548,7 +1548,7 @@ struct GeneralizePadOpPattern : public OpRewritePattern { /// into %arg1[0, 0, 0, 0] [1, 1, 2, %tile_dim_1] [1, 1, 1, 1] /// : tensor<2x?xf32> into tensor<1x1x2x?xf32> /// ``` -struct GeneralizeOuterUnitDimsPackOpPattern +struct DecomposeOuterUnitDimsPackOpPattern : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(tensor::PackOp packOp, @@ -1558,7 +1558,7 @@ struct GeneralizeOuterUnitDimsPackOpPattern /// Rewrites a tensor::UnPackOp into a sequence of rank-reduced extract_slice op /// + transpose op + insert_slice op, where the tensor::UnPackOp has outer dims /// being all 1s. -struct GeneralizeOuterUnitDimsUnPackOpPattern +struct DecomposeOuterUnitDimsUnPackOpPattern : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(tensor::UnPackOp unpackOp, @@ -1686,7 +1686,7 @@ void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns, /// Populates patterns to decompose tensor.pack and tensor.unpack Ops into e.g. /// tensor.pad, linalg.transpose, tensor.{insert|extract}_slice. Require all /// outer dims to be unit. -void populateGeneralizePatterns(RewritePatternSet &patterns); +void populateDecomposePackUnpackPatterns(RewritePatternSet &patterns); /// Populates patterns to transform linalg.conv_2d_xxx operations into /// linalg.generic (for img2col packing) and linalg.matmul. diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp index a00c609779c3a..ada80deacfdbf 100644 --- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp +++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp @@ -229,9 +229,9 @@ void transform::ApplyEraseUnnecessaryInputsPatternsOp::populatePatterns( linalg::populateEraseUnnecessaryInputsPatterns(patterns); } -void transform::ApplyGeneralizeTensorPackUnpackPatternsOp::populatePatterns( +void transform::ApplyDecomposeTensorPackUnpackPatternsOp::populatePatterns( RewritePatternSet &patterns) { - linalg::populateGeneralizePatterns(patterns); + linalg::populateDecomposePackUnpackPatterns(patterns); } void transform::ApplyFoldUnitExtentDimsViaReshapesPatternsOp::populatePatterns( diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp index c9eac66367559..d92543d726462 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -1138,7 +1138,7 @@ getPackUnpackRankReducedPerm(ArrayRef shape, return perm; } -LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite( +LogicalResult DecomposeOuterUnitDimsPackOpPattern::matchAndRewrite( tensor::PackOp packOp, PatternRewriter &rewriter) const { // TODO: support the case that outer dimensions are not all 1s. A // tensor.expand_shape will be generated in this case. @@ -1239,7 +1239,7 @@ LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite( return success(); } -LogicalResult GeneralizeOuterUnitDimsUnPackOpPattern::matchAndRewrite( +LogicalResult DecomposeOuterUnitDimsUnPackOpPattern::matchAndRewrite( tensor::UnPackOp unpackOp, PatternRewriter &rewriter) const { int64_t srcRank = unpackOp.getSourceRank(); int64_t destRank = unpackOp.getDestRank(); @@ -1619,7 +1619,7 @@ void linalg::populateDecomposeConvolutionPatterns(RewritePatternSet &patterns, patterns.getContext(), benefit); } -void linalg::populateGeneralizePatterns(RewritePatternSet &patterns) { +void linalg::populateDecomposePackUnpackPatterns(RewritePatternSet &patterns) { // TODO: Add and test patterns for tensor.unpack - patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); } diff --git a/mlir/test/Dialect/Linalg/generalize-tensor-pack-tile.mlir b/mlir/test/Dialect/Linalg/decompose-tensor-pack-tile.mlir similarity index 94% rename from mlir/test/Dialect/Linalg/generalize-tensor-pack-tile.mlir rename to mlir/test/Dialect/Linalg/decompose-tensor-pack-tile.mlir index 1fae311467bcf..ec761d9a49436 100644 --- a/mlir/test/Dialect/Linalg/generalize-tensor-pack-tile.mlir +++ b/mlir/test/Dialect/Linalg/decompose-tensor-pack-tile.mlir @@ -1,4 +1,7 @@ -// RUN: mlir-opt -split-input-file --transform-interpreter --canonicalize --test-linalg-transform-patterns="test-generalize-tensor-pack" %s | FileCheck %s +// RUN: mlir-opt -split-input-file -transform-interpreter --canonicalize \ +// RUN: -transform-preload-library='transform-library-paths=%p/td/decompose-pack.mlir' \ +// RUN: -transform-interpreter=entry-point=decompose_pack \ +// RUN: -transform-interpreter %s | FileCheck %s func.func @KCRS_to_KCRSsr(%arg0: tensor<1x1x128x64xf32>, %arg1: tensor<1x1x4x8x8x32xf32>) -> tensor<1x1x4x8x8x32xf32> { %0 = tensor.pack %arg0 inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<1x1x128x64xf32> -> tensor<1x1x4x8x8x32xf32> diff --git a/mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir b/mlir/test/Dialect/Linalg/decompose-tensor-pack.mlir similarity index 98% rename from mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir rename to mlir/test/Dialect/Linalg/decompose-tensor-pack.mlir index ad20541e301d3..4f986606ef93a 100644 --- a/mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir +++ b/mlir/test/Dialect/Linalg/decompose-tensor-pack.mlir @@ -1,4 +1,6 @@ -// RUN: mlir-opt --transform-preload-library='transform-library-paths=%p/td/generalize-pack.mlir' -split-input-file --transform-interpreter %s | FileCheck %s +// RUN: mlir-opt -split-input-file \ +// RUN: -transform-preload-library='transform-library-paths=%p/td/decompose-pack.mlir' \ +// RUN: -transform-interpreter=entry-point=decompose_pack %s | FileCheck %s func.func @simple_KCRS_to_KCRSsr(%arg0: tensor, %arg1: tensor<1x1x?x1xi32>) -> tensor<1x1x?x1xi32> { %c8 = arith.constant 8 : index diff --git a/mlir/test/Dialect/Linalg/generalize-tensor-unpack-tile.mlir b/mlir/test/Dialect/Linalg/decompose-tensor-unpack-tile.mlir similarity index 98% rename from mlir/test/Dialect/Linalg/generalize-tensor-unpack-tile.mlir rename to mlir/test/Dialect/Linalg/decompose-tensor-unpack-tile.mlir index c15859d898ec1..6d9709caf7093 100644 --- a/mlir/test/Dialect/Linalg/generalize-tensor-unpack-tile.mlir +++ b/mlir/test/Dialect/Linalg/decompose-tensor-unpack-tile.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt -split-input-file --transform-interpreter --canonicalize --test-linalg-transform-patterns="test-generalize-tensor-unpack" %s | FileCheck %s +// RUN: mlir-opt -split-input-file --transform-interpreter --canonicalize --test-linalg-transform-patterns="test-decompose-tensor-unpack" %s | FileCheck %s func.func @KCRSsr_to_KCRS(%arg0: tensor<1x1x4x8x8x32xf32>, %arg1: tensor<1x1x128x64xf32>) -> tensor<1x1x128x64xf32> { %0 = tensor.unpack %arg0 inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<1x1x4x8x8x32xf32> -> tensor<1x1x128x64xf32> diff --git a/mlir/test/Dialect/Linalg/generalize-tensor-unpack.mlir b/mlir/test/Dialect/Linalg/decompose-tensor-unpack.mlir similarity index 99% rename from mlir/test/Dialect/Linalg/generalize-tensor-unpack.mlir rename to mlir/test/Dialect/Linalg/decompose-tensor-unpack.mlir index 153ce68b8f086..8b15873473a97 100644 --- a/mlir/test/Dialect/Linalg/generalize-tensor-unpack.mlir +++ b/mlir/test/Dialect/Linalg/decompose-tensor-unpack.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-generalize-tensor-unpack" %s | FileCheck %s +// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-decompose-tensor-unpack" %s | FileCheck %s func.func @simple_KCRSsr_to_KCRS(%arg0: tensor<1x1x1x1x8x32xf32>, %arg1: tensor<1x1x32x8xf32>) -> tensor<1x1x32x8xf32> { %0 = tensor.unpack %arg0 inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<1x1x1x1x8x32xf32> -> tensor<1x1x32x8xf32> diff --git a/mlir/test/Dialect/Linalg/td/generalize-pack.mlir b/mlir/test/Dialect/Linalg/td/decompose-pack.mlir similarity index 70% rename from mlir/test/Dialect/Linalg/td/generalize-pack.mlir rename to mlir/test/Dialect/Linalg/td/decompose-pack.mlir index 62e5b779ff361..49c45e29d5a14 100644 --- a/mlir/test/Dialect/Linalg/td/generalize-pack.mlir +++ b/mlir/test/Dialect/Linalg/td/decompose-pack.mlir @@ -1,10 +1,10 @@ module @transforms attributes { transform.with_named_sequence } { - transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) { + transform.named_sequence @decompose_pack(%module: !transform.any_op {transform.readonly}) { %pack = transform.structured.match ops{["tensor.pack"]} in %module : (!transform.any_op) -> !transform.any_op %1 = transform.get_parent_op %pack {isolated_from_above} : (!transform.any_op) -> !transform.any_op transform.apply_patterns to %1 { - transform.apply_patterns.linalg.generalize_pack_unpack + transform.apply_patterns.linalg.decompose_pack_unpack } : !transform.any_op transform.yield diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir index bec1b9a4e9d82..0428ada86041d 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir @@ -1,6 +1,6 @@ // DEFINE: %{compile} = mlir-opt %s \ // DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule |\ -// DEFINE: mlir-opt --test-linalg-transform-patterns="test-generalize-tensor-pack"\ +// DEFINE: mlir-opt --test-linalg-transform-patterns="test-decompose-tensor-pack"\ // DEFINE: --test-transform-dialect-erase-schedule \ // DEFINE: -one-shot-bufferize="bufferize-function-boundaries" \ // DEFINE: -buffer-deallocation-pipeline="private-function-dynamic-ownership" \ diff --git a/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp b/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp index 5899f56da7345..c65e68eaf31f0 100644 --- a/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp +++ b/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp @@ -74,13 +74,13 @@ struct TestLinalgTransforms *this, "test-generalize-pad-tensor", llvm::cl::desc("Test transform pad tensor by copying with generic ops"), llvm::cl::init(false)}; - Option testGeneralizeTensorPackOp{ - *this, "test-generalize-tensor-pack", + Option testDecomposeTensorPackOp{ + *this, "test-decompose-tensor-pack", llvm::cl::desc("Test transform that generalizes pack ops into a sequence " "of tensor and Linalg ops"), llvm::cl::init(false)}; - Option testGeneralizeTensorUnPackOp{ - *this, "test-generalize-tensor-unpack", + Option testDecomposeTensorUnPackOp{ + *this, "test-decompose-tensor-unpack", llvm::cl::desc( "Test transform that generalizes unpack ops into a sequence " "of tensor and Linalg ops"), @@ -172,15 +172,15 @@ static void applyGeneralizePadTensorPatterns(func::FuncOp funcOp) { (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns)); } -static void applyGeneralizeTensorPackPatterns(func::FuncOp funcOp) { +static void applyDecomposeTensorPackPatterns(func::FuncOp funcOp) { RewritePatternSet patterns(funcOp.getContext()); - patterns.add(funcOp.getContext()); + patterns.add(funcOp.getContext()); (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns)); } -static void applyGeneralizeTensorUnPackPatterns(func::FuncOp funcOp) { +static void applyDecomposeTensorUnPackPatterns(func::FuncOp funcOp) { RewritePatternSet patterns(funcOp.getContext()); - patterns.add(funcOp.getContext()); + patterns.add(funcOp.getContext()); (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns)); } @@ -237,10 +237,10 @@ void TestLinalgTransforms::runOnOperation() { return applyLinalgToVectorPatterns(getOperation()); if (testGeneralizePadTensor) return applyGeneralizePadTensorPatterns(getOperation()); - if (testGeneralizeTensorPackOp) - return applyGeneralizeTensorPackPatterns(getOperation()); - if (testGeneralizeTensorUnPackOp) - return applyGeneralizeTensorUnPackPatterns(getOperation()); + if (testDecomposeTensorPackOp) + return applyDecomposeTensorPackPatterns(getOperation()); + if (testDecomposeTensorUnPackOp) + return applyDecomposeTensorUnPackPatterns(getOperation()); if (testSwapSubTensorPadTensor) return applyExtractSliceOfPadTensorSwapPattern(getOperation()); if (testBubbleUpExtractSliceOpPattern)