[mlir] Support dimension permutations in ProgressiveVectorToSCF

[springerm]

Summary

This commit adds support for dimension permutations in permutation maps of vector transfer ops.

Depends On D100946

[springerm]

(no change)

[springerm]

(no change)

[springerm]

(no change)

[springerm]

rebase

[springerm]

(no change)

[springerm]

(no change)

[nicolasvasilache]

Looks good, note that there is also a pass that rewrites the transfer into memref.transpose and identity permutation map.

Would be good to start thinking (in the background) about a mechanism where we can more easily expose the decision of either:

1. what you do here.
2. what I mention above.
3. lowering to vector.transpose

We will be interested in the interaction of the above with vector.shape_cast and vector.reshape as well as masking.

[springerm]

(rebase)

[springerm]

(new export)

mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp

static Value setAllocAtFunctionEntry(MemRefType type, Operation *op) {
   b.setInsertionPointToStart(&scope->getRegion(0).front());
   Value res = memref_alloca(type);
   return res;
 }
 
 /// Given a vector transfer op, calculate which dimension of the `source`
 /// memref should be unpacked in the next application of TransferOpConversion.
 template <typename OpTy>
-static int64_t unpackedDim(OpTy xferOp) {
-  return xferOp.getShapedType().getRank() - xferOp.getVectorType().getRank();
+static unsigned unpackedDim(OpTy xferOp) {
+  auto map = xferOp.permutation_map();
+  // TODO: Handle broadcast
+  auto expr = map.getResult(0).template dyn_cast<AffineDimExpr>();
+  assert(expr && "Expected AffineDimExpr in permutation map result");
+  return expr.getPosition();
+}
+
+/// Compute the permutation map for the new (N-1)-D vector transfer op. This
+/// map is identical to the current permutation map, but the first result is
+/// omitted.
+template <typename OpTy>
+static AffineMap unpackedPermutationMap(OpTy xferOp, OpBuilder &builder) {
+  auto map = xferOp.permutation_map();
+  return AffineMap::get(map.getNumDims(), 0, map.getResults().drop_front(),
+                        builder.getContext());
 }
 
 /// Calculate the indices for the new vector transfer op.
 ///
 /// E.g.: transfer_read %A[%a, %b, %c, %d] ... : vector<5x4x3xf32> ...
 ///       --> transfer_read %A[%a, %b + iv, %c, %d] ... vector<4x3f32>
 ///                                 ^^^^^^
 ///              `iv` is the iteration variable of the (new) surrounding loop.

 ///   (out-of-bounds case)
 /// }
 /// ```
 /// This function variant returns the value returned by `inBoundsCase` or
 /// `outOfBoundsCase`. The MLIR type of the return value must be specified in
 /// `resultTypes`.
 template <typename OpTy>
 static Value generateInBoundsCheck(
-    OpTy xferOp, Value iv, OpBuilder &builder, int64_t dim,
+    OpTy xferOp, Value iv, OpBuilder &builder, unsigned dim,
     TypeRange resultTypes,
     function_ref<Value(OpBuilder &, Location)> inBoundsCase,
     function_ref<Value(OpBuilder &, Location)> outOfBoundsCase = nullptr) {
   bool hasRetVal = !resultTypes.empty();
   if (!xferOp.isDimInBounds(0)) {
     auto memrefDim = memref_dim(xferOp.source(), std_constant_index(dim));
     using edsc::op::operator+;
     auto memrefIdx = xferOp.indices()[dim] + iv;
static void rewriteOp(PatternRewriter &rewriter, TransferReadOp xferOp,
     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 map = getTransferMinorIdentityMap(xferOp.getShapedType(), vecType);
     auto inBoundsAttr = dropFirstElem(rewriter, xferOp.in_boundsAttr());
-    auto newXfer = vector_transfer_read(vecType, xferOp.source(), xferIndices,
-                                        AffineMapAttr::get(map),
-                                        xferOp.padding(), Value(), inBoundsAttr)
-                       .value;
+    auto newXfer =
+        vector_transfer_read(
+            vecType, xferOp.source(), xferIndices,
+            AffineMapAttr::get(unpackedPermutationMap(xferOp, rewriter)),
 
     if (vecType.getRank() > kTargetRank)
       newXfer.getDefiningOp()->setAttr(kPassLabel, rewriter.getUnitAttr());
 
     memref_store(newXfer, buffer, storeIndices);
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension: Write
static void rewriteOp(PatternRewriter &rewriter, TransferWriteOp xferOp,
     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 map = getTransferMinorIdentityMap(xferOp.getShapedType(), vecType);
     auto inBoundsAttr = dropFirstElem(rewriter, xferOp.in_boundsAttr());
-    auto newXfer =
-        vector_transfer_write(Type(), vec, xferOp.source(), xferIndices,
-                              AffineMapAttr::get(map), Value(), inBoundsAttr);
+    auto newXfer = vector_transfer_write(
+        Type(), vec, xferOp.source(), xferIndices,
+        AffineMapAttr::get(unpackedPermutationMap(xferOp, rewriter)), Value(),
+        inBoundsAttr);
 
     if (vecType.getRank() > kTargetRank)
       newXfer.op->setAttr(kPassLabel, rewriter.getUnitAttr());
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension.
   static void handleOutOfBoundsDim(PatternRewriter &rewriter,
                                    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 checkPrepareXferOp(OpTy xferOp) {
   if (xferOp->hasAttr(kPassLabel))
     return failure();
   if (xferOp.getVectorType().getRank() <= kTargetRank)
     return failure();
   if (xferOp.mask())
     return failure();
-  if (!xferOp.permutation_map().isMinorIdentity())
-    return failure();
   return success();
 }
 
 /// Prepare a TransferReadOp for progressive lowering.
 ///
 /// 1. Allocate a temporary buffer.
 /// 2. Label the TransferReadOp, marking it eligible for progressive lowering.
 /// 3. Store the result of the TransferReadOp into the temporary buffer.

 } // namespace
 
 namespace mlir {
 
 void populateProgressiveVectorToSCFConversionPatterns(
     RewritePatternSet &patterns) {
   patterns.add<PrepareTransferReadConversion, PrepareTransferWriteConversion,
                TransferOpConversion<TransferReadOp>,
-               TransferOpConversion<TransferWriteOp>,
-               Strided1dTransferOpConversion<TransferReadOp>,
+               TransferOpConversion<TransferWriteOp>>(patterns.getContext());
+
+  if (kTargetRank == 1) {
+    patterns.add<Strided1dTransferOpConversion<TransferReadOp>,
                  Strided1dTransferOpConversion<TransferWriteOp>>(
         patterns.getContext());
   }
+}
 
 struct ConvertProgressiveVectorToSCFPass
     : public ConvertVectorToSCFBase<ConvertProgressiveVectorToSCFPass> {
   void runOnFunction() override {
     RewritePatternSet patterns(getFunction().getContext());
     populateProgressiveVectorToSCFConversionPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
 } // namespace mlir
 
 std::unique_ptr<Pass> mlir::createProgressiveConvertVectorToSCFPass() {
   return std::make_unique<ConvertProgressiveVectorToSCFPass>();
 }

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-2d.mlir

func @transfer_read_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     memref<?x?xf32>, vector<4x9xf32>
   vector.print %f: vector<4x9xf32>
   return
 }
 
+func @transfer_read_2d_transposed(
+    %A : memref<?x?xf32>, %base1: index, %base2: index) {
+  %fm42 = constant -42.0: f32
+  %f = vector.transfer_read %A[%base1, %base2], %fm42
+      {permutation_map = affine_map<(d0, d1) -> (d1, d0)>} :
+    memref<?x?xf32>, vector<4x9xf32>
+  vector.print %f: vector<4x9xf32>
+  return
+}
+
 func @transfer_write_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fn1 = constant -1.0 : f32
   %vf0 = splat %fn1 : vector<1x4xf32>
   vector.transfer_write %vf0, %A[%base1, %base2]
     {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     vector<1x4xf32>, memref<?x?xf32>
   return
 }
scf.for %j = %c0 to %second step %c1 {
         %fj = sitofp %j32 : i32 to f32
         %fres = addf %fi10, %fj : f32
         memref.store %fres, %A[%i, %j] : memref<?x?xf32>
     }
   }
   // On input, memory contains [[ 0, 1, 2, ...], [10, 11, 12, ...], ...]
   // Read shifted by 2 and pad with -42:
   call @transfer_read_2d(%A, %c1, %c2) : (memref<?x?xf32>, index, index) -> ()
+  // Same as above, but transposed
+  call @transfer_read_2d_transposed(%A, %c1, %c2)
+      : (memref<?x?xf32>, index, index) -> ()
   // Write into memory shifted by 3
   call @transfer_write_2d(%A, %c3, %c1) : (memref<?x?xf32>, index, index) -> ()
   // Read shifted by 0 and pad with -42:
   call @transfer_read_2d(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()
+  // Same as above, but transposed
+  call @transfer_read_2d_transposed(%A, %c0, %c0)
+      : (memref<?x?xf32>, index, index) -> ()
   return
 }
 
 // CHECK: ( ( 12, 13, -42, -42, -42, -42, -42, -42, -42 ), ( 22, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+// CHECK: ( ( 12, 22, -42, -42, -42, -42, -42, -42, -42 ), ( 13, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
 // CHECK: ( ( 0, 1, 2, 3, -42, -42, -42, -42, -42 ), ( 10, 11, 12, 13, -42, -42, -42, -42, -42 ), ( 20, 21, 22, 23, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+// CHECK: ( ( 0, 10, 20, -42, -42, -42, -42, -42, -42 ), ( 1, 11, 21, -42, -42, -42, -42, -42, -42 ), ( 2, 12, 22, -42, -42, -42, -42, -42, -42 ), ( 3, 13, 23, -42, -42, -42, -42, -42, -42 ) )

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-3d.mlir

func @transfer_read_3d(%A : memref<?x?x?x?xf32>,
                        %o: index, %a: index, %b: index, %c: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42
       : memref<?x?x?x?xf32>, vector<2x5x3xf32>
   vector.print %f: vector<2x5x3xf32>
   return
 }
 
+func @transfer_read_3d_transposed(%A : memref<?x?x?x?xf32>,
+                                  %o: index, %a: index, %b: index, %c: index) {
+  %fm42 = constant -42.0: f32
+  %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42
+      {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1)>}
+      : memref<?x?x?x?xf32>, vector<3x5x3xf32>
+  vector.print %f: vector<3x5x3xf32>
+  return
+}
+
 func @transfer_write_3d(%A : memref<?x?x?x?xf32>,
                         %o: index, %a: index, %b: index, %c: index) {
   %fn1 = constant -1.0 : f32
   %vf0 = splat %fn1 : vector<2x9x3xf32>
   vector.transfer_write %vf0, %A[%o, %a, %b, %c]
       : vector<2x9x3xf32>, memref<?x?x?x?xf32>
   return
 }
func @entry() {
   }
 
   call @transfer_read_3d(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
   call @transfer_write_3d(%A, %c0, %c0, %c1, %c1)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
   call @transfer_read_3d(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+  call @transfer_read_3d_transposed(%A, %c0, %c0, %c0, %c0)
+      : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
   return
 }
 
 // CHECK: ( ( ( 0, 0, -42 ), ( 2, 3, -42 ), ( 4, 6, -42 ), ( 6, 9, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, 33, -42 ), ( 24, 36, -42 ), ( 26, 39, -42 ), ( -42, -42, -42 ) ) )
 // CHECK: ( ( ( 0, 0, -42 ), ( 2, -1, -42 ), ( 4, -1, -42 ), ( 6, -1, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, -1, -42 ), ( 24, -1, -42 ), ( 26, -1, -42 ), ( -42, -42, -42 ) ) )
+// CHECK: ( ( ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ) ), ( ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ) ), ( ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ) ) )

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-to-loops.mlir

 // RUN: mlir-opt %s -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void  \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext,%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
+// RUN: mlir-opt %s -test-progressive-convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void  \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext,%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
 #map0 = affine_map<(d0, d1) -> (d1, d0)>
 #map1 = affine_map<(d0, d1) -> (d1)>
 
 func private @print_memref_f32(memref<*xf32>)
 
 func @alloc_2d_filled_f32(%arg0: index, %arg1: index) -> memref<?x?xf32> {
   %c0 = constant 0 : index
   %c1 = constant 1 : index