[mlir] Support masks in TransferOpReduceRank and TransferReadPermutationLowering

[springerm]

Summary

These two patterns allow for more efficient codegen in VectorToSCF.

Depends On D101745

[springerm]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3077–3079

I don't fully understand what this is doing. If there's a better way to implement the mask handling below (maybe reusing some of these functions), please let me know.

[springerm]

changed comment

[springerm]

minor change

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3077–3079

You should be able to directly permutation for the transpose. (you'll have to convert it to int64_t unfortunately as affineMap take unsigned but transpose take int64_t, this is something we should fix)

What this isPermutationOfMinorIdentityWithBroadcasting does is it picks a permutation to get to minor identity with broadcast. If there is a broadcast in the map there may be several potential permutation but the right thing to do is to apply the same transpose to masks and dimensions.

In the code we inverse the permutation map to apply it to the result of the transfer read as we want to convert it back to the original order but for masks you just want to directly apply the permutation.

3167

Don't we need to remove the leading dimension of the mask so that it matches the new rank of the transfer read? We should also add a test for this case.

[springerm]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3077–3079

Thanks for explaining.

I cannot use permutation directly, because mask ignores all broadcast dimensions. E.g., a 4D transfer read op would have a 1D mask if 3 dimensions are broadcasted. Therefore, I would somehow have to remove source dimensions from permutation and "compress" (re-index) the map.

Note, I added support for transfer op masks recently, and the commit handling broadcasts (D101745) is not submitted yet. So if the "vector type shape != mask type shape" thing is surprising, this is something we could still discuss in that commit.

3167

Since broadcast dims do not have a corresponding dimension in the mask vector, removing a broadcast does not require any changes to the mask vector.

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3077–3079

I see, I missed the part that mask skips the broadcast dimension. In this case it makes sense, I added some potential simplification below.

3095–3102

Can you skip all this by doing compressUnusedDims(map)? this will remove all the unused dimension then you can just loop through all the results and do

for (unsigned i = 0; i < map.getNumResults(); ++i) {
  if (auto expr = map.getResult(i).dyn_cast<AffineDimExpr>())
    maskTransposeIndices.push_back(expr.getPosition());
}

3099

should be cast<> instead of dyn_cast since you already know the type.

3167

Makes sense.

[springerm]

rebase

[springerm]

rebase

[nicolasvasilache]

Same as previous PRs, please add some minimal IR test.
Approving conditioned on that.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3095–3102

+1

[springerm]

address comments

[springerm]

Test case is included in mlir/test/Dialect/Vector/vector-transfer-lowering.mlir.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3095–3102

Nice, I didn't know that there's actually a compressUnusedDims helper function.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

struct TransferReadPermutationLowering
   using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferReadOp op,
                                 PatternRewriter &rewriter) const override {
     SmallVector<unsigned> permutation;
     AffineMap map = op.permutation_map();
     if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutation))
       return failure();
-
     AffineMap permutationMap =
         map.getPermutationMap(permutation, op.getContext());
     if (permutationMap.isIdentity())
       return failure();
-    if (op.mask())
-      return failure();
+
     // Caluclate the map of the new read by applying the inverse permutation.
     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(permutation)) {
       newVectorShape[pos.value()] = originalShape[pos.index()];
     }
+
+    Value newMask;
+    if (op.mask()) {
+      // Remove unused dims from the permutation map. E.g.:
+      // E.g.:  (d0, d1, d2, d3, d4, d5) -> (d5, 0, d3, 0, d2)
+      // comp = (d0, d1, d2) -> (d2, 0, d1, 0 d0)
+      auto comp = compressUnusedDims(map);
+      // Get positions of remaining result dims.
+      // E.g.:  (d0, d1, d2) -> (d2, 0, d1, 0 d0)
+      // maskTransposeIndices = [ 2,     1,    0]
+      SmallVector<int64_t> maskTransposeIndices;
+      for (unsigned i = 0; i < comp.getNumResults(); ++i) {
+        if (auto expr = comp.getResult(i).dyn_cast<AffineDimExpr>())
+          maskTransposeIndices.push_back(expr.getPosition());
+      }
+
+      newMask = rewriter.create<vector::TransposeOp>(op.getLoc(), op.mask(),
+                                                     maskTransposeIndices);
+    }
+
     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.in_bounds() ? *op.in_bounds() : ArrayAttr());
+        op.padding(), newMask, op.in_bounds() ? *op.in_bounds() : ArrayAttr());
+
     SmallVector<int64_t> transposePerm(permutation.begin(), permutation.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 {
-    if (op.mask())
-      return failure();
     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++;
     }
VectorType newReadType =
         VectorType::get(newShape, originalVecType.getElementType());
     ArrayAttr newInBounds =
         op.in_bounds()
             ? rewriter.getArrayAttr(
                   op.in_boundsAttr().getValue().take_back(reducedShapeRank))
             : ArrayAttr();
     Value newRead = rewriter.create<vector::TransferReadOp>(
         op.getLoc(), newReadType, op.source(), op.indices(), newMap,
-        op.padding(), newInBounds);
+        op.padding(), op.mask(), newInBounds);
     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.

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

 // 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
+  %m = constant 1 : i1
 
-  %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>
+  %mask0 = splat %m : vector<7x14xi1>
+  %0 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst, %mask0 {permutation_map = #map0} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: %[[MASK0:.*]] = vector.transpose {{.*}} : vector<7x14xi1> to vector<14x7xi1>
+// CHECK: vector.transfer_read {{.*}} %[[MASK0]] {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>
+  %mask1 = splat %m : vector<14x16xi1>
+  %1 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst, %mask1 {permutation_map = #map1} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: %[[MASK1:.*]] = vector.transpose {{.*}} : vector<14x16xi1> to vector<16x14xi1>
+// CHECK: vector.transfer_read {{.*}} %[[MASK1]] {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 {in_bounds = [true, true, false, true], permutation_map = #map2} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
-// CHECK: vector.transfer_read {{.*}} {in_bounds = [true, false, true], permutation_map = #[[$MAP1]]} : memref<?x?x?x?xf32>, vector<14x16x7xf32>
+  %mask2 = splat %m : vector<7x14xi1>
+  %2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst, %mask2 {in_bounds = [true, true, false, true], permutation_map = #map2} : memref<?x?x?x?xf32>, vector<7x14x8x16xf32>
+// CHECK: %[[MASK2:.*]] = vector.transpose {{.*}} : vector<7x14xi1> to vector<14x7xi1>
+// CHECK: vector.transfer_read {{.*}} %[[MASK2]] {in_bounds = [true, false, true], 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>