[mlir][Linalg] Add canonicalization to remove no-op linalg operations.

[mravishankar]

Summary

linalg.generic/indexed_generic operations on tensors whose body is
just yielding the (non-induction variable) arguments of the operation
can be canonicalized by replacing uses of the result with the
corresponding arguments.

[hanchung]

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

2148

auto

2148

I think we can use body.getTerminator() here. It will return an Operation *, and it's clearer to me.

2157

auto

[silvas]

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

2127

nit: would call this RemoveIdentityLinalgOps.

2133

Why restrict to only GenericOp/IndexedGenericOp? The justification in DeduplicateInputs::matchAndRewrite doesn't seem to apply.

2163

should this more properly be genericOp->getInput(argumentNumber - numIndexArgs)?

[mravishankar]

Address comments

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

2133

If a named op is a no-op, then we shouldnt be defining such named ops to begin with right? The Generic/IndexedGeneric ops are free form and after fusion, canonicalization, etc can become a no-op....

2163

There is an interesting corner case where you are returning one of the outs operands without any modification. That is still valid to do I think. So just using operand covers both these cases.

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

for (int i = 0, e = linalgOp.getNumInputs(); i < e; i++) {
       payload.eraseArgument(bbArgBaseOffset + reversed);
     }
 
     rewriter.replaceOp(op, newOp->getResults());
     return success();
   }
 };
 
+/// Remove generic/indexed_generic operations (on tensors) that are just copying
+/// the values from inputs to the results. Requirements are
+/// 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 RewritePattern {
+  RemoveIdentityLinalgOps(PatternBenefit benefit = 1)
+      : RewritePattern(benefit, MatchAnyOpTypeTag()) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    if (!isa<GenericOp, IndexedGenericOp>(op))
+      return failure();
+    LinalgOp genericOp = cast<LinalgOp>(op);
+    if (!genericOp.hasTensorSemantics())
+      return failure();
+    // Check all indexing maps are identity.
+    if (llvm::any_of(genericOp.getIndexingMaps(),
+                     [](AffineMap map) { return !map.isIdentity(); }))
+      return failure();
+
+    // Check that the body of the linalg operation is just a linalg.yield
+    // operation.
+    Block &body = op->getRegion(0).front();
+    if (!llvm::hasSingleElement(body))
+      return failure();
+    auto yieldOp = dyn_cast<linalg::YieldOp>(body.getTerminator());
+    if (!yieldOp)
+      return failure();
+
+    // Get the argument number of the returned values. That is the operand
+    // number to use for replacing uses of this operation.
+    unsigned numIndexArgs = genericOp.getNumPayloadInductionVariables();
+    SmallVector<Value, 4> returnedArgs;
+    for (Value yieldVal : yieldOp.values()) {
+      auto yieldArg = yieldVal.dyn_cast<BlockArgument>();
+      if (!yieldArg)
+        return failure();
+      unsigned argumentNumber = yieldArg.getArgNumber();
+      if (argumentNumber < numIndexArgs)
+        return failure();
+      returnedArgs.push_back(op->getOperand(argumentNumber - numIndexArgs));
+    }
+    rewriter.replaceOp(genericOp, returnedArgs);
+    return success();
+  }
+};
+
 /// Canonicalize a `linalgOp` -> `dim` pattern by replacing the `dim` arg
 /// with the corresponding output tensor argument of the linalg op.
 struct ReplaceDimOfLinalgResult : public OpRewritePattern<DimOp> {
   using OpRewritePattern<DimOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(DimOp dimOp,
                                 PatternRewriter &rewriter) const override {
     Value dimOpArg = dimOp.memrefOrTensor();
     auto linalgOp = dimOpArg.getDefiningOp<LinalgOp>();
     if (!linalgOp)
       return failure();
 
     auto results = linalgOp.getOperation()->getResults();
     int64_t id = std::distance(results.begin(), llvm::find(results, dimOpArg));
     auto outputTensors = linalgOp.getOutputTensors();
     rewriter.replaceOpWithNewOp<DimOp>(dimOp, outputTensors[id], dimOp.index());
     return success();
   }
 };
 } // namespace
 
 #define CANONICALIZERS_AND_FOLDERS(XXX)                                        \
   void XXX::getCanonicalizationPatterns(OwningRewritePatternList &results,     \
                                         MLIRContext *context) {                \
-    results.insert<DeduplicateInputs, EraseDeadLinalgOp, FoldTensorCastOp>();  \
+    results.insert<DeduplicateInputs, EraseDeadLinalgOp, FoldTensorCastOp,     \
+                   RemoveIdentityLinalgOps>();                                 \
     results.insert<ReplaceDimOfLinalgResult>(context);                         \
   }                                                                            \
                                                                                \
   LogicalResult XXX::fold(ArrayRef<Attribute>,                                 \
                           SmallVectorImpl<OpFoldResult> &) {                   \
     return foldMemRefCast(*this);                                              \
   }
 

mlir/test/Dialect/Linalg/canonicalize.mlir

func @dce_zero_memref(%arg0 : memref<0xf32>, %arg1: tensor<0xf32>) -> tensor<0xf32> {
   %1 = linalg.generic #trait outs(%arg1 : tensor<0xf32>) {
   ^bb(%0: f32) :
     linalg.yield %0 : f32
   } -> tensor<0xf32>
 
   return %1: tensor<0xf32>
 }
 // CHECK-LABEL: @dce_zero_memref
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<0xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<0xf32>
 //   CHECK-NOT:   linalg.copy
-//  CHECK-NEXT:   linalg.generic
+//  CHECK-NEXT:   return %[[ARG1]]
 
 // -----
 
 func @reshape_splat_constant_int32() -> tensor<2x4x2xi32>
 {
   %c0 = constant dense<42> : tensor<2x8xi32>
   %0 = linalg.tensor_reshape %c0
          [affine_map<(d0, d1, d2) -> (d0)>,
func @init_tensor_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
   return %1 : tensor<6x5x?xf32>
 }
 //      CHECK: func @init_tensor_reshape_collapse
 // CHECK-SAME:   %[[ARG0:.+]]: index
 //      CHECK:   %[[C28:.+]] = constant 28 : index
 //      CHECK:   %[[T0:.+]] = muli %[[ARG0]], %[[C28]]
 //      CHECK:   %[[T1:.+]] = linalg.init_tensor [6, 5, %[[T0]]]
 //      CHECK:   return %[[T1]]
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @remove_no_op(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>)
+  -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %0 = dim %arg0, %c0 : tensor<?x?x?xf32>
+  %1 = dim %arg0, %c1 : tensor<?x?x?xf32>
+  %2 = dim %arg0, %c2 : tensor<?x?x?xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4, %5 = linalg.generic {
+    indexing_maps = [#map, #map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+    outs(%3, %3 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32, %arg5 : f32):
+    linalg.yield %arg3, %arg2 : f32, f32
+  } -> tensor<?x?x?xf32>, tensor<?x?x?xf32>
+  return %4, %5 : tensor<?x?x?xf32>, tensor<?x?x?xf32>
+}
+// CHECK-LABEL: func @remove_no_op
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//       CHECK:     return %[[ARG1]], %[[ARG0]]