[mlir][Vector] Support 0-D vectors in `CreateMaskOp`

[michalt]

Summary

The 0-D case gets lowered in almost the same way that the 1-D case does
in VectorCreateMaskOpConversion. I also had to slightly update the
verifier for the op to always require exactly 1 operand in the 0-D case.

Depends On D115220

[ftynse]

mlir/lib/Dialect/Vector/VectorOps.cpp

3963

Please add a test for this message.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

1016

Nit: please use full sentences in comments, i.e., start with a capital letter and end with a full stop.

Also, you can do rewriter.notifyMatchFailure to indicate the failure reason and have it propagated through the debug infra instead of putting it an comment.

[michalt]

Rebase and address comments

[michalt]

Thanks for the review! Could I ask you to also land this? (I don't have commit access)

[nicolasvasilache]

Rebase by applying patch.

mlir/include/mlir/Dialect/Vector/VectorOps.td

def Vector_ConstantMaskOp :
   let extraClassDeclaration = [{
     static StringRef getMaskDimSizesAttrName() { return "mask_dim_sizes"; }
   }];
   let assemblyFormat = "$mask_dim_sizes attr-dict `:` type(results)";
 }
 
 def Vector_CreateMaskOp :
   Vector_Op<"create_mask", [NoSideEffect]>,
