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

[aartbik]

Summary

First, compact implementation of lowering to LLVM IR. A bit more
challenging than the constant mask due to the dynamic indices, of course.
I like to hear if there are more efficient ways of doing this in LLVM,
but this for now at least gives us a functional reference implementation.

[rriddle]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

1239

Can you add comments to what the expected output is given a mask input? Lowerings should really be documented.

[aartbik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

1239

Yes of course. Not sure how I missed that when writing this CL.... Fix coming up....

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 };
 } // namespace
 
 void LowerVectorToLLVMPass::runOnOperation() {
   // Perform progressive lowering of operations on slices and
   // all contraction operations. Also applies folding and DCE.
   {
     OwningRewritePatternList patterns;
+    populateVectorToVectorCanonicalizationPatterns(patterns, &getContext());
     populateVectorSlicesLoweringPatterns(patterns, &getContext());
     populateVectorContractLoweringPatterns(patterns, &getContext());
     applyPatternsAndFoldGreedily(getOperation(), patterns);
   }
 
   // Convert to the LLVM IR dialect.
   LLVMTypeConverter converter(&getContext());
   OwningRewritePatternList patterns;

mlir/lib/Dialect/Vector/VectorTransforms.cpp

for (int64_t d = 0; d < trueDim; d++) {
       result =
           rewriter.create<vector::InsertOp>(loc, dstType, trueVal, result, pos);
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
+class CreateMaskOpLowering : public OpRewritePattern<vector::CreateMaskOp> {
+public:
+  using OpRewritePattern<vector::CreateMaskOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::CreateMaskOp op,
+                                PatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto dstType = op.getResult().getType().cast<VectorType>();
+    auto eltType = dstType.getElementType();
+    int64_t rank = dstType.getRank();
+    Value idx = op.getOperand(0);
+
+    Value trueVal;
+    Value falseVal;
+    if (rank > 1) {
+      VectorType lowType =
+          VectorType::get(dstType.getShape().drop_front(), eltType);
+      trueVal = rewriter.create<vector::CreateMaskOp>(
+          loc, lowType, op.getOperands().drop_front());
+      falseVal = rewriter.create<ConstantOp>(loc, lowType,
+                                             rewriter.getZeroAttr(lowType));
+    }
+
+    Value result = rewriter.create<ConstantOp>(loc, dstType,
+                                               rewriter.getZeroAttr(dstType));
+    for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; d++) {
+      Value bnd = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(d));
+      Value val = rewriter.create<CmpIOp>(loc, CmpIPredicate::slt, bnd, idx);
+      if (rank > 1)
+        val = rewriter.create<SelectOp>(loc, val, trueVal, falseVal);
+      auto pos = rewriter.getI64ArrayAttr(d);
+      result =
+          rewriter.create<vector::InsertOp>(loc, dstType, val, 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::populateVectorContractLoweringPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context,
     VectorTransformsOptions parameters) {
   patterns.insert<ShapeCastOp2DDownCastRewritePattern,
                   ShapeCastOp2DUpCastRewritePattern, BroadcastOpLowering,
                   TransposeOpLowering, OuterProductOpLowering,
-                  ConstantMaskOpLowering>(context);
+                  ConstantMaskOpLowering, CreateMaskOpLowering>(context);
   patterns.insert<ContractionOpLowering>(parameters, context);
 }

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

 // 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>
 }
+
+// CHECK-LABEL: func @genbool_var_1d
+// CHECK-SAME: %[[A:.*0]]: index
+// CHECK-DAG:  %[[VF:.*]] = constant dense<false> : vector<3xi1>
+// CHECK-DAG:  %[[C0:.*]] = constant 0 : index
+// CHECK-DAG:  %[[C1:.*]] = constant 1 : index
+// CHECK-DAG:  %[[C2:.*]] = constant 2 : index
+// CHECK:      %[[T0:.*]] = cmpi "slt", %[[C0]], %[[A]] : index
+// CHECK:      %[[T1:.*]] = vector.insert %[[T0]], %[[VF]] [0] : i1 into vector<3xi1>
+// CHECK:      %[[T2:.*]] = cmpi "slt", %[[C1]], %[[A]] : index
+// CHECK:      %[[T3:.*]] = vector.insert %[[T2]], %[[T1]] [1] : i1 into vector<3xi1>
+// CHECK:      %[[T4:.*]] = cmpi "slt", %[[C2]], %[[A]] : index
+// CHECK:      %[[T5:.*]] = vector.insert %[[T4]], %[[T3]] [2] : i1 into vector<3xi1>
+// CHECK:      return %[[T5]] : vector<3xi1>
+
+func @genbool_var_1d(%arg0: index) -> vector<3xi1> {
+  %0 = vector.create_mask %arg0 : vector<3xi1>
+  return %0 : vector<3xi1>
+}
+
+// CHECK-LABEL: func @genbool_var_2d
+// CHECK-SAME: %[[A:.*0]]: index
+// CHECK-SAME: %[[B:.*1]]: index
+// CHECK-DAG:  %[[Z1:.*]] = constant dense<false> : vector<3xi1>
+// CHECK-DAG:  %[[Z2:.*]] = constant dense<false> : vector<2x3xi1>
+// CHECK-DAG:  %[[C0:.*]] = constant 0 : index
+// CHECK-DAG:  %[[C1:.*]] = constant 1 : index
+// CHECK-DAG:  %[[C2:.*]] = constant 2 : index
+// CHECK:      %[[T0:.*]] = cmpi "slt", %[[C0]], %[[B]] : index
+// CHECK:      %[[T1:.*]] = vector.insert %[[T0]], %[[Z1]] [0] : i1 into vector<3xi1>
+// CHECK:      %[[T2:.*]] = cmpi "slt", %[[C1]], %[[B]] : index
+// CHECK:      %[[T3:.*]] = vector.insert %[[T2]], %[[T1]] [1] : i1 into vector<3xi1>
+// CHECK:      %[[T4:.*]] = cmpi "slt", %[[C2]], %[[B]] : index
+// CHECK:      %[[T5:.*]] = vector.insert %[[T4]], %[[T3]] [2] : i1 into vector<3xi1>
+// CHECK:      %[[T6:.*]] = cmpi "slt", %[[C0]], %[[A]] : index
+// CHECK:      %[[T7:.*]] = select %[[T6]], %[[T5]], %[[Z1]] : vector<3xi1>
+// CHECK:      %[[T8:.*]] = vector.insert %7, %[[Z2]] [0] : vector<3xi1> into vector<2x3xi1>
+// CHECK:      %[[T9:.*]] = cmpi "slt", %[[C1]], %[[A]] : index
+// CHECK:      %[[T10:.*]] = select %[[T9]], %[[T5]], %[[Z1]] : vector<3xi1>
+// CHECK:      %[[T11:.*]] = vector.insert %[[T10]], %[[T8]] [1] : vector<3xi1> into vector<2x3xi1>
+// CHECK:      return %[[T11]] : vector<2x3xi1>
+
+func @genbool_var_2d(%arg0: index, %arg1: index) -> vector<2x3xi1> {
+  %0 = vector.create_mask %arg0, %arg1 : vector<2x3xi1>
+  return %0 : vector<2x3xi1>
+}

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

 
 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 [[
+
+func @genbool_1d_var_but_constant() -> vector<8xi1> {
+  %i = constant 0 : index
+  %v = vector.create_mask %i : vector<8xi1>
+  return %v : vector<8xi1>
+}
+// CHECK-LABEL: @genbool_1d_var_but_constant()
+// CHECK-NEXT: ret <8 x i1> zeroinitializer