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

[michalt]

Summary

To support creating both a mask with just a single true and false values,
I had to relax the restriction in the verifier that the rank is always equal to
the length of the attribute array, in other words, we now allow:

• vector.constant_mask [0] : vector<i1> which gets lowered to arith.constant dense<false> : vector<i1>
• vector.constant_mask [1] : vector<i1> which gets lowered to arith.constant dense<true> : vector<i1>

(the attribute list for the 0-D case must be a singleton containing
either 0 or 1)

[nicolasvasilache]

LGTM, modulo the extension of BuiltinTypes that I would just drop unless @rriddle wants it in.
Accepting conditioned on that being dropped or River weighing in positively, whichever comes first.

mlir/include/mlir/IR/Builders.h

124 ↗

(On Diff #391549)

I think this is too confusing to put in Builders.h and I would just avoid it.
It is also inconsistent, al APIs would have to be updated.

At the very minimum I'd spell it getZeroDimBoolVectorAttr(bool value) but I think @rriddle wouldn't buy it, never hurts to ask though.

Spelling it explicitly at use the corner-case use sites is better IMO.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

970

Would this work ?

bool value = ... ;
 rewriter.replaceOpWithNewOp<arith::ConstantOp>(
op, dstType, VectorType::get(ArrayRef<int64_t>{}, rewriter.getI1Type()), value));

(not sure whether bool -> ArrayRef<bool> is automatic, you may need ArrayRef<bool>{value})

[michalt]

Address review comments

mlir/include/mlir/IR/Builders.h

124 ↗

(On Diff #391549)

Yeah, I agree: creating it at the use site looks better. Thanks! :)

mlir/lib/Dialect/Vector/VectorTransforms.cpp

970

That looks way better too :D
(but I did need the ArrayRef<bool> to make it work)

[nicolasvasilache]

Thanks for your contribution @michalt !

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

def Vector_TypeCastOp :
   let assemblyFormat = [{
     $memref attr-dict `:` type($memref) `to` type($result)
   }];
 }
 
 def Vector_ConstantMaskOp :
   Vector_Op<"constant_mask", [NoSideEffect]>,
     Arguments<(ins I64ArrayAttr:$mask_dim_sizes)>,
