[mlir][vector] Add lowering of Transfer_read with leading dim broadcast and permutation map

[ThomasRaoux]

Summary

Convert transfer_read ops with permutation maps into simpler transfer_read with minority map + vector.braodcast and vector.transpose

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2849

I am not sure I understand this case, valid minor identity maps are (d0, d1, d2), (d1, d2) and (d2) ?

2850

s/permutate/permute/g (both comment and code please).

2852

Can you expose this new functionality in AffineMap.h?

2865

s/broadcasted/broadcast/g

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2849

Agree, but the idea is to get to a minor identity with broadcasting, meaning some dimensions can be replaced by 0 (this match the definition in isMinorIdentityWithBroadcasting). So (d1, 0) is (d1, d2) with broadcast on dimension d2. This means we can do a transfer read with no permute to <Nx1> and broadcast the lower dimension.

If the map is (d0, d1, d2) -> (d1) it would require first transforming it to (d0, d1, d2) -> (d1, 0) which needs more than a transpose. That's why I currently don't support it.

[ThomasRaoux]

Address review comments

[sgrechanik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2876

typo

mlir/lib/IR/AffineMap.cpp

152

I think that a bitvector (SmallBitVector?) could be better here, but I don't insist.

157

The "or" counterpart of this comment is missing.

[ThomasRaoux]

Address review comments

[ThomasRaoux]

mlir/lib/IR/AffineMap.cpp

152

Makes sense, changed it.

[nicolasvasilache]

mlir/include/mlir/IR/AffineMap.h

116

Just copy pasting my other comment here:

I think spelling it out a little more would be useful.
First I'd drop the N-D -> ND and focus on N-D -> (N-1)-D and/or N-D -> (N-2)-D as I think the asymmetry helps documentation.

Then I'd write something like:

* (d0, d1, d2) -> (d0, 0) => cannot be mapped to a minor identity by permutation + broadcast
* (d0, d1, d2) -> (0, d1) => maps to minor identity (d1, 0 = d2) with perm = [1, 0] and broadcast d2
...
a few more examples to build intuition
I would put that same documentation very loudly both in AffineMap.h and here

123

Would it be better API-wise to also pass the SmallBitVector &broadcastDims to fill that too?
Atm the broadcast is still implicit and subject to magic.

I'd also make it clear that it would be first permute by applying permuteDims then set the true dimensions of broadcastDims to 0.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2849

Right, so I think spelling it out a little more would be useful.
First I'd drop the N-D -> ND and focus on N-D -> (N-1)-D and/or N-D -> (N-2)-D as I think the asymmetry helps documentation.

Then I'd write something like:

* (d0, d1, d2) -> (d0, 0) => cannot be mapped to a minor identity by permutation + broadcast
* (d0, d1, d2) -> (0, d1) => maps to minor identity (d1, 0 = d2) with perm = [1, 0] and broadcast d2
...
a few more examples to build intuition

I would put that same documentation very loudly both in AffineMap.h and here

2853

Would also be nice to mention there is a dual transformation that takes vector.transfer + vector.transpose and turns it into linalg.transpose + vector.transfer to do stuff in memory.
This pattern and that transformations are related, it is target-specific what combination should be used.

2866

I think if isPermutationOfMinorIdentityWithBroadcasting also filled a broadcast bitvector, it would be a more natural read to create the map based on that.
It is not immediately clear to me what the swappy in-place behavior with modulos does.

2880

Can you use https://sourcegraph.com/github.com/llvm/llvm-project@d75a611/-/blob/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h#L127:6 ?

void applyPermutationToVector(SmallVector<T, N> &inVec,
                              ArrayRef<unsigned> permutation)

mlir/test/Dialect/Vector/vector-transfer-lowering.mlir

230

can you add a few whitespaces to isolate the different parts of the test ?

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2904

This is an interesting case where an internal broadcast is left.
Will be interesting to see how this composes with vector -> scf lowering (not for this revision).

2937

can this be written as something resembling newMask = rewriter.getBoolArrayAttr(op.masked().getAsValueRange<bool>().take_back(reducedShapeRank)); ?

mlir/lib/IR/AffineMap.cpp

144

same comment re adding a noisy, spelly comment.

[ThomasRaoux]

Address more review comments.

mlir/include/mlir/IR/AffineMap.h

116

Makes sense, I added more examples, hopefully this makes it easier to understand.

123

I tried to do it but I feel like it adds unnecessary complexity to this function and in my current case I am not able to use it. The broadcast dimensions can be accessed after permutation using the existing getMinorIdentityMap

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2866

I'm not sure using the broadcast bitvector is easier. Here we cannot move all the broadcast dim to the front as we still want a minor identity map after the broadcasts. I added some comments explaining better what this does. Hopefully it is clearer.

2880

I would need to reverse the permutation first and I don't see any easy way to do it? Otherwise I could addd a different helper applyReversePermutationToVector? Is it okay to depend on linalg utils here?

2904

Makes sense. Note that it already compose well with the existing TransferReadToVectorLoadLowering pattern which load to Nx1xM then broadcast the 1.

2937

Yes, finally found the magic formula.

[ThomasRaoux]

Simplify the logic based on comments

[ThomasRaoux]

@nicolasvasilache I removed the confusing part we were discussing (the "rotating" the permutation map). This simplifies the code and should make it more clear what it does. Please take another look.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2866

This part was confusing and unnecessary complicated indeed. I removed it. It is better handled in isPermutationOfMinorIdentityWithBroadcasting by making sure we start the map with 0s in case numResults > numDims

[sgrechanik]

mlir/include/mlir/IR/AffineMap.h

127

Could you write out the permutation for this last example too?

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2885

nit: probably it would be better to rename it "permutation" because it's a single permutation.

[ThomasRaoux]

Address review comments

mlir/include/mlir/IR/AffineMap.h

public:
   /// dimensions which are indicated by value 0 in the result. If
   /// `broadcastedDims` is not null, it will be populated with the indices of
   /// the broadcasted dimensions in the result array.
   /// Example: affine_map<(d0, d1, d2, d3, d4) -> (0, d2, 0, d4)>
   ///          (`broadcastedDims` will contain [0, 2])
   bool isMinorIdentityWithBroadcasting(
       SmallVectorImpl<unsigned> *broadcastedDims = nullptr) const;
 
+  /// Return true if this affine map can be converted to a minor identity with
+  /// broadcast by doing a permute. Return a permutation (there may be
+  /// several) to apply to get to a minor identity with broadcasts. Ex: (d0, d1,
+  /// d2) -> (0, d1) returns true and a permutation [1, 0] (d0, d1, d2) -> (d1)
+  /// returns false as it cannot be permuted into a minor identity with
+  /// broadcasting.
+  bool isPermutationOfMinorIdentityWithBroadcasting(
+      SmallVectorImpl<unsigned> &permutedDims) const;
+
   /// Returns true if this affine map is an empty map, i.e., () -> ().
   bool isEmpty() const;
 
   /// Returns true if this affine map is a single result constant function.
   bool isSingleConstant() const;
 
   /// Returns the constant result of this map. This methods asserts that the map
   /// has a single constant result.
   int64_t getSingleConstantResult() const;
 
   // Prints affine map to 'os'.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

LogicalResult matchAndRewrite(vector::TransferWriteOp write,
     if (write.hasMaskedDim())
       return failure();
     rewriter.replaceOpWithNewOp<vector::StoreOp>(
         write, write.vector(), write.source(), write.indices());
     return success();
   }
 };
 
