[mlir][VectorOps] Support vector transfer_read/write unrolling for memrefs with vector element type.

[andydavis1]

Summary

[mlir][VectorOps] Support vector transfer_read/write unrolling for memrefs with vector element type. When unrolling vector transfer read/write on memrefs with vector element type, the indices used to index the memref argument must be updated to reflect the unrolled operation. However, in the case of memrefs with vector element type, we need to be careful to only update the relevant memref indices.

For example, a vector transfer read with the following source/result types, memref<6x2x1xvector<2x4xf32>>, vector<2x1x2x4xf32>, should only update memref indices 1 and 2 during unrolling.

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[nicolasvasilache]

I don't see a commit message?

[andydavis1]

added commit message

[aartbik]

Make sure to add the labels

[mlir] [VectorOps]

to the commit message's title

[rriddle]

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

488

While you are here can you change this to ///?

517

Remove trivial braces.

548

Use ///

554

nit: Why not use Optional<int> instead?

[andydavis1]

addressing comments

[andydavis1]

addressing comments

[merge_guards_bot]

Unit tests: pass. 62322 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[aartbik]

note that the title still needs "[mlir] [VectorOps]"

[nicolasvasilache]

Please address Aart's comment and land.

Thanks much Andy!

[lebedev.ri]

Would it be please possible to prefix all the differentials related to MLIR with [MLIR] prefix?
This makes reading through llvm-commits mails much less infuriating.
http://llvm.org/docs/DeveloperPolicy.html#commit-messages

* When the changes are restricted to a specific part of the code (e.g. a back-end or optimization pass),
  it is customary to add a tag to the beginning of the line in square brackets.
  For example, “[SCEV] …” or “[OpenMP] …”. This helps email filters and searches for post-commit reviews.

[rriddle]

In D72965#1852331, @lebedev.ri wrote:

Would it be please possible to prefix all the differentials related to MLIR with [MLIR] prefix?
This makes reading through llvm-commits mails much less infuriating.
http://llvm.org/docs/DeveloperPolicy.html#commit-messages

* When the changes are restricted to a specific part of the code (e.g. a back-end or optimization pass),
  it is customary to add a tag to the beginning of the line in square brackets.
  For example, “[SCEV] …” or “[OpenMP] …”. This helps email filters and searches for post-commit reviews.

+1 @andydavis1 Can you please make sure to use the proper tags?

[andydavis1]

updated tags

[nicolasvasilache]

@lebedev.ri @rriddle if a non-enforced guideline makes it infuriating to interact with the system, maybe it is ripe for automation?

[andydavis1]

rebasing

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

getVectorElementwiseOpUnrollState(op, targetShape, iterationBounds, vectors,
                                       resultIndex);
   }
 
   // Unroll 'op' with 'iterationBounds' to 'targetShape'.
   return SmallVector<Value, 1>{unrollSingleResultStructuredOp(
       op, iterationBounds, vectors, resultIndex, targetShape, builder)};
 }
 
 // Generates slices of 'vectorType' according to 'sizes' and 'strides, and
 // calls 'fn' with linear index and indices for each slice.
 static void
