[mlir][VectorOps] Don't drop scalable dims when lowering transfer_reads/writes (in VectorToLLVM)

[benmxwl-arm]

Summary

This is a follow-on to D158753, and allows the lowering of a
transfer read/write of n-D vectors with a single trailing scalable dimension
to primitive vector ops.

The final conversion to LLVM depends on D158517 and D158752, without
these patches type conversion will fail (or an assert is hit in the LLVM
backend) if the final IR contains an array of scalable vectors.

This patch adds transform.apply_patterns.vector.lower_create_mask
which allows the lowering of vector.create_mask/constant_mask to be
tested independently of --convert-vector-to-llvm.

[benmxwl-arm]

Slightly reworked the testing to use a new transform dialect op transform.apply_patterns.vector.lower_create_mask.
This allows testing the create_mask lowering independently of -convert-vector-to-llvm, which avoids hitting
TypeConverter asserts (until the other patches are merged).

[awarzynski]

In D159482#4641968, @benmxwl-arm wrote:

Slightly reworked the testing to use a new transform dialect op transform.apply_patterns.vector.lower_create_mask.
This allows testing the create_mask lowering independently of -convert-vector-to-llvm, which avoids hitting
TypeConverter asserts (until the other patches are merged).

I really like this approach, thanks! While this does change the scope of this patch, I think that this is very straightforward and non-controversial. LGTM

mlir/test/Dialect/Vector/vector-scalable-create-mask-lowering.mlir

2

[benmxwl-arm]

mlir/test/Dialect/Vector/vector-scalable-create-mask-lowering.mlir

2

--split-input-file does not really work here since it'd require duplicating the transform.sequence for both tests (unless there's some way to add a common section?)

[awarzynski]

mlir/test/Dialect/Vector/vector-scalable-create-mask-lowering.mlir

2

I was hoping that this would just work ™ :)

You can try this, but I'm not convinced that that would be better:

• https://github.com/llvm/llvm-project/blob/119c489cc1a026feafe67c2b58c34b40dc2b6bd0/mlir/test/Dialect/Transform/test-interpreter-external.mlir#L1-L2

I am happy for you to land this as is.

[benmxwl-arm]

mlir/test/Dialect/Vector/vector-scalable-create-mask-lowering.mlir

2

I decided to keep this as-is; I think it's more consistent with the other tests to have everything in one file :)

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

def ApplyLowerContractionPatternsOp : Op<Transform_Dialect,
   let arguments = (ins DefaultValuedAttr<VectorContractLoweringAttr,
       "vector::VectorContractLowering::OuterProduct">:$lowering_strategy
   );
   let assemblyFormat = [{
     (`lowering_strategy` `=` $lowering_strategy^)? attr-dict
   }];
 }
 
+def ApplyLowerCreateMaskPatternsOp : Op<Transform_Dialect,
+    "apply_patterns.vector.lower_create_mask",
+    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Indicates that vector create_mask-like operations should be lowered to
+    finer-grained vector primitives.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 def ApplyLowerMasksPatternsOp : Op<Transform_Dialect,
     "apply_patterns.vector.lower_masks",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Indicates that vector.create_mask and vector.constant_mask operations
     should be lowered to finer-grained vector primitives.
 
     This is usually a late step that is run after bufferization as part of the

mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

void transform::ApplyFoldArithExtensionPatternsOp::populatePatterns(
   vector::populateFoldArithExtensionPatterns(patterns);
 }
 
 void transform::ApplyVectorReductionToContractPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   vector::populateVectorReductionToContractPatterns(patterns);
 }
 
+void transform::ApplyLowerCreateMaskPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  vector::populateVectorMaskOpLoweringPatterns(patterns);
+}
+
 void transform::ApplyRankReducingSubviewPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   vector::populateVectorTransferDropUnitDimsPatterns(patterns);
 }
 
 void transform::ApplyTransferPermutationPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);

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

