[mlir][vector] Add missing support for scalable vectors

[awarzynski]

Summary

This patch adds the missing logic so that the
TransferReadPermutationLowering can be used for scalable vectors. To
this end:

• TransferOp custom C++ builder is updated to support scalable vectors,
• TransferOpReduceRank is also updated to support scalable vectors.

This pattern is relevant when lowering linalg.matmul via
vector_multi_reduction for scalable vectors.

Depends on D157092

[dcaballe]

Thanks!

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

4953

nit: move ub to variable

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 //===----------------------------------------------------------------------===//
 // TransposeOp
 //===----------------------------------------------------------------------===//
 
 void vector::TransposeOp::build(OpBuilder &builder, OperationState &result,
                                 Value vector, ArrayRef<int64_t> transp) {
   VectorType vt = llvm::cast<VectorType>(vector.getType());
   SmallVector<int64_t, 4> transposedShape(vt.getRank());
-  for (unsigned i = 0; i < transp.size(); ++i)
+  SmallVector<bool, 4> transposedScalableDims(vt.getRank());
+  for (unsigned i = 0; i < transp.size(); ++i) {
     transposedShape[i] = vt.getShape()[transp[i]];
+    transposedScalableDims[i] = vt.getScalableDims()[transp[i]];
+  }
 
   result.addOperands(vector);
-  result.addTypes(VectorType::get(transposedShape, vt.getElementType()));
+  result.addTypes(VectorType::get(transposedShape, vt.getElementType(),
+                                  transposedScalableDims));
   result.addAttribute(TransposeOp::getTranspAttrName(result.name),
                       builder.getI64ArrayAttr(transp));
 }
 
 OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
   // Eliminate splat constant transpose ops.
   if (auto attr =
           llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getVector()))

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

LogicalResult matchAndRewrite(vector::TransferReadOp op,
 
     permutationMap = map.getPermutationMap(permutation, op.getContext());
     // 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());
+    ArrayRef<bool> originalScalableDims = op.getVectorType().getScalableDims();
+    SmallVector<bool> newScalableDims(originalShape.size());
     for (const auto &pos : llvm::enumerate(permutation)) {
       newVectorShape[pos.value()] = originalShape[pos.index()];
+      newScalableDims[pos.value()] = originalScalableDims[pos.index()];
     }
 
     // Transpose in_bounds attribute.
     ArrayAttr newInBoundsAttr =
         op.getInBounds() ? inverseTransposeInBoundsAttr(
                                rewriter, op.getInBounds().value(), permutation)
                          : ArrayAttr();
 
     // Generate new transfer_read operation.
-    VectorType newReadType =
-        VectorType::get(newVectorShape, op.getVectorType().getElementType());
+    VectorType newReadType = VectorType::get(
+        newVectorShape, op.getVectorType().getElementType(), newScalableDims);
     Value newRead = rewriter.create<vector::TransferReadOp>(
         op.getLoc(), newReadType, op.getSource(), op.getIndices(),
         AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
         newInBoundsAttr);
 
     // Transpose result of transfer_read.
     SmallVector<int64_t> transposePerm(permutation.begin(), permutation.end());
     rewriter.replaceOpWithNewOp<vector::TransposeOp>(op, newRead,
if (reducedShapeRank == 0) {
       }
       rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, originalVecType,
                                                        newRead);
       return success();
     }
 
     SmallVector<int64_t> newShape = llvm::to_vector<4>(
         originalVecType.getShape().take_back(reducedShapeRank));
+    auto newScalableDims =
+        originalVecType.getScalableDims().take_back(reducedShapeRank);
     // Vector rank cannot be zero. Handled by TransferReadToVectorLoadLowering.
     if (newShape.empty())
       return rewriter.notifyMatchFailure(op, "rank-reduced vector is 0-d");
 
-    VectorType newReadType =
-        VectorType::get(newShape, originalVecType.getElementType());
+    VectorType newReadType = VectorType::get(
+        newShape, originalVecType.getElementType(), newScalableDims);
     ArrayAttr newInBoundsAttr =
         op.getInBounds()
             ? rewriter.getArrayAttr(
                   op.getInBoundsAttr().getValue().take_back(reducedShapeRank))
             : ArrayAttr();
     Value newRead = rewriter.create<vector::TransferReadOp>(
         op.getLoc(), newReadType, op.getSource(), op.getIndices(),
         AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),

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-LABEL: func @lower_permutation_with_mask_fixed_width(
 //       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,
+func.func @lower_permutation_with_mask_fixed_width(%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
 }
 
+// CHECK-LABEL:   func.func @permutation_with_mask_scalable(
+// CHECK-SAME:      %[[ARG_0:.*]]: memref<?x?xf32>,
+// CHECK-SAME:      %[[IDX_1:.*]]: index,
+// CHECK-SAME:      %[[IDX_2:.*]]: index) -> vector<8x[4]x2xf32> {
+// CHECK:           %[[C0:.*]] = arith.constant 0 : index
+// CHECK:           %[[PASS_THROUGH:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[MASK:.*]] = vector.create_mask %[[IDX_2]], %[[IDX_1]] : vector<2x[4]xi1>
+// CHECK:           %[[T_READ:.*]] = vector.transfer_read %[[ARG_0]]{{\[}}%[[C0]], %[[C0]]], %[[PASS_THROUGH]], %[[MASK]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x[4]xf32>
+// CHECK:           %[[BCAST:.*]] = vector.broadcast %[[T_READ]] : vector<2x[4]xf32> to vector<8x2x[4]xf32>
+// CHECK:           %[[TRANSPOSE:.*]] = vector.transpose %[[BCAST]], [0, 2, 1] : vector<8x2x[4]xf32> to vector<8x[4]x2xf32>
+// CHECK:           return %[[TRANSPOSE]] : vector<8x[4]x2xf32>
+// CHECK:         }
+func.func @permutation_with_mask_scalable(%2: memref<?x?xf32>, %dim_1: index, %dim_2: index) -> (vector<8x[4]x2xf32>) {
+
+  %c0 = arith.constant 0 : index
+  %cst_0 = arith.constant 0.000000e+00 : f32
+
+  %mask = vector.create_mask %dim_2, %dim_1 : vector<2x[4]xi1>
+  %1 = vector.transfer_read %2[%c0, %c0], %cst_0, %mask 
+    {in_bounds = [true, true, true], permutation_map = affine_map<(d0, d1) -> (0, d1, d0)>}
+    : memref<?x?xf32>, vector<8x[4]x2xf32>
+  return %1 : vector<8x[4]x2xf32>
+}
+
 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
 }