-generateTransferOpSlices(VectorType vectorType, TupleType tupleType,
-                         ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides,
-                         ArrayRef<Value> indices, PatternRewriter &rewriter,
+generateTransferOpSlices(Type memrefElementType, VectorType vectorType,
+                         TupleType tupleType, ArrayRef<int64_t> sizes,
+                         ArrayRef<int64_t> strides, ArrayRef<Value> indices,
+                         PatternRewriter &rewriter,
                          function_ref<void(unsigned, ArrayRef<Value>)> fn) {
   // Compute strides w.r.t. to slice counts in each dimension.
   auto maybeDimSliceCounts = shapeRatio(vectorType.getShape(), sizes);
   assert(maybeDimSliceCounts.hasValue());
   auto sliceDimCounts = *maybeDimSliceCounts;
   auto basis = computeStrides(sliceDimCounts);
 
   int64_t numSlices = tupleType.size();
   unsigned numSliceIndices = indices.size();
+  // Compute 'indexOffset' at which to update 'indices', which is equal
+  // to the memref rank (indices.size) minus the effective 'vectorRank'.
+  // The effective 'vectorRank', is equal to the rank of the vector type
+  // minus the rank of the memref vector element type (if it has one).
+  //
+  // For example:
+  //
+  //   Given memref type 'memref<6x2x1xvector<2x4xf32>>' and vector
+  //   transfer_read/write ops which read/write vectors of type
+  //   'vector<2x1x2x4xf32>'. The memref rank is 3, and the effective
+  //   vector rank is 4 - 2 = 2, and so 'indexOffset' = 3 - 2 = 1.
+  //
+  unsigned vectorRank = vectorType.getRank();
+  if (auto memrefVectorElementType = memrefElementType.dyn_cast<VectorType>()) {
+    vectorRank -= memrefVectorElementType.getRank();
+  }
+  assert(numSliceIndices >= vectorRank);
+  unsigned indexOffset = numSliceIndices - vectorRank;
+
   auto *ctx = rewriter.getContext();
   for (unsigned i = 0; i < numSlices; ++i) {
     // De-linearize w.r.t. 'basis'.
     auto vectorOffsets = delinearize(i, basis);
     // Convert from unrolled vector-space offsets to element-space offsets.
     auto offsets = zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
                           vectorOffsets, sizes);
     // Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'.
     SmallVector<Value, 4> sliceIndices(numSliceIndices);
-    for (auto it : llvm::enumerate(indices)) {
+    for (unsigned j = 0; j < numSliceIndices; ++j) {
+      if (j < indexOffset) {
+        sliceIndices[j] = indices[j];
+      } else {
         auto expr = getAffineDimExpr(0, ctx) +
-                  getAffineConstantExpr(offsets[it.index()], ctx);
+                    getAffineConstantExpr(offsets[j - indexOffset], ctx);
         auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
-      sliceIndices[it.index()] = rewriter.create<AffineApplyOp>(
-          it.value().getLoc(), map, ArrayRef<Value>(it.value()));
+        sliceIndices[j] = rewriter.create<AffineApplyOp>(
+            indices[j].getLoc(), map, ArrayRef<Value>(indices[j]));
+      }
     }
     // Call 'fn' to generate slice 'i' at 'sliceIndices'.
     fn(i, sliceIndices);
   }
 }
 
+// Returns true if 'map' is a suffix of an identity affine map, false otherwise.
+// Example: affine_map<(d0, d1, d2, d3) -> (d2, d3)>
+static bool isIdentitySuffix(AffineMap map) {
+  if (map.getNumDims() < map.getNumResults())
+    return false;
+  ArrayRef<AffineExpr> results = map.getResults();
+  int lastPos = -1;
+  for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
+    auto expr = results[i].dyn_cast<AffineDimExpr>();
+    if (!expr)
+      return false;
+    int currPos = static_cast<int>(expr.getPosition());
+    if (lastPos >= 0 && currPos != lastPos + 1)
+      return false;
+    lastPos = currPos;
+  }
+  return true;
+}
+
 namespace {
 // Splits vector TransferReadOp into smaller TransferReadOps based on slicing
 // scheme of its unique ExtractSlicesOp user.
 struct SplitTransferReadOp : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(vector::TransferReadOp xferReadOp,
                                      PatternRewriter &rewriter) const override {
     // TODO(andydavis, ntv) Support spliting TransferReadOp with non-identity
     // permutation maps. Repurpose code from MaterializeVectors transformation.
-    if (!xferReadOp.permutation_map().isIdentity())
+    if (!isIdentitySuffix(xferReadOp.permutation_map()))
       return matchFailure();
     // Return unless the unique 'xferReadOp' user is an ExtractSlicesOp.
     Value xferReadResult = xferReadOp.getResult();
     auto extractSlicesOp =
         dyn_cast<vector::ExtractSlicesOp>(*xferReadResult.getUsers().begin());
     if (!xferReadResult.hasOneUse() || !extractSlicesOp)
       return matchFailure();
 
     // Get 'sizes' and 'strides' parameters from ExtractSlicesOp user.
     auto sourceVectorType = extractSlicesOp.getSourceVectorType();
     auto resultTupleType = extractSlicesOp.getResultTupleType();
     SmallVector<int64_t, 4> sizes;
     extractSlicesOp.getSizes(sizes);
     SmallVector<int64_t, 4> strides;
     extractSlicesOp.getStrides(strides);
     assert(llvm::all_of(strides, [](int64_t s) { return s == 1; }));
 
     Location loc = xferReadOp.getLoc();
+    auto memrefElementType =
+        xferReadOp.memref().getType().cast<MemRefType>().getElementType();
     int64_t numSlices = resultTupleType.size();
     SmallVector<Value, 4> vectorTupleValues(numSlices);
     SmallVector<Value, 4> indices(xferReadOp.indices().begin(),
                                   xferReadOp.indices().end());
     auto createSlice = [&](unsigned index, ArrayRef<Value> sliceIndices) {
       // Get VectorType for slice 'i'.
       auto sliceVectorType = resultTupleType.getType(index);
       // Create split TransferReadOp for 'sliceUser'.
       vectorTupleValues[index] = rewriter.create<vector::TransferReadOp>(
           loc, sliceVectorType, xferReadOp.memref(), sliceIndices,
           xferReadOp.permutation_map(), xferReadOp.padding());
     };
-    generateTransferOpSlices(sourceVectorType, resultTupleType, sizes, strides,
-                             indices, rewriter, createSlice);
+    generateTransferOpSlices(memrefElementType, sourceVectorType,
+                             resultTupleType, sizes, strides, indices, rewriter,
+                             createSlice);
 
     // Create tuple of splice xfer read operations.
     Value tupleOp = rewriter.create<vector::TupleOp>(loc, resultTupleType,
                                                      vectorTupleValues);
     // Replace 'xferReadOp' with result 'insertSlicesResult'.
     rewriter.replaceOpWithNewOp<vector::InsertSlicesOp>(
         xferReadOp, sourceVectorType, tupleOp, extractSlicesOp.sizes(),
         extractSlicesOp.strides());
     return matchSuccess();
   }
 };
 
 // Splits vector TransferWriteOp into smaller TransferWriteOps for each source.
 struct SplitTransferWriteOp : public OpRewritePattern<vector::TransferWriteOp> {
   using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(vector::TransferWriteOp xferWriteOp,
                                      PatternRewriter &rewriter) const override {
     // TODO(andydavis, ntv) Support spliting TransferWriteOp with non-identity
     // permutation maps. Repurpose code from MaterializeVectors transformation.
-    if (!xferWriteOp.permutation_map().isIdentity())
+    if (!isIdentitySuffix(xferWriteOp.permutation_map()))
       return matchFailure();
     // Return unless the 'xferWriteOp' 'vector' operand is an 'InsertSlicesOp'.
     auto *vectorDefOp = xferWriteOp.vector().getDefiningOp();
     auto insertSlicesOp = dyn_cast_or_null<vector::InsertSlicesOp>(vectorDefOp);
     if (!insertSlicesOp)
       return matchFailure();
 
     // Get TupleOp operand of 'insertSlicesOp'.
     auto tupleOp = dyn_cast_or_null<vector::TupleOp>(
         insertSlicesOp.vectors().getDefiningOp());
     if (!tupleOp)
       return matchFailure();
 
     // Get 'sizes' and 'strides' parameters from InsertSlicesOp user.
     auto sourceTupleType = insertSlicesOp.getSourceTupleType();
     auto resultVectorType = insertSlicesOp.getResultVectorType();
     SmallVector<int64_t, 4> sizes;
     insertSlicesOp.getSizes(sizes);
     SmallVector<int64_t, 4> strides;
     insertSlicesOp.getStrides(strides);
 
     Location loc = xferWriteOp.getLoc();
+    auto memrefElementType =
+        xferWriteOp.memref().getType().cast<MemRefType>().getElementType();
     SmallVector<Value, 4> indices(xferWriteOp.indices().begin(),
                                   xferWriteOp.indices().end());
     auto createSlice = [&](unsigned index, ArrayRef<Value> sliceIndices) {
       // Create split TransferWriteOp for source vector 'tupleOp.operand[i]'.
       rewriter.create<vector::TransferWriteOp>(
           loc, tupleOp.getOperand(index), xferWriteOp.memref(), sliceIndices,
           xferWriteOp.permutation_map());
     };
-    generateTransferOpSlices(resultVectorType, sourceTupleType, sizes, strides,
-                             indices, rewriter, createSlice);
+    generateTransferOpSlices(memrefElementType, resultVectorType,
+                             sourceTupleType, sizes, strides, indices, rewriter,
+                             createSlice);
 
     // Erase old 'xferWriteOp'.
     rewriter.eraseOp(xferWriteOp);
     return matchSuccess();
   }
 };
 
 // Patter rewrite which forward tuple elements to their users.

