[mlir][linalg] Introduce a separate EraseIdentityCopyOp Pattern.

[gysit]

Summary

Split out an EraseIdentityCopyOp from the existing RemoveIdentityLinalgOps pattern. Introduce an additional check to ensure the pattern checks the permutation maps match. This is a preparation step to specialize RemoveIdentityLinalgOps to GenericOp only.

[gysit]

Rebase.

[rriddle]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

438

An assert like this in a canonicalization pattern feels pretty sketchy. Canonicalization patterns can always be applied to an op, so when is this not true?

[gysit]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

438

The input and output operands are limited to memrefs in the op definition and the assertion is always true.

It is here since the canonicalization only makes sense in buffer land and it would be triggered if the op is modified to work on tensors (as most other structured operations). It thus encodes a necessary precondition to make the canonicalization correct.

[mehdi_amini]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

438

If I understand correctly, you're saying that this assert could be inserted in the verifier for CopyOp as well?

[gysit]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

438

The verification already checks the input and output operands are memrefs (they are of type AnyStridedMemRef). I think we are good there.

The purpose of the assert is to mark that this particular piece of code breaks once the operation supports tensor semantics. I personally can live with removing it but I would say it is a safety net and it also helps to understand the code. Comparing an output tensor otherwise may seem strange otherwise...

[mehdi_amini]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

438

My question was to validate my understanding: "in the current state of this op, this check could be added to the verifier and nothing would break"?

You're saying "The verification already checks the input and output operands are memrefs", but it isn't obvious to a reader (like River here) that this implies the same thing as this assert, hence my clarifying question.

[gysit]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

438

Ah, yes the hasBufferSemantics method checks that all operands are memrefs or scalars.

Sorry too deep in the topic :).

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

def CopyOp : LinalgStructured_Op<"copy", [CopyOpInterface]> {
   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
     AnyTypeOf<[AnyComplex, AnyFloat, AnySignlessInteger, AnyVector]>:$value,
     AnyShaped:$output);

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

void CopyOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), input(),
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Write::get(), output(),
                        SideEffects::DefaultResource::get());
 }
 
+namespace {
+/// Remove copy operations that copy data inplace. Requirements are:
+/// 1) The input and output values are identical.
+/// 2) The input and output permutation maps are identical.
+struct EraseIdentityCopyOp : public OpRewritePattern<CopyOp> {
+  using OpRewritePattern<CopyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CopyOp copyOp,
+                                PatternRewriter &rewriter) const override {
+    assert(copyOp.hasBufferSemantics());
+    if (copyOp.input() == copyOp.output() &&
+        copyOp.inputPermutation() == copyOp.outputPermutation()) {
+      rewriter.eraseOp(copyOp);
+      return success();
+    }
+    return failure();
+  }
+};
+} // namespace
+
+void CopyOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *context) {
+  results.add<EraseIdentityCopyOp>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // FillOp
 //===----------------------------------------------------------------------===//
 void FillOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block) {
   assert(block.getNumArguments() == 2 && "FillOp regionBuilder expects 2 args");
   b.create<linalg::YieldOp>(block.getArgument(0));
 }
 

 /// 1) All iterator types are parallel
 /// 2) The body contains just a yield operation with the yielded values being
 ///    the arguments corresponding to the operands.
 struct RemoveIdentityLinalgOps : public OpInterfaceRewritePattern<LinalgOp> {
   using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
 
   LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override {
-    if (auto copyOp = dyn_cast<CopyOp>(*op)) {
-      assert(copyOp.hasBufferSemantics());
-      if (copyOp.input() == copyOp.output() &&
-          copyOp.inputPermutation() == copyOp.outputPermutation()) {
-        rewriter.eraseOp(op);
-        return success();
-      }
-    }
-
     if (!isa<GenericOp>(op))
       return failure();
     if (!op.hasTensorSemantics())
       return failure();
     // Check all indexing maps are identity.
     if (llvm::any_of(op.getIndexingMaps(),
                      [](AffineMap map) { return !map.isIdentity(); }))
       return failure();

mlir/test/Dialect/Linalg/canonicalize.mlir

//   CHECK-NOT: linalg.copy
   linalg.copy(%arg0, %arg0): memref<2x3x?x4xf32>, memref<2x3x?x4xf32>
 
 //   CHECK: return
   return
 }
 
 // -----
 
+// CHECK-LABEL: @self_copy_with_permutation
+func @self_copy_with_permutation(%arg0 : memref<2x3x?x4xf32>) {
+
+//   CHECK: linalg.copy
+  linalg.copy(%arg0, %arg0)
+    {inputPermutation = affine_map<(i, j, k, l) -> (j, k, i, l)>,
+     outputPermuation = affine_map<(i, j, k, l) -> (i, j, k, l)>} : memref<2x3x?x4xf32>, memref<2x3x?x4xf32>
+
+//   CHECK: return
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func @fold_fill_reshape()
 func @fold_fill_reshape() -> tensor<6x4xf32> {
   %zero = constant 0.0 : f32
   // CHECK: %[[INIT:.+]] = linalg.init_tensor [6, 4] : tensor<6x4xf32>
   %init = linalg.init_tensor [1, 2, 3, 4] : tensor<1x2x3x4xf32>
   // CHECK: %[[FILL:.+]] = linalg.fill(%cst, %[[INIT]]) : f32, tensor<6x4xf32> -> tensor<6x4xf32>
   %fill = linalg.fill(%zero, %init) : f32, tensor<1x2x3x4xf32> -> tensor<1x2x3x4xf32>
   %reshape = linalg.tensor_collapse_shape %fill [[0, 1, 2], [3]]