@@ -804,18 +804,6 @@ FailureOr<linalg::ForallReductionTilingResult> linalg::tileReductionUsingForall(
804
804
return results;
805
805
}
806
806
807
- // Insert a tile `source` into the destination tensor `dest`. The position at
808
- // which the tile is inserted (as well as size of tile) is taken from a given
809
- // ExtractSliceOp `sliceOp`.
810
- static Value insertSliceIntoTensor (OpBuilder &b, Location loc,
811
- tensor::ExtractSliceOp sliceOp, Value source,
812
- Value dest) {
813
- return b.create <tensor::InsertSliceOp>(
814
- loc, sliceOp.getSource ().getType (), source, dest, sliceOp.getOffsets (),
815
- sliceOp.getSizes (), sliceOp.getStrides (), sliceOp.getStaticOffsets (),
816
- sliceOp.getStaticSizes (), sliceOp.getStaticStrides ());
817
- }
818
-
819
807
template <typename LoopTy>
820
808
FailureOr<TiledLinalgOp> static tileLinalgOpImpl (
821
809
RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options) {
@@ -851,93 +839,6 @@ mlir::linalg::tileLinalgOp(RewriterBase &b, LinalgOp op,
851
839
return failure ();
852
840
}
853
841
854
- // / Generate a loop nest around a given tensor::PadOp (for tiling). `newPadOp`
855
- // / and `loopNest` are output parameters that return the new (tiled)
856
- // / tensor::PadOp and the loop nest.
857
- static LogicalResult tilePadOp (RewriterBase &builder, tensor::PadOp op,
858
- tensor::PadOp &newPadOp, LoopNest &loopNest,
859
- const LinalgTilingOptions &options) {
860
- Location loc = op.getLoc ();
861
- OpBuilder::InsertionGuard g (builder);
862
- builder.setInsertionPoint (op);
863
-
864
- // Clone tensor::PadOp so that the existing op can be replaced more easily.
865
- newPadOp = cast<tensor::PadOp>(builder.clone (*op.getOperation ()));
866
- // Get rank and tile sizes.
867
- int64_t rank = op.getResultType ().getRank ();
868
- SmallVector<OpFoldResult> tileSizes =
869
- getAsOpFoldResult (options.tileSizeComputationFunction (builder, op));
870
- // Normalize untiled padding dimensions to 0.
871
- tileSizes.append (rank - tileSizes.size (), builder.getIndexAttr (0 ));
872
- // Compute lower and upper bounds of the loop nest.
873
- TilingInterface tilingInterface =
874
- dyn_cast<TilingInterface>(op.getOperation ());
875
- SmallVector<Range> ranges = tilingInterface.getIterationDomain (builder);
876
- SmallVector<Value> lbs, dims, steps;
877
- SmallVector<OpFoldResult> allDims;
878
- for (int64_t i = 0 ; i < rank; ++i) {
879
- allDims.push_back (ranges[i].size );
880
- if (!isZero (tileSizes[i])) {
881
- lbs.push_back (
882
- getValueOrCreateConstantIndexOp (builder, loc, ranges[i].offset ));
883
- dims.push_back (
884
- getValueOrCreateConstantIndexOp (builder, loc, ranges[i].size ));
885
- steps.push_back (
886
- getValueOrCreateConstantIndexOp (builder, loc, tileSizes[i]));
887
- }
888
- }
889
- SmallVector<Value> destinationTensors;
890
- if (failed (tensor::getOrCreateDestinations (builder, loc, tilingInterface,
891
- destinationTensors)))
892
- return failure ();
893
-
894
- loopNest = mlir::scf::buildLoopNest (
895
- builder, loc, lbs, /* ubs=*/ dims, steps, ValueRange (destinationTensors),
896
- [&](OpBuilder &b, Location loc, ValueRange localIvs,
897
- ValueRange iterArgs) -> scf::ValueVector {
898
- // Compute offsets and sizes of ExtractSliceOp.
899
- SmallVector<Value> localIVVector = llvm::to_vector (localIvs);
900
- SmallVector<OpFoldResult> offsets = computeTileOffsets (
901
- b, loc, getAsOpFoldResult (localIVVector), tileSizes);
902
- SmallVector<OpFoldResult> sizes =
903
- computeTileSizes (b, loc, tileSizes, allDims);
904
- // Create ExtractSliceOp: Extract a tile from the tensor::PadOp.
905
- // Note: The tensor::PadOp is located outside of the loop nest. It is
906
- // later moved inside by ExtractSliceOfPadTensorSwapPattern.
907
- auto map = AffineMap::getMultiDimIdentityMap (rank, b.getContext ());
908
- Value tiledOutput = makeTiledShape (
909
- b, loc, newPadOp->getResult (0 ), tileSizes, map, offsets, allDims,
910
- sizes, /* omitPartialTileCheck=*/ false );
911
- auto sliceOp = tiledOutput.getDefiningOp <tensor::ExtractSliceOp>();
912
- assert (sliceOp && " expected ExtractSliceOp" );
913
- // Insert the tile into the output tensor.
914
- Value yieldValue =
915
- insertSliceIntoTensor (b, loc, sliceOp, sliceOp, iterArgs[0 ]);
916
- return scf::ValueVector ({yieldValue});
917
- });
918
- return success ();
919
- }
920
-
921
- namespace {
922
- struct PadOpTilingPattern : public OpRewritePattern <tensor::PadOp> {
923
- PadOpTilingPattern (MLIRContext *ctx, LinalgTilingOptions opt)
924
- : OpRewritePattern<tensor::PadOp>(ctx), options(std::move(opt)) {}
925
-
926
- LogicalResult matchAndRewrite (tensor::PadOp op,
927
- PatternRewriter &rewriter) const override {
928
- tensor::PadOp newPadOp;
929
- LoopNest loopNest;
930
- if (failed (tilePadOp (rewriter, op, newPadOp, loopNest, options)))
931
- return failure ();
932
- // Replace all uses of the original tensor::PadOp.
933
- rewriter.replaceOp (op, loopNest.results .front ());
934
- return success ();
935
- }
936
-
937
- LinalgTilingOptions options;
938
- };
939
- } // namespace
940
-
941
842
namespace {
942
843
// / Helper classes for type list expansion.
943
844
template <typename ... OpTypes>
@@ -993,9 +894,3 @@ void mlir::linalg::populateLinalgTilingCanonicalizationPatterns(
993
894
#include " mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
994
895
>::insert (patterns);
995
896
}
996
-
997
- void mlir::linalg::populatePadTensorTilingPatterns (
998
- RewritePatternSet &patterns, const LinalgTilingOptions &options) {
999
- auto *ctx = patterns.getContext ();
1000
- patterns.add <PadOpTilingPattern>(ctx, options);
1001
- }
0 commit comments