+/// Lower transfer_read op with permutation into a transfer_read with a
+/// permutation map composed of leading zeros followed by a minor identiy +
+/// vector.transpose op.
+/// Ex: vector.transfer_read ...
+///         permutation_map: (d0, d1, d2, d3) -> (d3, d1, 0, 0)
+/// into:
+///     %v = vector.transfer_read ...
+///         permutation_map: (d0, d1, d2, d3) -> (0, d1, 0, d3)
+///     vector.transpose %v, [3, 1, 0, 2]
+struct TransferReadPermutationLowering
+    : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<unsigned> permutations;
+    AffineMap map = op.permutation_map();
+    if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutations))
+      return failure();
+
+    // Move the broadcast towards the higher dim.
+    if (map.getNumResults() > map.getNumInputs()) {
+      unsigned rotation = map.getNumResults() - map.getNumInputs();
+      for (unsigned &i : permutations) {
+        i = (i + rotation) % permutations.size();
+      }
+    }
+    AffineMap permutationMap =
+        map.getPermutationMap(permutations, op.getContext());
+    if (permutationMap.isIdentity())
+      return failure();
+    // Caluclate the map of the new read by applying the inverse permuation.
+    permutationMap = inversePermutation(permutationMap);
+    AffineMap newMap = permutationMap.compose(map);
+    // Apply the reverse transpose to deduce the type of the transfer_read.
+    ArrayRef<int64_t> originalShape = op.getVectorType().getShape();
+    SmallVector<int64_t> newVectorShape(originalShape.size());
+    for (auto pos : llvm::enumerate(permutations)) {
+      newVectorShape[pos.value()] = originalShape[pos.index()];
+    }
+    VectorType newReadType =
+        VectorType::get(newVectorShape, op.getVectorType().getElementType());
+    Value newRead = rewriter.create<vector::TransferReadOp>(
+        op.getLoc(), newReadType, op.source(), op.indices(), newMap,
+        op.padding(), op.masked() ? *op.masked() : ArrayAttr());
+    SmallVector<int64_t> transposePerm(permutations.begin(),
+                                       permutations.end());
+    rewriter.replaceOpWithNewOp<vector::TransposeOp>(op, newRead,
+                                                     transposePerm);
+    return success();
+  }
+};
+
+/// Lower transfer_read op with broadcast in the leading dimensions into
+/// transfer_read of lower rank + vector.broadcast.
+/// Ex: vector.transfer_read ...
+///         permutation_map: (d0, d1, d2, d3) -> (0, d1, 0, d3)
+/// into:
+///     %v = vector.transfer_read ...
+///         permutation_map: (d0, d1, d2, d3) -> (d1, 0, d3)
+///     vector.broadcast %v
+struct TransferOpReduceRank : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp op,
+                                PatternRewriter &rewriter) const override {
+    AffineMap map = op.permutation_map();
+    unsigned numLeadingBroadcast = 0;
+    for (auto expr : map.getResults()) {
+      auto dimExpr = expr.dyn_cast<AffineConstantExpr>();
+      if (!dimExpr || dimExpr.getValue() != 0)
+        break;
+      numLeadingBroadcast++;
+    }
+    // If there are no leading zeros in the map there is nothing to do.
+    if (numLeadingBroadcast == 0)
+      return failure();
+    VectorType originalVecType = op.getVectorType();
+    unsigned reducedShapeRank = originalVecType.getRank() - numLeadingBroadcast;
+    // Calculate new map, vector type and masks without the leading zeros.
+    AffineMap newMap = AffineMap::get(
+        map.getNumDims(), 0, map.getResults().take_back(reducedShapeRank),
+        op.getContext());
+    // Only remove the leading zeros if the rest of the map is a minor identity
+    // with broadasting. Otherwise we first want to permute the map.
+    if (!newMap.isMinorIdentityWithBroadcasting())
+      return failure();
+    SmallVector<int64_t> newShape = llvm::to_vector<4>(
+        originalVecType.getShape().take_back(reducedShapeRank));
+    VectorType newReadType =
+        VectorType::get(newShape, originalVecType.getElementType());
+    ArrayAttr newMask;
+    if (op.masked()) {
+      SmallVector<bool> masks(reducedShapeRank);
+      for (unsigned i = 0; i < reducedShapeRank; i++) {
+        masks[i] = op.isMaskedDim(i + numLeadingBroadcast);
+      }
+      newMask = rewriter.getBoolArrayAttr(masks);
+    }
+    Value newRead = rewriter.create<vector::TransferReadOp>(
+        op.getLoc(), newReadType, op.source(), op.indices(), newMap,
+        op.padding(), newMask);
+    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, originalVecType,
+                                                     newRead);
+    return success();
+  }
+};
+
 // Trims leading one dimensions from `oldType` and returns the result type.
 // Returns `vector<1xT>` if `oldType` only has one element.
 static VectorType trimLeadingOneDims(VectorType oldType) {
   ArrayRef<int64_t> oldShape = oldType.getShape();
   ArrayRef<int64_t> newShape =
       oldShape.drop_while([](int64_t dim) { return dim == 1; });
   // Make sure we have at least 1 dimension per vector type requirements.
   if (newShape.empty())
patterns.insert<TransposeOpLowering,
                   ContractionOpToMatmulOpLowering,
                   ContractionOpToOuterProductOpLowering>(parameters, context);
   // clang-format on
 }
 
 void mlir::vector::populateVectorTransferLoweringPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<TransferReadToVectorLoadLowering,
