[mlir][vector] Fix bug in `TransferWriteNonPermutationLowering`

[springerm]

Summary

This pattern expands the rank of the vector. However, the rank of the mask was not expanded.

[nicolasvasilache]

nice, thanks for finding and fixing!

mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp

static Value extendVectorRank(OpBuilder &builder, Location loc, Value vec,
   SmallVector<int64_t> newShape(addedRank, 1);
   newShape.append(originalVecType.getShape().begin(),
                   originalVecType.getShape().end());
   VectorType newVecType =
       VectorType::get(newShape, originalVecType.getElementType());
   return builder.create<vector::BroadcastOp>(loc, newVecType, vec);
 }
 
+/// Extend the rank of a vector Value by `addedRanks` by adding inner unit
+/// dimensions.
+static Value extendMaskRank(OpBuilder &builder, Location loc, Value vec,
+                            int64_t addedRank) {
+  Value broadcasted = extendVectorRank(builder, loc, vec, addedRank);
+  SmallVector<int64_t> permutation;
+  for (int64_t i = addedRank,
+               e = broadcasted.getType().cast<VectorType>().getRank();
+       i < e; ++i)
+    permutation.push_back(i);
+  for (int64_t i = 0; i < addedRank; ++i)
+    permutation.push_back(i);
+  return builder.create<vector::TransposeOp>(loc, broadcasted, permutation);
+}
+
 //===----------------------------------------------------------------------===//
 // populateVectorTransferPermutationMapLoweringPatterns
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// 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.
for (size_t i = 0; i < foundDim.size(); i++) {
       }
       if (!foundFirstDim)
         continue;
       // Once we found one outer dimension existing in the map keep track of all
       // the missing dimensions after that.
       missingInnerDim.push_back(i);
       exprs.push_back(rewriter.getAffineDimExpr(i));
     }
-    // Add unit dims at the beginning of the shape.
+    // Vector: add unit dims at the beginning of the shape.
     Value newVec = extendVectorRank(rewriter, op.getLoc(), op.getVector(),
                                     missingInnerDim.size());
+    // Mask: add unit dims at the end of the shape.
+    Value newMask;
+    if (op.getMask())
+      newMask = extendMaskRank(rewriter, op.getLoc(), op.getMask(),
+                               missingInnerDim.size());
     exprs.append(map.getResults().begin(), map.getResults().end());
     AffineMap newMap =
         AffineMap::get(map.getNumDims(), 0, exprs, op.getContext());
     ArrayAttr newInBoundsAttr;
     if (op.getInBounds()) {
       // All the new dimensions added are inbound.
       SmallVector<bool> newInBoundsValues(missingInnerDim.size(), true);
       for (Attribute attr : op.getInBounds().value().getValue()) {
         newInBoundsValues.push_back(cast<BoolAttr>(attr).getValue());
       }
       newInBoundsAttr = rewriter.getBoolArrayAttr(newInBoundsValues);
     }
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         op, newVec, op.getSource(), op.getIndices(), AffineMapAttr::get(newMap),
-        op.getMask(), newInBoundsAttr);
+        newMask, newInBoundsAttr);
     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)

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

+// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @lower_permutation_with_mask(
+//       CHECK:   %[[vec:.*]] = arith.constant dense<-2.000000e+00> : vector<7x1xf32>
+//       CHECK:   %[[mask:.*]] = arith.constant dense<[true, false, true, false, true, true, true]> : vector<7xi1>
+//       CHECK:   %[[b:.*]] = vector.broadcast %[[mask]] : vector<7xi1> to vector<1x7xi1>
+//       CHECK:   %[[tp:.*]] = vector.transpose %[[b]], [1, 0] : vector<1x7xi1> to vector<7x1xi1>
+//       CHECK:   vector.transfer_write %[[vec]], %{{.*}}[%{{.*}}, %{{.*}}], %[[tp]] {in_bounds = [false, true]} : vector<7x1xf32>, memref<?x?xf32>
+func.func @lower_permutation_with_mask(%A : memref<?x?xf32>, %base1 : index,
+                                       %base2 : index) {
+  %fn1 = arith.constant -2.0 : f32
+  %vf0 = vector.splat %fn1 : vector<7xf32>
+  %mask = arith.constant dense<[1, 0, 1, 0, 1, 1, 1]> : vector<7xi1>
+  vector.transfer_write %vf0, %A[%base1, %base2], %mask
+    {permutation_map = affine_map<(d0, d1) -> (d0)>, in_bounds = [false]}
+    : vector<7xf32>, memref<?x?xf32>
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !transform.any_op):
+  %f = transform.structured.match ops{["func.func"]} in %module_op
+    : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns to %f {
+    transform.apply_patterns.vector.transfer_permutation_patterns
+  } : !transform.any_op
+}