[mlir] Use SCF for loops in ProgressiveVectorToSCF

[springerm]

Summary

Use SCF for loops instead of Affine for loops.

Depends On D101007

[springerm]

(no change)

[springerm]

(no change)

[springerm]

rebase

[springerm]

(no change)

[springerm]

(no change)

[nicolasvasilache]

Does this need some test update?

[springerm]

(rebase)

[springerm]

In D101013#2711177, @nicolasvasilache wrote:

Does this need some test update?

There are currently no tests that check for generated ops. Only integration tests. So no need to update tests.

mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp

generateInBoundsCheck(
       [&](OpBuilder &builder, Location loc) {
         if (outOfBoundsCase)
           outOfBoundsCase(builder, loc);
         return Value();
       });
 }
 
 /// Given an ArrayAttr, return a copy where the first element is dropped.
-static ArrayAttr dropFirstElem(PatternRewriter &rewriter, ArrayAttr attr) {
+static ArrayAttr dropFirstElem(OpBuilder &builder, ArrayAttr attr) {
   if (!attr)
     return attr;
-  return ArrayAttr::get(rewriter.getContext(), attr.getValue().drop_front());
+  return ArrayAttr::get(builder.getContext(), attr.getValue().drop_front());
 }
 
 /// Codegen strategy, depending on the operation.
 template <typename OpTy>
 struct Strategy;
 
 /// Code strategy for vector TransferReadOp.
 template <>
struct Strategy<TransferReadOp> {
   ///   %vec = vector.transfer_read %A[%a+%i, %b+%j, %c], %cst
   ///       : memref<?x?x?xf32>, vector<3xf32>
   ///   memref.store %vec, %casted[%i, %j] : memref<5x4xvector<3xf32>>
   /// }
   /// ```
   ///
   /// Note: The loop and type cast are generated in TransferOpConversion.
   ///       The original TransferReadOp and store op are deleted in `cleanup`.
-  static void rewriteOp(PatternRewriter &rewriter, TransferReadOp xferOp,
-                        Value buffer, Value iv) {
+  static void rewriteOp(OpBuilder &builder, TransferReadOp xferOp, Value buffer,
+                        Value iv) {
     SmallVector<Value, 8> storeIndices;
     getStoreIndices(xferOp, storeIndices);
     storeIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto bufferType = buffer.getType().dyn_cast<ShapedType>();
     auto vecType = bufferType.getElementType().dyn_cast<VectorType>();
-    auto inBoundsAttr = dropFirstElem(rewriter, xferOp.in_boundsAttr());
+    auto inBoundsAttr = dropFirstElem(builder, xferOp.in_boundsAttr());
     auto newXfer =
         vector_transfer_read(
             vecType, xferOp.source(), xferIndices,
-            AffineMapAttr::get(unpackedPermutationMap(xferOp, rewriter)),
+            AffineMapAttr::get(unpackedPermutationMap(xferOp, builder)),
             xferOp.padding(), Value(), inBoundsAttr)
             .value;
 
     if (vecType.getRank() > kTargetRank)
-      newXfer.getDefiningOp()->setAttr(kPassLabel, rewriter.getUnitAttr());
+      newXfer.getDefiningOp()->setAttr(kPassLabel, builder.getUnitAttr());
 
     memref_store(newXfer, buffer, storeIndices);
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension: Write
   /// padding value to the temporary buffer.
-  static void handleOutOfBoundsDim(PatternRewriter &rewriter,
+  static void handleOutOfBoundsDim(OpBuilder & /*builder*/,
                                    TransferReadOp xferOp, Value buffer,
                                    Value iv) {
     SmallVector<Value, 8> storeIndices;
     getStoreIndices(xferOp, storeIndices);
     storeIndices.push_back(iv);
 
     auto bufferType = buffer.getType().dyn_cast<ShapedType>();
     auto vecType = bufferType.getElementType().dyn_cast<VectorType>();
struct Strategy<TransferWriteOp> {
   /// accesses on the to-be-unpacked dimension.
   ///
   /// 1. Load an (N-1)-d vector from the (already `vector.type_cast`ed) buffer,
   ///    using the loop iteration variable `iv`.
   /// 2. Generate a new (N-1)-d TransferWriteOp, writing the loaded vector back
   ///    to memory.
   ///
   /// Note: For more details, see comments on Strategy<TransferReadOp>.
-  static void rewriteOp(PatternRewriter &rewriter, TransferWriteOp xferOp,
+  static void rewriteOp(OpBuilder &builder, TransferWriteOp xferOp,
                         Value buffer, Value iv) {
     SmallVector<Value, 8> loadIndices;
     getLoadIndices(xferOp, loadIndices);
     loadIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto vec = memref_load(buffer, loadIndices);
     auto vecType = vec.value.getType().dyn_cast<VectorType>();
-    auto inBoundsAttr = dropFirstElem(rewriter, xferOp.in_boundsAttr());
+    auto inBoundsAttr = dropFirstElem(builder, xferOp.in_boundsAttr());
     auto newXfer = vector_transfer_write(
         Type(), vec, xferOp.source(), xferIndices,
-        AffineMapAttr::get(unpackedPermutationMap(xferOp, rewriter)), Value(),
+        AffineMapAttr::get(unpackedPermutationMap(xferOp, builder)), Value(),
         inBoundsAttr);
 
     if (vecType.getRank() > kTargetRank)
-      newXfer.op->setAttr(kPassLabel, rewriter.getUnitAttr());
+      newXfer.op->setAttr(kPassLabel, builder.getUnitAttr());
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension.
-  static void handleOutOfBoundsDim(PatternRewriter &rewriter,
-                                   TransferWriteOp xferOp, Value buffer,
-                                   Value iv) {}
+  static void handleOutOfBoundsDim(OpBuilder &builder, TransferWriteOp xferOp,
+                                   Value buffer, Value iv) {}
 
   /// Cleanup after rewriting the op.
   static void cleanup(PatternRewriter &rewriter, TransferWriteOp xferOp) {
     rewriter.eraseOp(xferOp);
   }
 };
 
 template <typename OpTy>
LogicalResult matchAndRewrite(OpTy xferOp,
 
     ScopedContext scope(rewriter, xferOp.getLoc());
     // How the buffer can be found depends on OpTy.
     auto buffer = Strategy<OpTy>::getBuffer(xferOp);
     auto bufferType = buffer.getType().template dyn_cast<MemRefType>();
     auto castedType = unpackOneDim(bufferType);
     auto casted = vector_type_cast(castedType, buffer);
 
+    // Loop bounds and step.
     auto lb = std_constant_index(0).value;
     auto ub =
         std_constant_index(castedType.getDimSize(castedType.getRank() - 1))
             .value;
-    affineLoopBuilder(lb, ub, 1, [&](Value iv) {
+    auto step = std_constant_index(1).value;
+
+    // Generate for loop.
+    rewriter.create<scf::ForOp>(
+        xferOp.getLoc(), lb, ub, step, ValueRange(),
+        [&](OpBuilder &b, Location loc, Value iv, ValueRange /*loopState*/) {
+          ScopedContext scope(b, loc);
           generateInBoundsCheck(
-          xferOp, iv, rewriter, unpackedDim(xferOp),
+              xferOp, iv, b, unpackedDim(xferOp),
               /*inBoundsCase=*/
-          [&](OpBuilder & /*b*/, Location loc) {
-            Strategy<OpTy>::rewriteOp(rewriter, xferOp, casted, iv);
+              [&](OpBuilder &b, Location /*loc*/) {
+                Strategy<OpTy>::rewriteOp(b, xferOp, casted, iv);
               },
               /*outOfBoundsCase=*/
-          [&](OpBuilder & /*b*/, Location loc) {
-            Strategy<OpTy>::handleOutOfBoundsDim(rewriter, xferOp, casted, iv);
+              [&](OpBuilder &b, Location /*loc*/) {
+                Strategy<OpTy>::handleOutOfBoundsDim(b, xferOp, casted, iv);
               });
+          b.create<scf::YieldOp>(loc);
         });
 
     Strategy<OpTy>::cleanup(rewriter, xferOp);
     return success();
   }
 };
 
 /// Compute the indices into the memref for the LoadOp/StoreOp generated as
 /// part of Strided1dTransferOpConversion. Return the memref dimension on which