[mlir] Support tensor types in non-unrolled VectorToSCF

[springerm]

Summary

Support for tensor types in the unrolled version will follow in a separate commit.

Add a new pass option to activate lowering of transfer ops with tensor types (default: deactivated).

[springerm]

no change

[springerm]

no change

[springerm]

add tensor lowering pass option

[springerm]

rebase

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

if (!isBroadcast) {
     using edsc::op::operator+;
     indices[dim.getValue()] = adaptor.indices()[dim.getValue()] + iv;
   }
 }
 
 static void maybeYieldValue(bool hasRetVal, OpBuilder builder, Location loc,
                             Value value) {
   if (hasRetVal) {
+    assert(value && "Expected non-empty value");
     builder.create<scf::YieldOp>(loc, value);
   } else {
     builder.create<scf::YieldOp>(loc);
   }
 }
 
 /// Generates a boolean Value that is true if the iv-th bit in xferOp's mask
 /// is set to true. No such check is generated under following circumstances:

 /// rank.
 template <typename OpTy>
 static void maybeApplyPassLabel(OpBuilder &builder, OpTy newXferOp,
                                 unsigned targetRank) {
   if (newXferOp.getVectorType().getRank() > targetRank)
     newXferOp->setAttr(kPassLabel, builder.getUnitAttr());
 }
 
