[mlir][linalg] Use explicit replace in canonicalization pattern (NFC).

[gysit]

Summary

Introduce an explicit replaceOp call to enable the tracking of the producer LinalgOp.

[nicolasvasilache]

Thanks for fixing this!

[rriddle]

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

1649

Why is the SmallVector necessary here? Can you not just do: rewriter.replaceOp(linalgOp, newOp->getResults())

[gysit]

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

1649

One of the replacement values needs to be set to castBack (see next line). The castBack value is needed if the result has more uses than just castOp.

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

LogicalResult matchAndRewrite(tensor::CastOp castOp,
     outputOperands[resultNumber] = newOperand;
     newOperands.append(outputOperands.begin(), outputOperands.end());
 
     SmallVector<Type> resultTypes(linalgOp->result_type_begin(),
                                   linalgOp->result_type_end());
     resultTypes[resultNumber] = resultType;
     Operation *newOp = linalgOp.clone(rewriter, loc, resultTypes, newOperands);
 
-    if (!resultValue.hasOneUse()) {
-      SmallVector<Value> results(newOp->result_begin(), newOp->result_end());
-      // Create a tensor.cast operation back to the original type.
-      Value castBack = rewriter.create<tensor::CastOp>(
-          loc, resultValue.getType(), newOp->getResult(resultNumber));
-      results[resultNumber] = castBack;
-      // Replace all uses except the use in the cast op that is matched by the
-      // pattern. Note that this cast is from a more static shape to a more
-      // dynamic shape. These are expected to be pulled into their consumers.
-      rewriter.replaceOpWithIf(linalgOp, results,
-                               [&castOp](OpOperand &use) -> bool {
-                                 return use.getOwner() != castOp.getOperation();
-                               });
-    }
+
+    SmallVector<Value> results(newOp->result_begin(), newOp->result_end());
+    rewriter.replaceOp(linalgOp, results);
     rewriter.replaceOp(castOp, newOp->getResult(resultNumber));
     return success();
   }
 };
 
 /// For each of the operand in `operands` this function maps the static sizes of
 /// dimensions to their affine dim expressions.
 static void populateMap(LinalgOp linalgOp, ArrayRef<OpOperand *> operands,