Revert "Revert "[mlir] Purge linalg.copy and use memref.copy instead.""

This reverts commit 25bf6a2.

Author: pifon2a

mlir/docs/Dialects/Linalg/_index.md

@@ -545,7 +545,6 @@ seem generally appealing.
 Additionally, `linalg` provides a small subset of commonly named operations:
 
 ```
-* `linalg.copy`,
 * `linalg.fill`,
 * `linalg.dot`,
 * `linalg.matmul`,

mlir/include/mlir/Conversion/LinalgToStandard/LinalgToStandard.h

@@ -39,25 +39,6 @@ class LinalgOpToLibraryCallRewrite
                                 PatternRewriter &rewriter) const override;
 };
 
-/// Rewrite pattern specialization for CopyOp, kicks in when both input and
-/// output permutations are left unspecified or are the identity.
-class CopyOpToLibraryCallRewrite : public OpRewritePattern<CopyOp> {
-public:
-  using OpRewritePattern<CopyOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(CopyOp op,
-                                PatternRewriter &rewriter) const override;
-};
-
-/// Rewrite CopyOp with permutations into a sequence of TransposeOp and
-/// permutation-free CopyOp. This interplays with TransposeOpConversion and
-/// LinalgConversion<CopyOp> to create a path to the LLVM dialect.
-class CopyTransposeRewrite : public OpRewritePattern<CopyOp> {
-public:
-  using OpRewritePattern<CopyOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(CopyOp op,
-                                PatternRewriter &rewriter) const override;
-};
-
 /// Populate the given list with patterns that convert from Linalg to Standard.
 void populateLinalgToStandardConversionPatterns(RewritePatternSet &patterns);
 

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

@@ -42,7 +42,7 @@ def BufferDeallocation : Pass<"buffer-deallocation", "FuncOp"> {
         }: memref<2xf32>, memref<2xf32>
         br ^bb3(%0 : memref<2xf32>)
       ^bb3(%1: memref<2xf32>):
-        "linalg.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
+        "memref.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
         return
       }
     }
@@ -58,7 +58,7 @@ def BufferDeallocation : Pass<"buffer-deallocation", "FuncOp"> {
         cond_br %arg0, ^bb1, ^bb2
       ^bb1:  // pred: ^bb0
         %0 = memref.alloc() : memref<2xf32>
-        linalg.copy(%arg1, %0) : memref<2xf32>, memref<2xf32>
+        memref.copy(%arg1, %0) : memref<2xf32>, memref<2xf32>
         br ^bb3(%0 : memref<2xf32>)
       ^bb2:  // pred: ^bb0
         %1 = memref.alloc() : memref<2xf32>
@@ -72,11 +72,11 @@ def BufferDeallocation : Pass<"buffer-deallocation", "FuncOp"> {
           linalg.yield %4 : f32
         }: memref<2xf32>, memref<2xf32>
         %2 = memref.alloc() : memref<2xf32>
-        linalg.copy(%1, %2) : memref<2xf32>, memref<2xf32>
+        memref.copy(%1, %2) : memref<2xf32>, memref<2xf32>
         dealloc %1 : memref<2xf32>
         br ^bb3(%2 : memref<2xf32>)
       ^bb3(%3: memref<2xf32>):  // 2 preds: ^bb1, ^bb2
-        linalg.copy(%3, %arg2) : memref<2xf32>, memref<2xf32>
+        memref.copy(%3, %arg2) : memref<2xf32>, memref<2xf32>
         dealloc %3 : memref<2xf32>
         return
       }

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

@@ -16,7 +16,6 @@
 
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
-include "mlir/Interfaces/CopyOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
@@ -57,119 +56,6 @@ class LinalgStructured_Op<string mnemonic, list<Trait> props>
 //===----------------------------------------------------------------------===//
 // Named Linalg ops, implemented as special configurations of generic ops.
 //===----------------------------------------------------------------------===//
-// At the moment these are not declarative and require a bunch of C++ code.
-// In the future, these should be migrated to a declarative specification.
-def CopyOp : LinalgStructured_Op<"copy", [CopyOpInterface]> {
-  let description = [{
-    Copies the data in the input view into the output view.
-
-    Usage:
-
-    ```mlir
-    linalg.copy(%arg0, %arg1) : memref<?xf32, stride_specification>,
-                                memref<?xf32, stride_specification>
-    ```
-
-    One possible lowering to loop form is:
-
-    ```mlir
-    %0 = linalg.dim %arg0, 0 : index
-    scf.for %i0 = %c0 to %0 step %c1 {
-      %1 = load %arg0[%i0] : memref<?xf32, stride_specification>
-      store %1, %arg1[%i0] : memref<?xf32, stride_specification>
-    }
-    ```
-
-    Optionally, can take `input_permutation` and `output_permutation` attributes
-    to reorder the dimensions of the input and output views.
-
-    Usage:
-
-    ```mlir
-    linalg.copy(%arg0, %arg1) {inputPermutation : (i, j, k) -> (i, k, j),
-                               outputPermutation : (i, j, k) -> (k, j, i)} :
-      memref<?x?x?xf32, stride_specification>,
-      memref<?x?x?xf32, stride_specification>
-    ```
-
-    One possible lowering to loop form is:
-
-    ```mlir
-    %0 = linalg.dim %arg0, 0
-    %1 = linalg.dim %arg0, 1
-    %2 = linalg.dim %arg0, 2
-    scf.for %i0 = %c0 to %{{.*}} step %c1 {
-      scf.for %i1 = %c0 to %{{.*}} step %c1 {
-        scf.for %i2 = %c0 to %{{.*}} step %c1 {
-          %3 = load %arg0[%i0, %i2, %i1] :
-                  memref<?x?x?xf32, stride_specification>
-          store %3, %arg1[%i2, %i1, %i0] :
-                  memref<?x?x?xf32, stride_specification>
-    ```
-
-    The views are expected to be compatible for correctness but this is not
-    enforced at the moment.
-  }];
-
-  let arguments = (ins
-    AnyStridedMemRef:$input,
-    AnyStridedMemRef:$output,
-    OptionalAttr<AffineMapAttr>:$inputPermutation,
-    OptionalAttr<AffineMapAttr>:$outputPermutation);
-  let regions = (region AnyRegion:$region);
-
-  let builders = [
-    OpBuilder<(ins "Value":$input, "Value":$output,
-      CArg<"AffineMap", "AffineMap()">:$inputPermutation,
-      CArg<"AffineMap", "AffineMap()">:$outputPermutation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>];
-
-  let extraClassDeclaration = structuredOpsDecls # [{
-    ValueRange inputs() { return getOperands().take_front(); }
-    ValueRange outputs() { return getOperands().take_back(); }
-
-    // Rank-polymorphic.
-    //   filling_value -> O(ivs) with parallel iterators.
-    ArrayAttr iterator_types() {
-      int64_t nPar = getRank(getInputOperand(0));
-      return Builder(getContext()).getStrArrayAttr(
-        SmallVector<StringRef, 8>(nPar, getParallelIteratorTypeName()));
-    }
-
-    // I(input_perm(ivs)) -> O(output_perm(ivs))
-    ArrayAttr indexing_maps() {
-      MLIRContext *context = getContext();
-      auto maybeInputMap = inputPermutation();
-      auto maybeOutputMap = outputPermutation();
-      int64_t inputRank = getRank(getInputOperand(0));
-      int64_t outputRank = getRank(getOutputOperand(0));
-      return Builder(getContext()).getAffineMapArrayAttr({
-          extractOrIdentityMap(maybeInputMap, inputRank, context),
-          extractOrIdentityMap(maybeOutputMap, outputRank, context)});
-    }
-
-    Value getSource() { return input();}
-    Value getTarget() { return output(); }
-
-    static void regionBuilder(ImplicitLocOpBuilder &b, Block &block);
-    static std::function<void(ImplicitLocOpBuilder &b, Block &block)>
-    getRegionBuilder() {
-      return &regionBuilder;
-    }
-    static unsigned getNumRegionArgs() { return 2; }
-  }];
-  let verifier = [{ return ::verify(*this); }];
-
-  let assemblyFormat = [{
-    `(` $input `,` $output `)` attr-dict `:`
-        type($input) `,` type($output)
-      custom<CopyOpRegion>($region, ref(type($input)), ref(type($input)))
-  }];
-
-  let hasCanonicalizer = 1;
-  let hasFolder = 1;
-  let skipDefaultBuilders = 1;
-}
 
 def FillOp : LinalgStructured_Op<"fill", []> {
   let arguments = (ins

mlir/include/mlir/Dialect/Linalg/Passes.td

@@ -52,7 +52,7 @@ def LinalgComprehensiveModuleBufferize :
     Option<"useAlloca", "use-alloca", "bool",
            /*default=*/"false",
            "Use stack allocations for memrefs (for testing purposes only)">,
-    Option<"useLinalgCopy", "use-linalg-copy", "bool",
+    Option<"useLinalgCopy", "use-memref.copy", "bool",
            /*default=*/"false",
            "Use a copy operation implemented as a Linalg op.">,
     Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",

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

@@ -349,7 +349,7 @@ struct LinalgPromotionOptions {
     return *this;
   }
   /// Callback function to do the copy of data to and from the promoted
-  /// subview. If None then a linalg.copy is used.
+  /// subview. If None then a memref.copy is used.
   Optional<CopyCallbackFn> copyInFn = None;
   Optional<CopyCallbackFn> copyOutFn = None;
   LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
@@ -390,6 +390,9 @@ FailureOr<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
 /// Emit a suitable vector form for a Linalg op with fully static shape.
 LogicalResult vectorize(RewriterBase &builder, LinalgOp linalgOp);
 
+/// Emit a suitable vector form for a Copy op with fully static shape.
+LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp);
+
 /// Emit a loop nest of `scf.for` with the proper body for `linalgOp`.
 FailureOr<LinalgLoops> linalgOpToLoops(PatternRewriter &rewriter,
                                        LinalgOp linalgOp);
@@ -934,6 +937,15 @@ struct LinalgVectorizationPattern : public OpInterfaceRewritePattern<LinalgOp> {
   LinalgTransformationFilter filter;
 };
 
+/// `filter` controls LinalgTransformMarker matching and update when specified.
+/// See `vectorizeLinalgOp` for more details.
+struct CopyVectorizationPattern : public OpRewritePattern<memref::CopyOp> {
+  using OpRewritePattern<memref::CopyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::CopyOp copyOp,
+                                PatternRewriter &rewriter) const override;
+};
+
 /// Return vector::CombiningKind for the given op.
 llvm::Optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
 
@@ -1206,7 +1218,7 @@ void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
 ///    %subView = subview %allocOrView ...
 ///    [optional] linalg.fill(%allocOrView, %cst) ...
 ///    ...
-///    linalg.copy(%in, %subView) ...
+///    memref.copy(%in, %subView) ...
 ///    vector.transfer_read %allocOrView[...], %cst ...
 /// ```
 /// into
@@ -1217,8 +1229,8 @@ void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
 ///    ...
 ///    vector.transfer_read %in[...], %cst ...
 /// ```
-/// Where there is no interleaved use between linalg.copy and transfer_read as
-/// well as no interleaved use between linalg.fill and linalg.copy (if
+/// Where there is no interleaved use between memref.copy and transfer_read as
+/// well as no interleaved use between linalg.fill and memref.copy (if
 /// linalg.fill is specified).
 /// This is a custom rewrite to forward partial reads (with optional fills) to
 /// vector.transfer_read.
@@ -1237,7 +1249,7 @@ struct LinalgCopyVTRForwardingPattern
 ///    %subView = subview %allocOrView...
 ///    ...
 ///    vector.transfer_write %..., %allocOrView[...]
-///    linalg.copy(%subView, %out)
+///    memref.copy(%subView, %out)
 /// ```
 /// into
 /// ```
@@ -1247,7 +1259,7 @@ struct LinalgCopyVTRForwardingPattern
 ///    ...
 ///    vector.transfer_write %..., %out[...]
 /// ```
-/// Where there is no interleaved use between transfer_write and linalg.copy.
+/// Where there is no interleaved use between transfer_write and memref.copy.
 /// This is a custom rewrite to forward partial writes to vector.transfer_write.
 struct LinalgCopyVTWForwardingPattern
     : public OpRewritePattern<vector::TransferWriteOp> {

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

@@ -407,6 +407,7 @@ def CopyOp : MemRef_Op<"copy",
   }];
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
   let verifier = ?;
 }
 

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

@@ -96,10 +96,6 @@ createTypeCanonicalizedMemRefOperands(OpBuilder &b, Location loc,
 
 LogicalResult mlir::linalg::LinalgOpToLibraryCallRewrite::matchAndRewrite(
     LinalgOp op, PatternRewriter &rewriter) const {
-  // Only LinalgOp for which there is no specialized pattern go through this.
-  if (isa<CopyOp>(op))
-    return failure();
-
   auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
   if (!libraryCallName)
     return failure();
@@ -113,65 +109,13 @@ LogicalResult mlir::linalg::LinalgOpToLibraryCallRewrite::matchAndRewrite(
   return success();
 }
 
-LogicalResult mlir::linalg::CopyOpToLibraryCallRewrite::matchAndRewrite(
-    CopyOp op, PatternRewriter &rewriter) const {
-  auto inputPerm = op.inputPermutation();
-  if (inputPerm.hasValue() && !inputPerm->isIdentity())
-    return failure();
-  auto outputPerm = op.outputPermutation();
-  if (outputPerm.hasValue() && !outputPerm->isIdentity())
-    return failure();
-
-  auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
-  if (!libraryCallName)
-    return failure();
-
-  rewriter.replaceOpWithNewOp<mlir::CallOp>(
-      op, libraryCallName.getValue(), TypeRange(),
-      createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(),
-                                            op.getOperands()));
-  return success();
-}
-
-LogicalResult mlir::linalg::CopyTransposeRewrite::matchAndRewrite(
-    CopyOp op, PatternRewriter &rewriter) const {
-  Value in = op.input(), out = op.output();
-
-  // If either inputPerm or outputPerm are non-identities, insert transposes.
-  auto inputPerm = op.inputPermutation();
-  if (inputPerm.hasValue() && !inputPerm->isIdentity())
-    in = rewriter.create<memref::TransposeOp>(op.getLoc(), in,
-                                              AffineMapAttr::get(*inputPerm));
-  auto outputPerm = op.outputPermutation();
-  if (outputPerm.hasValue() && !outputPerm->isIdentity())
-    out = rewriter.create<memref::TransposeOp>(op.getLoc(), out,
-                                               AffineMapAttr::get(*outputPerm));
-
-  // If nothing was transposed, fail and let the conversion kick in.
-  if (in == op.input() && out == op.output())
-    return failure();
-
-  auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
-  if (!libraryCallName)
-    return failure();
-
-  rewriter.replaceOpWithNewOp<mlir::CallOp>(
-      op, libraryCallName.getValue(), TypeRange(),
-      createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(), {in, out}));
-  return success();
-}
 
 /// Populate the given list with patterns that convert from Linalg to Standard.
 void mlir::linalg::populateLinalgToStandardConversionPatterns(
     RewritePatternSet &patterns) {
   // TODO: ConvOp conversion needs to export a descriptor with relevant
   // attribute values such as kernel striding and dilation.
-  // clang-format off
-  patterns.add<
-      CopyOpToLibraryCallRewrite,
-      CopyTransposeRewrite,
-      LinalgOpToLibraryCallRewrite>(patterns.getContext());
-  // clang-format on
+  patterns.add<LinalgOpToLibraryCallRewrite>(patterns.getContext());
 }
 
 namespace {

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

@@ -914,9 +914,10 @@ struct MemRefCopyOpLowering : public ConvertOpToLLVMPattern<memref::CopyOp> {
     auto sourcePtr = promote(unrankedSource);
     auto targetPtr = promote(unrankedTarget);
 
+    unsigned bitwidth = mlir::DataLayout::closest(op).getTypeSizeInBits(
+        srcType.getElementType());
     auto elemSize = rewriter.create<LLVM::ConstantOp>(
-        loc, getIndexType(),
-        rewriter.getIndexAttr(srcType.getElementTypeBitWidth() / 8));
+        loc, getIndexType(), rewriter.getIndexAttr(bitwidth / 8));
     auto copyFn = LLVM::lookupOrCreateMemRefCopyFn(
         op->getParentOfType<ModuleOp>(), getIndexType(), sourcePtr.getType());
     rewriter.create<LLVM::CallOp>(loc, copyFn,