-    Results<(outs VectorOf<[I1]>)> {
+    Results<(outs VectorOfAnyRankOf<[I1]>)> {
   let summary = "creates a constant 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 'mask_dim_sizes'
     array attribute argument. Each element of the 'mask_dim_sizes' array,
     specifies an exclusive upper bound [0, mask-dim-size-element-value)
     for a unique dimension in the vector result. The conjunction of the ranges

mlir/lib/Dialect/Vector/VectorOps.cpp

void vector::TransposeOp::getTransp(SmallVectorImpl<int64_t> &results) {
   populateFromInt64AttrArray(transp(), results);
 }
 
 //===----------------------------------------------------------------------===//
 // ConstantMaskOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ConstantMaskOp &op) {
-  // Verify that array attr size matches the rank of the vector result.
   auto resultType = op.getResult().getType().cast<VectorType>();
+  // Check the corner case of 0-D vectors first.
+  if (resultType.getRank() == 0) {
+    if (op.mask_dim_sizes().size() != 1)
+      return op->emitError("array attr must have length 1 for 0-D vectors");
+    auto dim = op.mask_dim_sizes()[0].cast<IntegerAttr>().getInt();
+    if (dim != 0 && dim != 1)
+      return op->emitError(
+          "mask dim size must be either 0 or 1 for 0-D vectors");
+    return success();
+  }
+
+  // Verify that array attr size matches the rank of the vector result.
   if (static_cast<int64_t>(op.mask_dim_sizes().size()) != resultType.getRank())
     return op.emitOpError(
         "must specify array attr of size equal vector result rank");
   // Verify that each array attr element is in bounds of corresponding vector
   // result dimension size.
   auto resultShape = resultType.getShape();
   SmallVector<int64_t, 4> maskDimSizes;
   for (auto it : llvm::enumerate(op.mask_dim_sizes())) {

mlir/lib/Dialect/Vector/VectorTransforms.cpp

public:
   using OpRewritePattern<vector::ConstantMaskOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::ConstantMaskOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
     auto dstType = op.getType();
     auto eltType = dstType.getElementType();
     auto dimSizes = op.mask_dim_sizes();
-    int64_t rank = dimSizes.size();
+    int64_t rank = dstType.getRank();
+
+    if (rank == 0) {
+      assert(dimSizes.size() == 1 &&
+             "Expected exactly one dim size for a 0-D vector");
+      bool value = dimSizes[0].cast<IntegerAttr>().getInt() == 1;
+      rewriter.replaceOpWithNewOp<arith::ConstantOp>(
+          op, dstType,
+          DenseIntElementsAttr::get(
+              VectorType::get(ArrayRef<int64_t>{}, rewriter.getI1Type()),
+              ArrayRef<bool>{value}));
+      return success();
+    }
+
     int64_t trueDim = std::min(dstType.getDimSize(0),
                                dimSizes[0].cast<IntegerAttr>().getInt());
 
     if (rank == 1) {
       // Express constant 1-D case in explicit vector form:
       //   [T,..,T,F,..,F].
       SmallVector<bool, 4> values(dstType.getDimSize(0));
       for (int64_t d = 0; d < trueDim; d++)

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

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>
 }
 
 // -----
 
+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>
+
+// -----
+
+func @genbool_0d_t() -> vector<i1> {
+  %0 = vector.constant_mask [1] : vector<i1>
+  return %0 : vector<i1>
+}
+// CHECK-LABEL: func @genbool_0d_t
+// CHECK: %[[VAL_0:.*]] = arith.constant dense<true> : vector<i1>
+// CHECK: return %[[VAL_0]] : vector<i1>
+
+// -----
+
 func @genbool_1d() -> vector<8xi1> {
   %0 = vector.constant_mask [4] : vector<8xi1>
   return %0 : vector<8xi1>
 }
 // CHECK-LABEL: func @genbool_1d
 // CHECK: %[[VAL_0:.*]] = arith.constant dense<[true, true, true, true, false, false, false, false]> : vector<8xi1>
 // CHECK: return %[[VAL_0]] : vector<8xi1>
 

mlir/test/Dialect/Vector/invalid.mlir

func @create_mask() {
   %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>
 }
 
 
 // -----
 
+func @constant_mask_0d_no_attr() {
+  // expected-error@+1 {{array attr must have length 1 for 0-D vectors}}
+  %0 = vector.constant_mask [] : vector<i1>
+}
+
+// -----
+
+func @constant_mask_0d_bad_attr() {
+  // expected-error@+1 {{mask dim size must be either 0 or 1 for 0-D vectors}}
+  %0 = vector.constant_mask [2] : vector<i1>
+}
+
+// -----
+
 func @constant_mask() {
   // expected-error@+1 {{must specify array attr of size equal vector result rank}}
   %0 = vector.constant_mask [3, 2, 7] : vector<4x3xi1>
 }
 
 // -----
 
 func @constant_mask_out_of_bounds() {

mlir/test/Dialect/Vector/ops.mlir

func @create_vector_mask() {
   // CHECK-NEXT: %[[C3:.*]] = arith.constant 3 : index
   %c3 = arith.constant 3 : index
   // CHECK-NEXT: vector.create_mask %[[C3]], %[[C2]] : vector<4x3xi1>
   %0 = vector.create_mask %c3, %c2 : vector<4x3xi1>
 
   return
 }
 
+// CHECK-LABEL: @constant_vector_mask_0d
+func @constant_vector_mask_0d() {
+  // CHECK: vector.constant_mask [0] : vector<i1>
+  %0 = vector.constant_mask [0] : vector<i1>
+  // CHECK: vector.constant_mask [1] : vector<i1>
+  %1 = vector.constant_mask [1] : vector<i1>
+  return
+}
+
 // CHECK-LABEL: @constant_vector_mask
 func @constant_vector_mask() {
   // CHECK: vector.constant_mask [3, 2] : vector<4x3xi1>
   %0 = vector.constant_mask [3, 2] : vector<4x3xi1>
   return
 }
 
 // CHECK-LABEL: @vector_print

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

func @bitcast_0d() {
   %4 = vector.extractelement %3[] : vector<f32>
   %5 = arith.bitcast %4 : f32 to i32
   // CHECK: 42
   vector.print %5: i32
   return
 }
 
 
+func @constant_mask_0d() {
+  %1 = vector.constant_mask [0] : vector<i1>
+  // CHECK: ( 0 )
+  vector.print %1: vector<i1>
+  %2 = vector.constant_mask [1] : vector<i1>
+  // CHECK: ( 1 )
+  vector.print %2: 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
+
+  // Warning: these must be called in their textual order of definition in the
+  // file to not mess up FileCheck.
   call  @splat_0d(%4) : (f32) -> ()
   call  @broadcast_0d(%4) : (f32) -> ()
-
   call  @bitcast_0d() : () -> ()
+  call  @constant_mask_0d() : () -> ()
 
   return
 }