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

This reverts commit 016956b.
Reverting it to fix NVidia build without being in a hurry.

Author: pifon2a

mlir/docs/Dialects/Linalg/_index.md

@@ -545,6 +545,7 @@ 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,6 +39,25 @@ 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>):
-        "memref.copy"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
+        "linalg.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>
-        memref.copy(%arg1, %0) : memref<2xf32>, memref<2xf32>
+        linalg.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>
-        memref.copy(%1, %2) : memref<2xf32>, memref<2xf32>
+        linalg.copy(%1, %2) : memref<2xf32>, memref<2xf32>
         dealloc %1 : memref<2xf32>
         br ^bb3(%2 : memref<2xf32>)
       ^bb3(%3: memref<2xf32>):  // 2 preds: ^bb1, ^bb2
-        memref.copy(%3, %arg2) : memref<2xf32>, memref<2xf32>
+        linalg.copy(%3, %arg2) : memref<2xf32>, memref<2xf32>
         dealloc %3 : memref<2xf32>
         return
       }

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

@@ -16,6 +16,7 @@
 
 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"
 
@@ -56,6 +57,119 @@ class LinalgStructured_Op<string mnemonic, list<OpTrait> 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-memref.copy", "bool",
+    Option<"useLinalgCopy", "use-linalg-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 memref.copy is used.
+  /// subview. If None then a linalg.copy is used.
   Optional<CopyCallbackFn> copyInFn = None;
   Optional<CopyCallbackFn> copyOutFn = None;
   LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
@@ -390,9 +390,6 @@ 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);
@@ -937,15 +934,6 @@ 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);
 
@@ -1218,7 +1206,7 @@ void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
 ///    %subView = subview %allocOrView ...
 ///    [optional] linalg.fill(%allocOrView, %cst) ...
 ///    ...
-///    memref.copy(%in, %subView) ...
+///    linalg.copy(%in, %subView) ...
 ///    vector.transfer_read %allocOrView[...], %cst ...
 /// ```
 /// into
@@ -1229,8 +1217,8 @@ void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
 ///    ...
 ///    vector.transfer_read %in[...], %cst ...
 /// ```
-/// 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
+/// 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
 /// linalg.fill is specified).
 /// This is a custom rewrite to forward partial reads (with optional fills) to
 /// vector.transfer_read.
@@ -1249,7 +1237,7 @@ struct LinalgCopyVTRForwardingPattern
 ///    %subView = subview %allocOrView...
 ///    ...
 ///    vector.transfer_write %..., %allocOrView[...]
-///    memref.copy(%subView, %out)
+///    linalg.copy(%subView, %out)
 /// ```
 /// into
 /// ```
@@ -1259,7 +1247,7 @@ struct LinalgCopyVTRForwardingPattern
 ///    ...
 ///    vector.transfer_write %..., %out[...]
 /// ```
-/// Where there is no interleaved use between transfer_write and memref.copy.
+/// Where there is no interleaved use between transfer_write and linalg.copy.
 /// This is a custom rewrite to forward partial writes to vector.transfer_write.
 struct LinalgCopyVTWForwardingPattern
     : public OpRewritePattern<vector::TransferWriteOp> {

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

@@ -96,6 +96,10 @@ 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();
@@ -109,12 +113,65 @@ 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.
-  patterns.add<LinalgOpToLibraryCallRewrite>(patterns.getContext());
+  // clang-format off
+  patterns.add<
+      CopyOpToLibraryCallRewrite,
+      CopyTransposeRewrite,
+      LinalgOpToLibraryCallRewrite>(patterns.getContext());
+  // clang-format on
 }
 
 namespace {