[mlir][linalg] Cleanup LinalgOp usage in test passes.

[gysit]

Summary

Replace the uses of deprecated Structured Op Interface methods in TestLinalgElementwiseFusion.cpp, TestLinalgFusionTransforms.cpp, and Transforms.cpp. The patch is based on https://reviews.llvm.org/D103394.

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

 static LogicalResult rewriteAsPaddedOp(PatternRewriter &rewriter,
                                        TiledLinalgOp &res,
                                        const LinalgTilingOptions &options) {
   LinalgOp opToPad = res.op;
   Location loc = opToPad->getLoc();
 
   // If the op is fully static, it does not need padding.
   // TODO: there are cases where we may still want to pad to larger sizes.
-  if (llvm::all_of(opToPad.getShapedOperands(), [](Value v) {
-        return v.getType().cast<RankedTensorType>().hasStaticShape();
-      }))
+  assert(opToPad.hasTensorSemantics() &&
+         "expected operation to have tensor semantics");
+  if (!opToPad.hasDynamicShape())
     return success();
 
   OpBuilder::InsertionGuard g(rewriter);
   // Set IP after op because we also take the dims of the original output.
   rewriter.setInsertionPointAfter(opToPad);
   // Make a copy of the shaped operands and update it.
   SmallVector<Value> newOperands;
-  newOperands.reserve(opToPad.getNumShapedOperands());
-  for (OpOperand &operand : opToPad.getShapedOpOperands()) {
+  newOperands.reserve(opToPad.getNumInputsAndOutputs());
+  for (OpOperand *opOperand : opToPad.getInputAndOutputOperands()) {
     Value paddedOperand;
     // If padding was requested but the shape cannot be bounded statically then
     // the pattern fails to apply.
-    if (failed(padOperandToSmallestStaticBoundingBox(rewriter, opToPad, operand,
-                                                     options, paddedOperand))) {
+    if (failed(padOperandToSmallestStaticBoundingBox(
+            rewriter, opToPad, *opOperand, options, paddedOperand))) {
       return failure();
     }
-    newOperands.push_back(paddedOperand ? paddedOperand : operand.get());
+    newOperands.push_back(paddedOperand ? paddedOperand : opOperand->get());
   }
 
   // Clone `opToPad` to operate on the statically padded shapes.
   auto resultTensorTypes =
       ValueRange(newOperands).take_back(opToPad.getNumOutputs()).getTypes();
   ValueRange otherOperands = opToPad.getAssumedNonShapedOperands();
   newOperands.append(otherOperands.begin(), otherOperands.end());
   linalg::LinalgOp paddedOp =

mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp

 
 namespace mlir {
 
 static void addOperands(Operation *op, SetVector<Value> &operandSet) {
   if (!op)
     return;
   TypeSwitch<Operation *, void>(op)
       .Case<linalg::LinalgOp>([&](linalg::LinalgOp linalgOp) {
-        operandSet.insert(linalgOp.getInputs().begin(),
-                          linalgOp.getInputs().end());
+        SmallVector<Value> inputOperands = linalgOp.getInputOperands();
+        operandSet.insert(inputOperands.begin(), inputOperands.end());
       })
       .Default([&](Operation *operation) {
         operandSet.insert(operation->operand_begin(), operation->operand_end());
       });
 }
 
 template <int limit = 3>
 static bool setFusedOpOperandLimit(const OpResult &producer,

mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp

f.walk([&](LinalgOp op) {
     // TODO: support multi-results.
     if (op->getNumResults() <= 1)
       linalgOps.push_back(op);
   });
 
   // Tile and Fuse for tensors inputs (TODO: all tensor operands).
   bool changed = false;
   for (LinalgOp linalgOp : llvm::reverse(linalgOps)) {
-    for (OpOperand &opOperand : linalgOp.getShapedOpOperands()) {
-      if (opOperand.get().getType().isa<MemRefType>()) {
+    for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
+      if (opOperand->get().getType().isa<MemRefType>()) {
         // TODO: LinalgDependenceGraph should be able to update itself.
         // The current naive and expensive reconstruction of the graph should be
         // removed.
         linalg::Aliases aliases;
         linalg::LinalgDependenceGraph graph(aliases, linalgOps);
-        if (auto info = fuseProducerOfBuffer(b, opOperand, graph)) {
+        if (auto info = fuseProducerOfBuffer(b, *opOperand, graph)) {
           auto *originalOp = info->originalProducer.getOperation();
           eraseSet.insert(originalOp);
           auto *originalOpInLinalgOpsVector =
               std::find(linalgOps.begin(), linalgOps.end(), originalOp);
           *originalOpInLinalgOpsVector = info->fusedProducer.getOperation();
           changed = true;
         }
       } else {
-        assert(opOperand.get().getType().isa<RankedTensorType>());
+        assert(opOperand->get().getType().isa<RankedTensorType>());
         // Tile and Fuse tensor input.
-        if (opOperand.getOperandNumber() >= linalgOp.getNumInputs())
+        if (opOperand->getOperandNumber() >= linalgOp.getNumInputs())
           continue;
-        if (auto info = fuseProducerOfTensor(b, opOperand)) {
+        if (auto info = fuseProducerOfTensor(b, *opOperand)) {
           auto *originalOp = info->originalProducer.getOperation();
           auto *originalOpInLinalgOpsVector =
               std::find(linalgOps.begin(), linalgOps.end(), originalOp);
           *originalOpInLinalgOpsVector = info->fusedProducer.getOperation();
           // Don't mark for erasure in the tensor case, let DCE handle this.
           changed = true;
         }
       }