-    Arguments<(ins Variadic<Index>:$operands)>, Results<(outs VectorOf<[I1]>)> {
+    Arguments<(ins Variadic<Index>:$operands)>,
+    Results<(outs VectorOfAnyRankOf<[I1]>)> {
   let summary = "creates a vector mask";
   let description = [{
     Creates and returns a vector mask where elements of the result vector
     are set to '0' or '1', based on whether the element indices are contained
     within a hyper-rectangular region specified by the operands. Specifically,
     each operand specifies a range [0, operand-value) for a unique dimension in
     the vector result. The conjunction of the operand ranges define a
     hyper-rectangular region within which elements values are set to 1

mlir/lib/Dialect/Vector/VectorOps.cpp

static LogicalResult verify(ConstantMaskOp &op) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // CreateMaskOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(CreateMaskOp op) {
+  auto vectorType = op.getResult().getType().cast<VectorType>();
   // Verify that an operand was specified for each result vector each dimension.
-  if (op.getNumOperands() !=
-      op.getResult().getType().cast<VectorType>().getRank())
+  if (vectorType.getRank() == 0) {
+    if (op->getNumOperands() != 1)
+      return op.emitOpError(
+          "must specify exactly one operand for 0-D create_mask");
+  } else if (op.getNumOperands() !=
+             op.getResult().getType().cast<VectorType>().getRank()) {
     return op.emitOpError(
         "must specify an operand for each result vector dimension");
+  }
   return success();
 }
 
 namespace {
 
 // Pattern to rewrite a CreateMaskOp with a ConstantMaskOp.
 class CreateMaskFolder final : public OpRewritePattern<CreateMaskOp> {
 public:

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 ///   %x = ....
 /// until a one-dimensional vector is reached.
 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>();
+    int64_t rank = dstType.getRank();
+    if (rank <= 1)
+      return rewriter.notifyMatchFailure(
+          op, "0-D and 1-D vectors are handled separately");
+
+    auto loc = op.getLoc();
     auto eltType = dstType.getElementType();
     int64_t dim = dstType.getDimSize(0);
-    int64_t rank = dstType.getRank();
     Value idx = op.getOperand(0);
 
-    if (rank == 1)
-      return failure(); // leave for lowering
-
     VectorType lowType =
         VectorType::get(dstType.getShape().drop_front(), eltType);
     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));
static Value createCastToIndexLike(PatternRewriter &rewriter, Location loc,
   if (targetIntegerType.getWidth() > valueIntegerType.getWidth())
     return rewriter.create<arith::ExtSIOp>(loc, targetIntegerType, value);
   return rewriter.create<arith::TruncIOp>(loc, targetIntegerType, value);
 }
 
 // Helper that returns a vector comparison that constructs a mask:
 //     mask = [0,1,..,n-1] + [o,o,..,o] < [b,b,..,b]
 //
+// If `dim == 0` then the result will be a 0-D vector.
+//
 // NOTE: The LLVM::GetActiveLaneMaskOp intrinsic would provide an alternative,
 //       much more compact, IR for this operation, but LLVM eventually
 //       generates more elaborate instructions for this intrinsic since it
 //       is very conservative on the boundary conditions.
 static Value buildVectorComparison(PatternRewriter &rewriter, Operation *op,
                                    bool indexOptimizations, int64_t dim,
                                    Value b, Value *off = nullptr) {
   auto loc = op->getLoc();
   // If we can assume all indices fit in 32-bit, we perform the vector
   // comparison in 32-bit to get a higher degree of SIMD parallelism.
   // Otherwise we perform the vector comparison using 64-bit indices.
-  Value indices;
-  Type idxType;
-  if (indexOptimizations) {
-    indices = rewriter.create<arith::ConstantOp>(
-        loc, rewriter.getI32VectorAttr(
-                 llvm::to_vector<4>(llvm::seq<int32_t>(0, dim))));
-    idxType = rewriter.getI32Type();
+  Type idxType =
+      indexOptimizations ? rewriter.getI32Type() : rewriter.getI64Type();
+  DenseIntElementsAttr indicesAttr;
+  if (dim == 0 && indexOptimizations) {
+    indicesAttr = DenseIntElementsAttr::get(
+        VectorType::get(ArrayRef<int64_t>{}, idxType), ArrayRef<int32_t>{0});
+  } else if (dim == 0) {
+    indicesAttr = DenseIntElementsAttr::get(
+        VectorType::get(ArrayRef<int64_t>{}, idxType), ArrayRef<int64_t>{0});
+  } else if (indexOptimizations) {
+    indicesAttr = rewriter.getI32VectorAttr(
+        llvm::to_vector<4>(llvm::seq<int32_t>(0, dim)));
   } else {
-    indices = rewriter.create<arith::ConstantOp>(
-        loc, rewriter.getI64VectorAttr(
-                 llvm::to_vector<4>(llvm::seq<int64_t>(0, dim))));
-    idxType = rewriter.getI64Type();
+    indicesAttr = rewriter.getI64VectorAttr(
+        llvm::to_vector<4>(llvm::seq<int64_t>(0, dim)));
   }
+  Value indices = rewriter.create<arith::ConstantOp>(loc, indicesAttr);
   // Add in an offset if requested.
   if (off) {
     Value o = createCastToIndexLike(rewriter, loc, idxType, *off);
     Value ov = rewriter.create<SplatOp>(loc, indices.getType(), o);
     indices = rewriter.create<arith::AddIOp>(loc, ov, indices);
   }
   // Construct the vector comparison.
   Value bound = createCastToIndexLike(rewriter, loc, idxType, b);
LogicalResult matchAndRewrite(ConcreteOp xferOp,
 
     return success();
   }
 
 private:
   const bool indexOptimizations;
 };
 
-/// Conversion pattern for a vector.create_mask (1-D only).
+/// Conversion pattern for a `vector.create_mask` (0-D and 1-D only).
 class VectorCreateMaskOpConversion
     : public OpRewritePattern<vector::CreateMaskOp> {
 public:
   explicit VectorCreateMaskOpConversion(MLIRContext *context,
                                         bool enableIndexOpt)
       : mlir::OpRewritePattern<vector::CreateMaskOp>(context),
         indexOptimizations(enableIndexOpt) {}
 
   LogicalResult matchAndRewrite(vector::CreateMaskOp op,
                                 PatternRewriter &rewriter) const override {
     auto dstType = op.getType();
     int64_t rank = dstType.getRank();
-    if (rank == 1) {
-      rewriter.replaceOp(
-          op, buildVectorComparison(rewriter, op, indexOptimizations,
-                                    dstType.getDimSize(0), op.getOperand(0)));
-      return success();
-    }
-    return failure();
-  }
+    if (rank > 1)
+      return failure();
+                                  rank == 0 ? 0 : dstType.getDimSize(0),
+                                  op.getOperand(0)));
 
 private:
   const bool indexOptimizations;
 };
 
 // Drop inner most contiguous unit dimensions from transfer_read operand.
 class DropInnerMostUnitDims : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;

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

 // CHECK: %[[VAL_2:.*]] = builtin.unrealized_conversion_cast %[[VAL_1]] : vector<4x4xi1> to !llvm.array<4 x vector<4xi1>>
 // CHECK: %[[VAL_3:.*]] = llvm.insertvalue %[[VAL_0]], %[[VAL_2]][0] : !llvm.array<4 x vector<4xi1>>
 // CHECK: %[[VAL_4:.*]] = llvm.insertvalue %[[VAL_0]], %[[VAL_3]][1] : !llvm.array<4 x vector<4xi1>>
 // CHECK: %[[VAL_5:.*]] = builtin.unrealized_conversion_cast %[[VAL_4]] : !llvm.array<4 x vector<4xi1>> to vector<4x4xi1>
 // CHECK: return %[[VAL_5]] : vector<4x4xi1>
 
 // -----
 
+func @create_mask_0d(%a : index) -> vector<i1> {
+  %v = vector.create_mask %a : vector<i1>
+  return %v: vector<i1>
+}
+
+// CHECK-LABEL: func @create_mask_0d
+// CHECK-SAME: %[[arg:.*]]: index
+// CHECK:  %[[indices:.*]] = arith.constant dense<0> : vector<i32>
+// CHECK:  %[[arg_i32:.*]] = arith.index_cast %[[arg]] : index to i32
+// CHECK:  %[[bounds:.*]] = splat %[[arg_i32]] : vector<i32>
+// CHECK:  %[[result:.*]] = arith.cmpi slt, %[[indices]], %[[bounds]] : vector<i32>
+// CHECK:  return %[[result]] : vector<i1>
+// -----
+
+func @create_mask_1d(%a : index) -> vector<4xi1> {
+  %v = vector.create_mask %a : vector<4xi1>
+  return %v: vector<4xi1>
+}
+
+// CHECK-LABEL: func @create_mask_1d
+// CHECK-SAME: %[[arg:.*]]: index
+// CHECK:  %[[indices:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xi32>
+// CHECK:  %[[arg_i32:.*]] = arith.index_cast %[[arg]] : index to i32
+// CHECK:  %[[bounds:.*]] = splat %[[arg_i32]] : vector<4xi32>
+// CHECK:  %[[result:.*]] = arith.cmpi slt, %[[indices]], %[[bounds]] : vector<4xi32>
+// CHECK:  return %[[result]] : vector<4xi1>
+
+// -----
+
 func @flat_transpose(%arg0: vector<16xf32>) -> vector<16xf32> {
   %0 = vector.flat_transpose %arg0 { rows = 4: i32, columns = 4: i32 }
      : vector<16xf32> -> vector<16xf32>
   return %0 : vector<16xf32>
 }
 
 // CHECK-LABEL: func @flat_transpose
 // CHECK-SAME:  %[[A:.*]]: vector<16xf32>

mlir/test/Dialect/Vector/invalid.mlir

 // expected-error@+1 {{invalid accumulator/result vector shape, expected: 'vector<3x2xf32>'}}
   %0 = vector.contract #contraction_trait %arg0, %arg1, %arg2
     : vector<2x1xf32>, vector<1x3xf32> into vector<2x3xf32>
   return %0 : vector<2x3xf32>
 }
 
 // -----
 
+func @create_mask_0d_no_operands() {
+  %c1 = arith.constant 1 : index
+  // expected-error@+1 {{must specify exactly one operand for 0-D create_mask}}
+  %0 = vector.create_mask : vector<i1>
+}
+
+// -----
+
+func @create_mask_0d_many_operands() {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  // expected-error@+1 {{must specify exactly one operand for 0-D create_mask}}
+  %0 = vector.create_mask %c1, %c2, %c3 : vector<i1>
+}
+
+// -----
+
 func @create_mask() {
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index
   // expected-error@+1 {{must specify an operand for each result vector dimension}}
   %0 = vector.create_mask %c3, %c2 : vector<4x3x7xi1>
 }
 
 

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

func @arith_cmpi_0d(%smaller : vector<i32>, %bigger : vector<i32>) {
 
   %2 = arith.cmpi eq, %smaller, %bigger : vector<i32>
   // CHECK: ( 0 )
   vector.print %2: vector<i1>
 
   return
 }
 
+func @create_mask_0d(%zero : index, %one : index) {
+  %zero_mask = vector.create_mask %zero : vector<i1>
+  // CHECK: ( 0 )
+  vector.print %zero_mask : vector<i1>
+
+  %one_mask = vector.create_mask %one : vector<i1>
+  // CHECK: ( 1 )
+  vector.print %one_mask : vector<i1>
+
+  return
+}
+
 func @entry() {
   %0 = arith.constant 42.0 : f32
   %1 = arith.constant dense<0.0> : vector<f32>
   %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
   call  @extract_element_0d(%2) : (vector<f32>) -> ()
 
   %3 = arith.constant dense<42.0> : vector<f32>
   call  @print_vector_0d(%3) : (vector<f32>) -> ()
 
   %4 = arith.constant 42.0 : f32
   call  @splat_0d(%4) : (f32) -> ()
   call  @broadcast_0d(%4) : (f32) -> ()
 
   call  @bitcast_0d() : () -> ()
 
   %smaller = arith.constant dense<42> : vector<i32>
   %bigger = arith.constant dense<4242> : vector<i32>
   call  @arith_cmpi_0d(%smaller, %bigger) : (vector<i32>, vector<i32>) -> ()
 
+  %zero_idx = arith.constant 0 : index
+  %one_idx = arith.constant 1 : index
+  call  @create_mask_0d(%zero_idx, %one_idx) : (index, index) -> ()
+
   return
 }