+/// Return true if this TensorWriteOp operates on a source tensor.
+static bool isTensorOp(TransferWriteOp xferOp) {
+  if (xferOp->getNumResults() > 0) {
+    assert(xferOp.getShapedType().isa<RankedTensorType>() &&
+           "Expected that TransferWriteOp with result has tensor source");
+    return true;
+  }
+  return false;
+}
+
 namespace lowering_n_d {
 
 /// Helper data structure for data and mask buffers.
 struct BufferAllocs {
   Value dataBuffer;
   Value maskBuffer;
 };
 
struct Strategy<TransferReadOp> {
   /// }
   /// ```
   ///
   /// Note: The loop and type cast are generated in TransferOpConversion.
   ///       The original TransferReadOp and store op are deleted in `cleanup`.
   /// Note: The `mask` operand is set in TransferOpConversion.
   static TransferReadOp rewriteOp(OpBuilder &builder,
                                   VectorTransferToSCFOptions options,
-                                  TransferReadOp xferOp, Value buffer,
-                                  Value iv) {
+                                  TransferReadOp xferOp, Value buffer, Value iv,
+                                  ValueRange /*loopState*/) {
     SmallVector<Value, 8> storeIndices;
     getBufferIndices(xferOp, storeIndices);
     storeIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto bufferType = buffer.getType().dyn_cast<ShapedType>();
maybeApplyPassLabel(builder,
                         options.targetRank);
 
     memref_store(newXfer, buffer, storeIndices);
     return newXfer.getDefiningOp<TransferReadOp>();
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension: Write
   /// padding value to the temporary buffer.
-  static void handleOutOfBoundsDim(OpBuilder & /*builder*/,
+  static Value handleOutOfBoundsDim(OpBuilder & /*builder*/,
                                     TransferReadOp xferOp, Value buffer,
-                                   Value iv) {
+                                    Value iv, ValueRange /*loopState*/) {
     SmallVector<Value, 8> storeIndices;
     getBufferIndices(xferOp, storeIndices);
     storeIndices.push_back(iv);
 
     auto bufferType = buffer.getType().dyn_cast<ShapedType>();
     auto vecType = bufferType.getElementType().dyn_cast<VectorType>();
     auto vec = std_splat(vecType, xferOp.padding());
     memref_store(vec, buffer, storeIndices);
+
+    return Value();
   }
 
   /// Cleanup after rewriting the op.
-  static void cleanup(PatternRewriter &rewriter, TransferReadOp xferOp) {
+  static void cleanup(PatternRewriter &rewriter, TransferReadOp xferOp,
+                      scf::ForOp /*forOp*/) {
     rewriter.eraseOp(getStoreOp(xferOp));
     rewriter.eraseOp(xferOp);
   }
+
+  /// Return the initial loop state for the generated scf.for loop.
+  static Value initialLoopState(TransferReadOp xferOp) { return Value(); }
 };
 
 /// Codegen strategy for vector TransferWriteOp.
 template <>
 struct Strategy<TransferWriteOp> {
   /// Find the temporary buffer allocation. All labeled TransferWriteOps are
   /// used like this, where %buf is either the buffer allocation or a type cast
   /// of the buffer allocation:
struct Strategy<TransferWriteOp> {
   ///    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 TransferWriteOp rewriteOp(OpBuilder &builder,
                                    VectorTransferToSCFOptions options,
                                    TransferWriteOp xferOp, Value buffer,
-                                   Value iv) {
+                                   Value iv, ValueRange loopState) {
     SmallVector<Value, 8> loadIndices;
     getBufferIndices(xferOp, loadIndices);
     loadIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto vec = memref_load(buffer, loadIndices);
     auto inBoundsAttr = dropFirstElem(builder, xferOp.in_boundsAttr());
+    auto source = loopState.empty() ? xferOp.source() : loopState[0];
+    Type type = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();
     auto newXfer = vector_transfer_write(
-        Type(), vec, xferOp.source(), xferIndices,
+        type, vec, source, xferIndices,
         AffineMapAttr::get(unpackedPermutationMap(xferOp, builder)), Value(),
         inBoundsAttr);
 
     maybeApplyPassLabel(builder, newXfer.op, options.targetRank);
 
     return newXfer;
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension.
-  static void handleOutOfBoundsDim(OpBuilder &builder, TransferWriteOp xferOp,
-                                   Value buffer, Value iv) {}
+  static Value handleOutOfBoundsDim(OpBuilder &builder, TransferWriteOp xferOp,
+                                    Value buffer, Value iv,
+                                    ValueRange loopState) {
+    return isTensorOp(xferOp) ? loopState[0] : Value();
+  }
 
   /// Cleanup after rewriting the op.
-  static void cleanup(PatternRewriter &rewriter, TransferWriteOp xferOp) {
+  static void cleanup(PatternRewriter &rewriter, TransferWriteOp xferOp,
+                      scf::ForOp forOp) {
+    if (isTensorOp(xferOp)) {
+      assert(forOp->getNumResults() == 1 && "Expected one for loop result");
+      rewriter.replaceOp(xferOp, forOp->getResult(0));
+    } else {
       rewriter.eraseOp(xferOp);
     }
+  }
+
+  /// Return the initial loop state for the generated scf.for loop.
+  static Value initialLoopState(TransferWriteOp xferOp) {
+    return isTensorOp(xferOp) ? xferOp.source() : Value();
+  }
 };
 
 template <typename OpTy>
 LogicalResult checkPrepareXferOp(OpTy xferOp, unsigned targetRank) {
   if (xferOp->hasAttr(kPassLabel))
     return failure();
   if (xferOp.getVectorType().getRank() <= targetRank)
     return failure();

 ///    ```
 ///    %casted = vector.type_cast %0
 ///        : memref<5xvector<4x3xf32>> to memref<5x4xvector<3xf32>>
 ///    ```
 /// 2. Generate a for loop and rewrite the transfer op according to the
 ///    corresponding Strategy<OpTy>. If the to-be-unpacked dimension can be
 ///    out-of-bounds, generate an if-check and handle both cases separately.
 /// 3. Clean up according to the corresponding Strategy<OpTy>.
+///
+/// Note: If the transfer op is a TransferWriteOp and operates on a tensor
+/// source (as opposed to a memref source), then each iteration of the generated
+/// scf.for loop yields the new tensor value. E.g.:
+/// ```
+/// %result = scf.for i = 0 to 5 {
+///   %0 = memref.load %buffer[i] : memref<5xvector<4x3xf32>>
+///   %1 = vector.transfer_write %0, %source[...]
+///       : vector<4x3xf32>, tensor<5x4x3xf32>
+///   scf.yield %1 : tensor<5x4x3xf32>
+/// }
+/// ```
 template <typename OpTy>
 struct TransferOpConversion : public VectorToSCFPattern<OpTy> {
   using VectorToSCFPattern<OpTy>::VectorToSCFPattern;
 
   LogicalResult matchAndRewrite(OpTy xferOp,
                                 PatternRewriter &rewriter) const override {
     if (!xferOp->hasAttr(kPassLabel))
       return failure();
LogicalResult matchAndRewrite(OpTy xferOp,
     }
 
     // Loop bounds and step.
     auto lb = std_constant_index(0).value;
     auto ub = std_constant_index(
                   castedDataType.getDimSize(castedDataType.getRank() - 1))
                   .value;
     auto step = std_constant_index(1).value;
+    // TransferWriteOps that operate on tensors return the modified tensor and
+    // require a loop state.
+    auto loopState = Strategy<OpTy>::initialLoopState(xferOp);
 
     // Generate for loop.
-    rewriter.create<scf::ForOp>(
-        xferOp.getLoc(), lb, ub, step, ValueRange(),
-        [&](OpBuilder &b, Location loc, Value iv, ValueRange /*loopState*/) {
+    auto result = rewriter.create<scf::ForOp>(
+        xferOp.getLoc(), lb, ub, step,
+        loopState ? ValueRange(loopState) : ValueRange(),
+        [&](OpBuilder &b, Location loc, Value iv, ValueRange loopState) {
           ScopedContext scope(b, loc);
-          generateInBoundsCheck(
+          Type stateType = loopState.empty() ? Type() : loopState[0].getType();
+
+          auto result = generateInBoundsCheck(
               xferOp, iv, b, unpackedDim(xferOp),
+              stateType ? TypeRange(stateType) : TypeRange(),
               /*inBoundsCase=*/
               [&](OpBuilder &b, Location /*loc*/) {
                 // Create new transfer op.
                 OpTy newXfer = Strategy<OpTy>::rewriteOp(
-                    b, this->options, xferOp, castedDataBuffer, iv);
+                    b, this->options, xferOp, castedDataBuffer, iv, loopState);
 
                 // If old transfer op has a mask: Set mask on new transfer op.
                 // Special case: If the mask of the old transfer op is 1D and
                 // the
                 //               unpacked dim is not a broadcast, no mask is
                 //               needed on the new transfer op.
                 if (xferOp.mask() && (xferOp.isBroadcastDim(0) ||
                                       xferOp.getMaskType().getRank() > 1)) {
                   OpBuilder::InsertionGuard guard(b);
                   b.setInsertionPoint(newXfer); // Insert load before newXfer.
 
                   SmallVector<Value, 8> loadIndices;
                   Strategy<OpTy>::getBufferIndices(xferOp, loadIndices);
                   // In case of broadcast: Use same indices to load from memref
                   // as before.
                   if (!xferOp.isBroadcastDim(0))
                     loadIndices.push_back(iv);
 
                   auto mask = memref_load(castedMaskBuffer, loadIndices);
                   rewriter.updateRootInPlace(
                       newXfer, [&]() { newXfer.maskMutable().assign(mask); });
                 }
+
+                return loopState.empty() ? Value() : newXfer->getResult(0);
               },
               /*outOfBoundsCase=*/
               [&](OpBuilder &b, Location /*loc*/) {
-                Strategy<OpTy>::handleOutOfBoundsDim(b, xferOp,
-                                                     castedDataBuffer, iv);
+                return Strategy<OpTy>::handleOutOfBoundsDim(
+                    b, xferOp, castedDataBuffer, iv, loopState);
               });
-          b.create<scf::YieldOp>(loc);
+
+          maybeYieldValue(!loopState.empty(), b, loc, result);
         });
 
-    Strategy<OpTy>::cleanup(rewriter, xferOp);
+    Strategy<OpTy>::cleanup(rewriter, xferOp, result);
     return success();
   }
 };
 
 } // namespace lowering_n_d
 
 namespace lowering_n_d_unrolled {
 

mlir/test/Conversion/VectorToSCF/tensor-transfer-ops.mlir

+// RUN: mlir-opt %s -convert-vector-to-scf -split-input-file -allow-unregistered-dialect | FileCheck %s
+
+// CHECK-LABEL: func @transfer_read_2d(
+//       CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<vector<4x9xf32>>
+//       CHECK: %[[CASTED:.*]] = vector.type_cast %[[ALLOC]] : memref<vector<4x9xf32>> to memref<4xvector<9xf32>>
+//       CHECK: scf.for {{.*}} {
+//       CHECK:   %[[READ:.*]] = vector.transfer_read %{{.*}}[{{.*}}], %cst {in_bounds = [true]} : tensor<?x?xf32>, vector<9xf32>
+//       CHECK:   memref.store %[[READ]], %[[CASTED]][%{{.*}}] : memref<4xvector<9xf32>>
+//       CHECK: }
+//       CHECK: %[[LOADED:.*]] = memref.load %[[ALLOC]][] : memref<vector<4x9xf32>>
+//       CHECK: return %[[LOADED]] : vector<4x9xf32>
+func @transfer_read_2d(%A : tensor<?x?xf32>, %base1 : index, %base2 : index)
+    -> (vector<4x9xf32>){
+  %p = constant -42.0: f32
+  %f = vector.transfer_read %A[%base1, %base2], %p {in_bounds = [true, true]}
+      : tensor<?x?xf32>, vector<4x9xf32>
+  return %f : vector<4x9xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @transfer_write_2d(
+//       CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<vector<2x3xf32>>
+//       CHECK: memref.store {{.*}}, %[[ALLOC]][] : memref<vector<2x3xf32>>
+//       CHECK: %[[CASTED:.*]] = vector.type_cast %[[ALLOC]] : memref<vector<2x3xf32>> to memref<2xvector<3xf32>>
+//       CHECK: %[[RESULT:.*]] = scf.for {{.*}} iter_args(%[[STATE:.*]] = %{{.*}}) -> (tensor<?x?xf32>) {
+//       CHECK:   %[[LOADED:.*]] = memref.load %[[CASTED]][%{{.*}}] : memref<2xvector<3xf32>>
+//       CHECK:   %[[WRITE:.*]] = vector.transfer_write %[[LOADED]], %[[STATE]][{{.*}}] {in_bounds = [true]} : vector<3xf32>, tensor<?x?xf32>
+//       CHECK:   scf.yield %[[WRITE]] : tensor<?x?xf32>
+//       CHECK: }
+//       CHECK: return %[[RESULT]] : tensor<?x?xf32>
+func @transfer_write_2d(%A : tensor<?x?xf32>, %vec : vector<2x3xf32>,
+                        %base1 : index, %base2 : index) -> (tensor<?x?xf32>) {
+  %t = vector.transfer_write %vec, %A[%base1, %base2] {in_bounds = [true, true]}
+      : vector<2x3xf32>, tensor<?x?xf32>
+  return %t : tensor<?x?xf32>
+}
+