-                  TransferWriteToVectorStoreLowering>(context);
+                  TransferWriteToVectorStoreLowering,
+                  TransferReadPermutationLowering, TransferOpReduceRank>(
+      context);
 }

mlir/lib/IR/AffineMap.cpp

if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
         return false;
     } else {
       return false;
     }
   }
   return true;
 }
 
+/// Returns true if this affine map can be transformed into a minor identity
+/// with broadcasting using a permuation.
+bool AffineMap::isPermutationOfMinorIdentityWithBroadcasting(
+    SmallVectorImpl<unsigned> &permutedDims) const {
+  unsigned projectionStart =
+      getNumResults() < getNumInputs() ? getNumInputs() - getNumResults() : 0;
+  permutedDims.clear();
+  SmallVector<unsigned> broadcastDims;
+  permutedDims.resize(getNumResults(), 0);
+  llvm::SmallDenseSet<unsigned> dimFound;
+  for (auto idxAndExpr : llvm::enumerate(getResults())) {
+    unsigned resIdx = idxAndExpr.index();
+    AffineExpr expr = idxAndExpr.value();
+    if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
+      // Each result may be either a constant 0 (broadcast dimension).
+      if (constExpr.getValue() != 0)
+        return false;
+      broadcastDims.push_back(resIdx);
+    } else if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) {
+      if (dimExpr.getPosition() < projectionStart)
+        return false;
+      permutedDims[resIdx] = dimExpr.getPosition() - projectionStart;
+      dimFound.insert(dimExpr.getPosition() - projectionStart);
+    } else {
+      return false;
+    }
+  }
+  // Find a permuation for the broadcast dimension. Since they are broadcasted
+  // any valid permutation is acceptable. We just permute the dim into a slot
+  // without an existing dimension.
+  unsigned pos = 0;
+  for (auto dim : broadcastDims) {
+    while (dimFound.count(pos)) {
+      pos++;
+    }
+    permutedDims[dim] = pos++;
+  }
+  return true;
+}
+
 /// Returns an AffineMap representing a permutation.
 AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
                                        MLIRContext *context) {
   assert(!permutation.empty() &&
          "Cannot create permutation map from empty permutation vector");
   SmallVector<AffineExpr, 4> affExprs;
   for (auto index : permutation)
     affExprs.push_back(getAffineDimExpr(index, context));

mlir/test/Dialect/Vector/vector-transfer-lowering.mlir

 // CHECK-NEXT:    }
 
 #broadcast = affine_map<(d0, d1, d2, d3, d4) -> (d1, 0, 0, d4)>
 func @transfer_broadcasting_complex(%mem : memref<10x20x30x8x8xf32>, %i : index) -> vector<3x2x4x5xf32> {
   %cf0 = constant 0.0 : f32
   %res = vector.transfer_read %mem[%i, %i, %i, %i, %i], %cf0 {masked = [false, false, false, false], permutation_map = #broadcast} : memref<10x20x30x8x8xf32>, vector<3x2x4x5xf32>
   return %res : vector<3x2x4x5xf32>
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, 0, 0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (0, d1, 0, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d3, d1, 0, 0)>
+#map3 = affine_map<(d0, d1) -> (d1, d0, 0, 0)>
+#map4 = affine_map<(d0, d1) -> (0, d1, 0, d0)>
+#map5 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3, d0)>
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, 0, 0)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, 0, d3)>
+
+// CHECK-LABEL: func @transfer_read_permutations
+func @transfer_read_permutations(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?x?x?xf32>)
+    -> (vector<7x14x8x16xf32>, vector<7x14x8x16xf32>, vector<7x14x8x16xf32>,
+       vector<7x14x8x16xf32>, vector<7x14x8x16xf32>, vector<7x14x8x16xf32>) {
+// CHECK-DAG: %[[CF0:.*]] = constant 0.000000e+00 : f32
+// CHECK-DAG: %[[C0:.*]] = constant 0 : index
+  %cst = constant 0.000000e+00 : f32
+  %c0 = constant 0 : index
+  %0 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst {permutation_map = #map0} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: vector.transfer_read {{.*}} {permutation_map = #[[$MAP0]]} : memref<?x?x?x?xf32>, vector<14x7x8x16xf32>
+// CHECK: vector.transpose %{{.*}}, [1, 0, 2, 3] : vector<14x7x8x16xf32> to vector<7x14x8x16xf32>
+  %1 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst {permutation_map = #map1} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: vector.transfer_read {{.*}} {permutation_map = #[[$MAP0]]} : memref<?x?x?x?xf32>, vector<16x14x7x8xf32>
+// CHECK: vector.transpose %{{.*}}, [2, 1, 3, 0] : vector<16x14x7x8xf32> to vector<7x14x8x16xf32>
+  %2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst {masked = [false, false, true, false], permutation_map = #map2} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: vector.transfer_read {{.*}} {masked = [false, true, false], permutation_map = #[[$MAP1]]} : memref<?x?x?x?xf32>, vector<14x16x7xf32>
+// CHECK: vector.broadcast %{{.*}} : vector<14x16x7xf32> to vector<8x14x16x7xf32>
+// CHECK: vector.transpose %{{.*}}, [3, 1, 0, 2] : vector<8x14x16x7xf32> to vector<7x14x8x16xf32>
+  %3 = vector.transfer_read %arg0[%c0, %c0], %cst {permutation_map = #map3} : memref<?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: vector.transfer_read %{{.*}}[%[[C0]], %[[C0]]], %[[CF0]] : memref<?x?xf32>, vector<14x7xf32>
+// CHECK: vector.broadcast %{{.*}} : vector<14x7xf32> to vector<8x16x14x7xf32>
+// CHECK: vector.transpose %{{.*}}, [3, 2, 0, 1] : vector<8x16x14x7xf32> to vector<7x14x8x16xf32>
+  %4 = vector.transfer_read %arg0[%c0, %c0], %cst {permutation_map = #map4} : memref<?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: vector.transfer_read %{{.*}}[%[[C0]], %[[C0]]], %[[CF0]] : memref<?x?xf32>, vector<16x14xf32>
+// CHECK: vector.broadcast %{{.*}} : vector<16x14xf32> to vector<7x8x16x14xf32>
+// CHECK: vector.transpose %{{.*}}, [0, 3, 1, 2] : vector<7x8x16x14xf32> to vector<7x14x8x16xf32>
+  %5 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst {permutation_map = #map5} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]], %[[CF0]] : memref<?x?x?x?xf32>, vector<16x14x7x8xf32>
+// CHECK: vector.transpose %{{.*}}, [2, 1, 3, 0] : vector<16x14x7x8xf32> to vector<7x14x8x16xf32>
+  return %0, %1, %2, %3, %4, %5 : vector<7x14x8x16xf32>, vector<7x14x8x16xf32>,
+         vector<7x14x8x16xf32>, vector<7x14x8x16xf32>, vector<7x14x8x16xf32>,
+         vector<7x14x8x16xf32>
+}