[mlir] [VectorOps] Implement vector.constant_mask lowering to LLVM IR

[aartbik]

Summary

Makes this operation runnable on CPU by generating MLIR instructions
that are eventually folded into an LLVM IR constant for the mask.

[ftynse]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

1223

There is some mechanism a pattern can tell the conversion infra it creates operations that the same pattern can match again, but this iteration is finite. I suppose it's worth using even if the pattern seems to work as is.

[aartbik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

1223

Thanks, yes, I saw the recent developments on that (hasBoundedRewriteRecursion() right?). For the vector->vector progressive lowering, we use "applyPatternsAndFoldGreedily" so that it does not really matter, yet (other than of course making sure that the iteration terminates, for most cases by stripping of a rank or so). But it does not hurt to add, although probably in a new CL for all the rules in this file so far.

[aartbik]

adhered to camelCaseRestriction style

[nicolasvasilache]

Thanks Aart!

[ftynse]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

1223

Last time I looked at it applyPatternsAndFoldGreedily applied the same pattern no more than 10 times. We may want to remove or raise this bound if hasBoundedRewriteRecursion tells it's okay.

[aartbik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

1223

Ah, thanks for that background!

mlir/lib/Dialect/Vector/VectorTransforms.cpp

for (int64_t d = 0, e = resType.getDimSize(0); d < e; ++d) {
       }
       result = rewriter.create<vector::InsertOp>(loc, resType, m, result, pos);
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
+/// Progressive lowering of ConstantMaskOp.
+/// One:
+///   %x = vector.constant_mask_op [a,b]
+/// is replaced by:
+///   %z = zero-result
+///   %l = vector.constant_mask_op [b]
+///   %4 = vector.insert %l, %z[0]
+///   ..
+///   %x = vector.insert %l, %..[a-1]
+/// which will be folded at LLVM IR level.
+class ConstantMaskOpLowering : public OpRewritePattern<vector::ConstantMaskOp> {
+public:
+  using OpRewritePattern<vector::ConstantMaskOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ConstantMaskOp op,
+                                PatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto dstType = op.getResult().getType().cast<VectorType>();
+    auto eltType = dstType.getElementType();
+    auto dim_sizes = op.mask_dim_sizes();
+    int64_t rank = dim_sizes.size();
+    int64_t trueDim = dim_sizes[0].cast<IntegerAttr>().getInt();
+
+    Value trueVal;
+    if (rank == 1) {
+      trueVal = rewriter.create<ConstantOp>(
+          loc, eltType, rewriter.getIntegerAttr(eltType, 1));
+    } else {
+      VectorType lowType =
+          VectorType::get(dstType.getShape().drop_front(), eltType);
+      SmallVector<int64_t, 4> new_dim_sizes;
+      for (int64_t r = 1; r < rank; r++)
+        new_dim_sizes.push_back(dim_sizes[r].cast<IntegerAttr>().getInt());
+      trueVal = rewriter.create<vector::ConstantMaskOp>(
+          loc, lowType, rewriter.getI64ArrayAttr(new_dim_sizes));
+    }
+
+    Value result = rewriter.create<ConstantOp>(loc, dstType,
+                                               rewriter.getZeroAttr(dstType));
+    for (int64_t d = 0; d < trueDim; d++) {
+      auto pos = rewriter.getI64ArrayAttr(d);
+      result =
+          rewriter.create<vector::InsertOp>(loc, dstType, trueVal, result, pos);
+    }
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+
 /// Progressive lowering of ContractionOp.
 /// One:
 ///   %x = vector.contract with at least one free/batch dimension
 /// is replaced by:
 ///   %a = vector.contract with one less free/batch dimension
 ///   %b = vector.contract with one less free/batch dimension
 ///   ..
 ///   %x = combine %a %b ..
void mlir::vector::populateVectorSlicesLoweringPatterns(
   patterns.insert<ExtractSlicesOpLowering, InsertSlicesOpLowering>(context);
 }
 
 void mlir::vector::populateVectorContractLoweringPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context,
     VectorTransformsOptions parameters) {
   patterns.insert<ShapeCastOp2DDownCastRewritePattern,
                   ShapeCastOp2DUpCastRewritePattern, BroadcastOpLowering,
-                  TransposeOpLowering, OuterProductOpLowering>(context);
+                  TransposeOpLowering, OuterProductOpLowering,
+                  ConstantMaskOpLowering>(context);
   patterns.insert<ContractionOpLowering>(parameters, context);
 }

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

 //       CHECK: %[[DIM:.*]] = llvm.extractvalue %{{.*}}[3, 0] :
 //  CHECK-SAME: !llvm<"{ float addrspace(3)*, float addrspace(3)*, i64, [1 x i64], [1 x i64] }">
 //
 // 3. Check address apce for GEP is correct.
 //       CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :
 //  CHECK-SAME: (!llvm<"float addrspace(3)*">, !llvm.i64) -> !llvm<"float addrspace(3)*">
 //       CHECK: %[[vecPtr_b:.*]] = llvm.addrspacecast %[[gep_b]] :
 //  CHECK-SAME: !llvm<"float addrspace(3)*"> to !llvm<"<17 x float>*">
+
+func @genbool_1d() -> vector<8xi1> {
+  %0 = vector.constant_mask [4] : vector<8xi1>
+  return %0 : vector<8xi1>
+}
+// CHECK-LABEL: func @genbool_1d
+// CHECK: %[[T0:.*]] = llvm.mlir.constant(1 : i1) : !llvm.i1
+// CHECK: %[[T1:.*]] = llvm.mlir.constant(dense<false> : vector<8xi1>) : !llvm<"<8 x i1>">
+// CHECK: %[[T2:.*]] = llvm.mlir.constant(0 : i64) : !llvm.i64
+// CHECK: %[[T3:.*]] = llvm.insertelement %[[T0]], %[[T1]][%[[T2]] : !llvm.i64] : !llvm<"<8 x i1>">
+// CHECK: %[[T4:.*]] = llvm.mlir.constant(1 : i64) : !llvm.i64
+// CHECK: %[[T5:.*]] = llvm.insertelement %[[T0]], %[[T3]][%[[T4]] : !llvm.i64] : !llvm<"<8 x i1>">
+// CHECK: %[[T6:.*]] = llvm.mlir.constant(2 : i64) : !llvm.i64
+// CHECK: %[[T7:.*]] = llvm.insertelement %[[T0]], %[[T5]][%[[T6]] : !llvm.i64] : !llvm<"<8 x i1>">
+// CHECK: %[[T8:.*]] = llvm.mlir.constant(3 : i64) : !llvm.i64
+// CHECK: %[[T9:.*]] = llvm.insertelement %[[T0]], %[[T7]][%[[T8]] : !llvm.i64] : !llvm<"<8 x i1>">
+// CHECK: llvm.return %9 : !llvm<"<8 x i1>">

mlir/test/Dialect/Vector/vector-contract-transforms.mlir

 // CHECK:      %[[T22:.*]] = vector.insert %[[T19]], %[[T21]] [2] : vector<2xf32> into vector<3x2xf32>
 // CHECK:      %[[T23:.*]] = vector.insert %[[T22]], %[[T17]] [3] : vector<3x2xf32> into vector<4x3x2xf32>
 // CHECK:      return %[[T23]] : vector<4x3x2xf32>
 
 func @broadcast_stretch_in_middle(%arg0: vector<4x1x2xf32>) -> vector<4x3x2xf32> {
   %0 = vector.broadcast %arg0 : vector<4x1x2xf32> to vector<4x3x2xf32>
   return %0 : vector<4x3x2xf32>
 }
+
+// CHECK-LABEL: func @genbool_1d
+// CHECK: %[[TT:.*]] = constant 1 : i1
+// CHECK: %[[C1:.*]] = constant dense<false> : vector<8xi1>
+// CHECK: %[[T0.*]] = vector.insert %[[TT]], %[[C1]] [0] : i1 into vector<8xi1>
+// CHECK: %[[T1.*]] = vector.insert %[[TT]], %[[T0]] [1] : i1 into vector<8xi1>
+// CHECK: %[[T2.*]] = vector.insert %[[TT]], %[[T1]] [2] : i1 into vector<8xi1>
+// CHECK: %[[T3.*]] = vector.insert %[[TT]], %[[T2]] [3] : i1 into vector<8xi1>
+// CHECK: return %[[T3]] : vector<8xi1>
+
+func @genbool_1d() -> vector<8xi1> {
+  %0 = vector.constant_mask [4] : vector<8xi1>
+  return %0 : vector<8xi1>
+}
+
+// CHECK-LABEL: func @genbool_2d
+// CHECK: %[[TT:.*]] = constant 1 : i1
+// CHECK: %[[C1:.*]] = constant dense<false> : vector<4xi1>
+// CHECK: %[[C2:.*]] = constant dense<false> : vector<4x4xi1>
+// CHECK: %[[T0:.*]] = vector.insert %[[TT]], %[[C1]] [0] : i1 into vector<4xi1>
+// CHECK: %[[T1:.*]] = vector.insert %[[TT]], %[[T0]] [1] : i1 into vector<4xi1>
+// CHECK: %[[T2:.*]] = vector.insert %[[T1]], %[[C2]] [0] : vector<4xi1> into vector<4x4xi1>
+// CHECK: %[[T3:.*]] = vector.insert %[[T1]], %[[T2]] [1] : vector<4xi1> into vector<4x4xi1>
+// CHECK: return %[[T3]] : vector<4x4xi1>
+
+func @genbool_2d() -> vector<4x4xi1> {
+  %v = vector.constant_mask [2, 2] : vector<4x4xi1>
+  return %v: vector<4x4xi1>
+}
+
+// CHECK-LABEL: func @genbool_3d
+// CHECK: %[[Tt:.*]] = constant 1 : i1
+// CHECK: %[[C1:.*]] = constant dense<false> : vector<4xi1>
+// CHECK: %[[C2:.*]] = constant dense<false> : vector<3x4xi1>
+// CHECK: %[[C3:.*]] = constant dense<false> : vector<2x3x4xi1>
+// CHECK: %[[T0:.*]] = vector.insert %[[TT]], %[[C1]] [0] : i1 into vector<4xi1>
+// CHECK: %[[T1:.*]] = vector.insert %[[TT]], %[[T0]] [1] : i1 into vector<4xi1>
+// CHECK: %[[T2:.*]] = vector.insert %[[TT]], %[[T1]] [2] : i1 into vector<4xi1>
+// CHECK: %[[T3:.*]] = vector.insert %[[T2]], %[[C2]] [0] : vector<4xi1> into vector<3x4xi1>
+// CHECK: %[[T4:.*]] = vector.insert %[[T3]], %[[C3]] [0] : vector<3x4xi1> into vector<2x3x4xi1>
+// CHECK: return %[[T4]] : vector<2x3x4xi1>
+
+func @genbool_3d() -> vector<2x3x4xi1> {
+  %v = vector.constant_mask [1, 1, 3] : vector<2x3x4xi1>
+  return %v: vector<2x3x4xi1>
+}

mlir/test/Target/vector-to-llvm-ir.mlir

+// RUN: mlir-opt %s -convert-vector-to-llvm | mlir-translate -mlir-to-llvmir | FileCheck %s
+
+func @genbool_1d() -> vector<8xi1> {
+  %0 = vector.constant_mask [4] : vector<8xi1>
+  return %0 : vector<8xi1>
+}
+// CHECK-LABEL: @genbool_1d()
+// CHECK-NEXT: ret <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 false, i1 false, i1 false, i1 false>
+
+func @genbool_2d() -> vector<4x4xi1> {
+  %v = vector.constant_mask [2, 2] : vector<4x4xi1>
+  return %v: vector<4x4xi1>
+}
+// CHECK-LABEL: @genbool_2d()
+// CHECK-NEXT: ret [4 x <4 x i1>] [<4 x i1> <i1 true, i1 true, i1 false, i1 false>, <4 x i1> <i1 true, i1 true, i1 false, i1 false>, <4 x i1> zeroinitializer, <4 x i1> zeroinitializer]
+
+func @genbool_3d() -> vector<2x3x4xi1> {
+  %v = vector.constant_mask [1, 1, 3] : vector<2x3x4xi1>
+  return %v: vector<2x3x4xi1>
+}
+// CHECK-LABEL: @genbool_3d()
+// CHECK-NEXT: ret [2 x [3 x <4 x i1>]] {{\[+}}3 x <4 x i1>] [<4 x i1> <i1 true, i1 true, i1 true, i1 false>, <4 x i1> zeroinitializer, <4 x i1> zeroinitializer], [3 x <4 x i1>] zeroinitializer]
+// note: awkward syntax to match [[