public:
   LogicalResult matchAndRewrite(vector::CreateMaskOp op,
                                 PatternRewriter &rewriter) const override {
     auto dstType = cast<VectorType>(op.getResult().getType());
     int64_t rank = dstType.getRank();
     if (rank <= 1)
       return rewriter.notifyMatchFailure(
           op, "0-D and 1-D vectors are handled separately");
 
+    if (dstType.getScalableDims().front())
+      return rewriter.notifyMatchFailure(
+          op, "Cannot unroll leading scalable dim in dstType");
+
     auto loc = op.getLoc();
-    auto eltType = dstType.getElementType();
     int64_t dim = dstType.getDimSize(0);
     Value idx = op.getOperand(0);
 
-    VectorType lowType =
-        VectorType::get(dstType.getShape().drop_front(), eltType);
+    VectorType lowType = VectorType::Builder(dstType).dropDim(0);
     Value trueVal = rewriter.create<vector::CreateMaskOp>(
         loc, lowType, op.getOperands().drop_front());
     Value falseVal = rewriter.create<arith::ConstantOp>(
         loc, lowType, rewriter.getZeroAttr(lowType));
     Value result = rewriter.create<arith::ConstantOp>(
         loc, dstType, rewriter.getZeroAttr(dstType));
     for (int64_t d = 0; d < dim; d++) {
       Value bnd =

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

LogicalResult matchAndRewrite(vector::TransferReadOp read,
 
     // If there is broadcasting involved then we first load the unbroadcasted
     // vector, and then broadcast it with `vector.broadcast`.
     ArrayRef<int64_t> vectorShape = read.getVectorType().getShape();
     SmallVector<int64_t> unbroadcastedVectorShape(vectorShape.begin(),
                                                   vectorShape.end());
     for (unsigned i : broadcastedDims)
       unbroadcastedVectorShape[i] = 1;
-    VectorType unbroadcastedVectorType = VectorType::get(
+    VectorType unbroadcastedVectorType = read.getVectorType().cloneWith(
         unbroadcastedVectorShape, read.getVectorType().getElementType());
 
     // `vector.load` supports vector types as memref's elements only when the
     // resulting vector type is the same as the element type.
     auto memrefElTy = memRefType.getElementType();
     if (isa<VectorType>(memrefElTy) && memrefElTy != unbroadcastedVectorType)
       return rewriter.notifyMatchFailure(read, "incompatible element type");
 

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

func.func @transfer_read_1d_mask(%A : memref<?xf32>, %base : index) -> vector<5xf32> {
   %m = arith.constant dense<[0, 0, 1, 0, 1]> : vector<5xi1>
   %f7 = arith.constant 7.0: f32
   %f = vector.transfer_read %A[%base], %f7, %m : memref<?xf32>, vector<5xf32>
   return %f: vector<5xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @transfer_read_1d_scalable_mask
+// CHECK: %[[passtru:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf32>
+// CHECK: %[[r:.*]] = llvm.intr.masked.load %{{.*}}, %{{.*}}, %[[passtru]] {alignment = 4 : i32} : (!llvm.ptr, vector<[4]xi1>, vector<[4]xf32>) -> vector<[4]xf32>
+// CHECK: return %[[r]] : vector<[4]xf32>
+func.func @transfer_read_1d_scalable_mask(%arg0: memref<1x?xf32>, %mask: vector<[4]xi1>) -> vector<[4]xf32> {
+  %c0 = arith.constant 0 : index
+  %pad = arith.constant 0.0 : f32
+  %vec = vector.transfer_read %arg0[%c0, %c0], %pad, %mask {in_bounds = [true]} : memref<1x?xf32>, vector<[4]xf32>
+  return %vec : vector<[4]xf32>
+}
+
+// -----
+// CHECK-LABEL: func @transfer_write_1d_scalable_mask
+// CHECK: llvm.intr.masked.store %{{.*}}, %{{.*}}, %{{.*}} {alignment = 4 : i32} : vector<[4]xf32>, vector<[4]xi1> into !llvm.ptr
+func.func @transfer_write_1d_scalable_mask(%arg0: memref<1x?xf32>, %vec: vector<[4]xf32>, %mask: vector<[4]xi1>) {
+  %c0 = arith.constant 0 : index
+  vector.transfer_write %vec, %arg0[%c0, %c0], %mask {in_bounds = [true]} : vector<[4]xf32>, memref<1x?xf32>
+  return
+}
+
+// -----
+
 func.func @genbool_0d_f() -> vector<i1> {
   %0 = vector.constant_mask [0] : vector<i1>
   return %0 : vector<i1>
 }
 // CHECK-LABEL: func @genbool_0d_f
 // CHECK: %[[VAL_0:.*]] = arith.constant dense<false> : vector<i1>
 // CHECK: return %[[VAL_0]] : vector<i1>
 

mlir/test/Dialect/Vector/vector-scalable-create-mask-lowering.mlir

+// RUN: mlir-opt %s --test-transform-dialect-interpreter | FileCheck %s
+
+// CHECK-LABEL: func.func @create_mask_2d_trailing_scalable(
+// CHECK-SAME: %[[arg:.*]]: index) -> vector<3x[4]xi1> {
+// CHECK-NEXT: %[[zero_mask_1d:.*]] = arith.constant dense<false> : vector<[4]xi1>
+// CHECK-NEXT: %[[zero_mask_2d:.*]] = arith.constant dense<false> : vector<3x[4]xi1>
+// CHECK-NEXT: %[[create_mask_1d:.*]] = vector.create_mask %[[arg]] : vector<[4]xi1>
+// CHECK-NEXT: %[[res_0:.*]] = vector.insert %[[create_mask_1d]], %[[zero_mask_2d]] [0] : vector<[4]xi1> into vector<3x[4]xi1>
+// CHECK-NEXT: %[[res_1:.*]] = vector.insert %[[create_mask_1d]], %[[res_0]] [1] : vector<[4]xi1> into vector<3x[4]xi1>
+// CHECK-NEXT: %[[res_2:.*]] = vector.insert %[[zero_mask_1d]], %[[res_1]] [2] : vector<[4]xi1> into vector<3x[4]xi1>
+// CHECK-NEXT: return %[[res_2]] : vector<3x[4]xi1>
+func.func @create_mask_2d_trailing_scalable(%a: index) -> vector<3x[4]xi1> {
+  %c2 = arith.constant 2 : index
+  %mask = vector.create_mask %c2, %a : vector<3x[4]xi1>
+  return %mask : vector<3x[4]xi1>
+}
+
+// -----
+
+/// The following cannot be lowered as the current lowering requires unrolling
+/// the leading dim.
+
+// CHECK-LABEL: func.func @cannot_create_mask_2d_leading_scalable(
+// CHECK-SAME: %[[arg:.*]]: index) -> vector<[4]x4xi1> {
+// CHECK: %{{.*}} = vector.create_mask %[[arg]], %{{.*}} : vector<[4]x4xi1>
+func.func @cannot_create_mask_2d_leading_scalable(%a: index) -> vector<[4]x4xi1> {
+  %c1 = arith.constant 1 : index
+  %mask = vector.create_mask %a, %c1 : vector<[4]x4xi1>
+  return %mask : vector<[4]x4xi1>
+}
+
+transform.sequence failures(suppress) {
+^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.lower_create_mask
+  } : !transform.any_op
+}