[mlir][SCF][bufferize][NFC] Bufferize scf.for terminator separately

[springerm]

Summary

This allows for better type inference during bufferization and is in preparation of supporting memory spaces.

Depends On D128581

[nicolasvasilache]

Can we split into 1PR per op plz, also drop the NFC as tests are changing.

[springerm]

update

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

for (int64_t idx = 0;
       yieldValues.push_back(alloc);
     }
 
     rewriter.updateRootInPlace(
         yieldOp, [&]() { yieldOp.getResultsMutable().assign(yieldValues); });
     return success();
   }
 
+  BaseMemRefType getBufferType(Operation *op, BlockArgument bbArg,
+                               const BufferizationOptions &options) const {
+    auto forOp = cast<scf::ForOp>(op);
+    return bufferization::getBufferType(
+        forOp.getOpOperandForRegionIterArg(bbArg).get(), options);
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto forOp = cast<scf::ForOp>(op);
     Block *oldLoopBody = &forOp.getLoopBody().front();
 
     // Indices of all iter_args that have tensor type. These are the ones that
     // are bufferized.
     DenseSet<int64_t> indices = getTensorIndices(forOp.getInitArgs());
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Set up new iter_args. The loop body uses tensors, so wrap the (memref)
     // iter_args of the new loop in ToTensorOps.
     rewriter.setInsertionPointToStart(loopBody);
     SmallVector<Value> iterArgs =
         getBbArgReplacements(rewriter, newForOp.getRegionIterArgs(), indices);
     iterArgs.insert(iterArgs.begin(), newForOp.getInductionVar());
 
-    // Erase terminator if present.
-    if (iterArgs.size() == 1)
-      rewriter.eraseOp(loopBody->getTerminator());
-
     // Move loop body to new loop.
     rewriter.mergeBlocks(oldLoopBody, loopBody, iterArgs);
 
-    // Update scf.yield of new loop.
-    auto yieldOp = cast<scf::YieldOp>(loopBody->getTerminator());
-    rewriter.setInsertionPoint(yieldOp);
-    SmallVector<Value> yieldValues = getYieldedValues(
-        rewriter, yieldOp.getResults(), initArgsTypes, indices, options);
-    yieldOp.getResultsMutable().assign(yieldValues);
-
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newForOp->getResults());
 
     return success();
   }
 
   /// Assert that yielded values of an scf.for op are equivalent to their
   /// corresponding bbArgs. In that case, the buffer relations of the
struct YieldOpInterface
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto yieldOp = cast<scf::YieldOp>(op);
     if (!isa<scf::ExecuteRegionOp, scf::IfOp, scf::ForOp, scf::WhileOp>(
             yieldOp->getParentOp()))
       return yieldOp->emitError("unsupported scf::YieldOp parent");
 
-    // TODO: Bufferize scf.yield inside scf.while/scf.for here.
-    // (Currently bufferized together with scf.while/scf.for.)
-    if (isa<scf::ForOp, scf::WhileOp>(yieldOp->getParentOp()))
+    // TODO: Bufferize scf.yield inside scf.while here. (Currently bufferized
+    // together with scf.while.)
+    if (isa<scf::WhileOp>(yieldOp->getParentOp()))
       return success();
 
     SmallVector<Value> newResults;
     for (const auto &it : llvm::enumerate(yieldOp.getResults())) {
       Value value = it.value();
       if (value.getType().isa<TensorType>()) {
         Value buffer = getBuffer(rewriter, value, options);
+        if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp())) {
+          BaseMemRefType resultType =
+              cast<BufferizableOpInterface>(forOp.getOperation())
+                  .getBufferType(forOp.getRegionIterArgs()[it.index()],
+                                 options);
+          buffer = castBuffer(rewriter, buffer, resultType);
+        }
         newResults.push_back(buffer);
       } else {
         newResults.push_back(value);
       }
     }
 
     replaceOpWithNewBufferizedOp<scf::YieldOp>(rewriter, op, newResults);
     return success();