mlir/test/Dialect/VectorOps/vector-transforms.mlir

 // RUN: mlir-opt %s -test-vector-to-vector-conversion | FileCheck %s
 
 // CHECK-DAG: #[[MAP0:map[0-9]+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[MAP1:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
 
 // CHECK-LABEL: func @add4x2
 //      CHECK: %[[ES1:.*]] = vector.extract_slices %{{.*}}, [2, 2], [1, 1] : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
 // CHECK-NEXT: %[[ES2:.*]] = vector.extract_slices %{{.*}}, [2, 2], [1, 1] : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
 // CHECK-NEXT: %[[TG1:.*]] = vector.tuple_get %[[ES1]], 0 : tuple<vector<2x2xf32>, vector<2x2xf32>>
 // CHECK-NEXT: %[[TG2:.*]] = vector.tuple_get %[[ES2]], 0 : tuple<vector<2x2xf32>, vector<2x2xf32>>
 // CHECK-NEXT: %[[A1:.*]] = addf %[[TG1]], %[[TG2]] : vector<2x2xf32>
 // CHECK-NEXT: %[[TG3:.*]] = vector.tuple_get %[[ES1]], 1 : tuple<vector<2x2xf32>, vector<2x2xf32>>
affine.for %arg3 = 0 to %arg1 step 4 {
       %4 = vector.transfer_read %0[%arg2, %arg3], %cst  {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
       %5 = vector.transfer_read %1[%arg2, %arg3], %cst  {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
       %6 = addf %4, %5 : vector<4x4xf32>
       vector.transfer_write %6, %2[%arg2, %arg3] {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : vector<4x4xf32>, memref<?x?xf32>
     }
   }
   return
 }
+
+// CHECK-LABEL: func @vector_transfers_vector_element_type
+//      CHECK: %[[C0:.*]] = constant 0 : index
+//      CHECK: %[[C1:.*]] = constant 1 : index
+//      CHECK: %[[VTR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {permutation_map = #[[MAP1]]} : memref<6x2x1xvector<2x4xf32>>, vector<1x1x2x4xf32>
+// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C1]], %[[C0]]], %{{.*}} {permutation_map = #[[MAP1]]} : memref<6x2x1xvector<2x4xf32>>, vector<1x1x2x4xf32>
+// CHECK-NEXT: vector.transfer_write %[[VTR0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {permutation_map = #[[MAP1]]} : vector<1x1x2x4xf32>, memref<6x2x1xvector<2x4xf32>>
+// CHECK-NEXT: vector.transfer_write %[[VTR1]], %{{.*}}[%[[C0]], %[[C1]], %[[C0]]] {permutation_map = #[[MAP1]]} : vector<1x1x2x4xf32>, memref<6x2x1xvector<2x4xf32>>
+
+func @vector_transfers_vector_element_type() {
+  %c0 = constant 0 : index
+  %cf0 = constant 0.000000e+00 : f32
+  %vf0 = splat %cf0 : vector<2x4xf32>
+
+  %0 = alloc() : memref<6x2x1xvector<2x4xf32>>
+
+  %1 = vector.transfer_read %0[%c0, %c0, %c0], %vf0
+      {permutation_map = affine_map<(d0, d1, d2) -> (d1, d2)>}
+        : memref<6x2x1xvector<2x4xf32>>, vector<2x1x2x4xf32>
+
+  %2 = vector.extract_slices %1, [1, 1, 2, 4], [1, 1, 1, 1]
+    : vector<2x1x2x4xf32> into tuple<vector<1x1x2x4xf32>, vector<1x1x2x4xf32>>
+  %3 = vector.tuple_get %2, 0 : tuple<vector<1x1x2x4xf32>, vector<1x1x2x4xf32>>
+  %4 = vector.tuple_get %2, 1 : tuple<vector<1x1x2x4xf32>, vector<1x1x2x4xf32>>
+  %5 = vector.tuple %3, %4 : vector<1x1x2x4xf32>, vector<1x1x2x4xf32>
+  %6 = vector.insert_slices %5, [1, 1, 2, 4], [1, 1, 1, 1]
+    : tuple<vector<1x1x2x4xf32>, vector<1x1x2x4xf32>> into vector<2x1x2x4xf32>
+
+  vector.transfer_write %6, %0[%c0, %c0, %c0]
+    {permutation_map = affine_map<(d0, d1, d2) -> (d1, d2)>}
+      : vector<2x1x2x4xf32>, memref<6x2x1xvector<2x4xf32>>
+
+  return
+}