[mlir] [VectorOps] Improve vector.constant_mask lowering

[aartbik]

Summary

Use direct vector constants for the 1-D case. This approach
scales much better than generating elaborate insertion operations
that are eventually folded into a constant. We could of course
generalize the 1-D case to higher ranks, but this simplification
already helps in scaling some microbenchmarks that would formerly
crash on the intermediate IR length.

[reidtatge]

LGTM

[aartbik]

@mehdi_amini : my first use of "aart-land" to submit this :-)

mlir/include/mlir/IR/Builders.h

public:
   IntegerAttr getI64IntegerAttr(int64_t value);
   IntegerAttr getIndexAttr(int64_t value);
 
   /// Signed and unsigned integer attribute getters.
   IntegerAttr getSI32IntegerAttr(int32_t value);
   IntegerAttr getUI32IntegerAttr(uint32_t value);
 
   /// Vector-typed DenseIntElementsAttr getters. `values` must not be empty.
+  DenseIntElementsAttr getBoolVectorAttr(ArrayRef<bool> values);
   DenseIntElementsAttr getI32VectorAttr(ArrayRef<int32_t> values);
   DenseIntElementsAttr getI64VectorAttr(ArrayRef<int64_t> values);
 
   /// Tensor-typed DenseIntElementsAttr getters. `values` can be empty.
   /// These are generally preferable for representing general lists of integers
   /// as attributes.
   DenseIntElementsAttr getI32TensorAttr(ArrayRef<int32_t> values);
   DenseIntElementsAttr getI64TensorAttr(ArrayRef<int64_t> values);

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 /// One:
 ///   %x = vector.constant_mask [a,b]
 /// is replaced by:
 ///   %z = zero-result
 ///   %l = vector.constant_mask [b]
 ///   %4 = vector.insert %l, %z[0]
 ///   ..
 ///   %x = vector.insert %l, %..[a-1]
-/// which will be folded at LLVM IR level.
+/// until a one-dimensional vector is reached. All these operations
+/// 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 dimSizes = op.mask_dim_sizes();
     int64_t rank = dimSizes.size();
     int64_t trueDim = dimSizes[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> newDimSizes;
-      for (int64_t r = 1; r < rank; r++)
-        newDimSizes.push_back(dimSizes[r].cast<IntegerAttr>().getInt());
-      trueVal = rewriter.create<vector::ConstantMaskOp>(
-          loc, lowType, rewriter.getI64ArrayAttr(newDimSizes));
-    }
-
+      SmallVector<bool, 4> values(dstType.getDimSize(0));
+      for (int64_t d = 0; d < trueDim; d++)
+        values[d] = true;
+      rewriter.replaceOpWithNewOp<ConstantOp>(
+          op, dstType, rewriter.getBoolVectorAttr(values));
+      return success();
+    }
+
+    Value trueVal = rewriter.create<vector::ConstantMaskOp>(
     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);

mlir/lib/IR/Builders.cpp

 IntegerAttr Builder::getIndexAttr(int64_t value) {
   return IntegerAttr::get(getIndexType(), APInt(64, value));
 }
 
 IntegerAttr Builder::getI64IntegerAttr(int64_t value) {
   return IntegerAttr::get(getIntegerType(64), APInt(64, value));
 }
 
+DenseIntElementsAttr Builder::getBoolVectorAttr(ArrayRef<bool> values) {
+  return DenseIntElementsAttr::get(
+      VectorType::get(static_cast<int64_t>(values.size()), getI1Type()),
+      values);
+}
+
 DenseIntElementsAttr Builder::getI32VectorAttr(ArrayRef<int32_t> values) {
   return DenseIntElementsAttr::get(
       VectorType::get(static_cast<int64_t>(values.size()), getIntegerType(32)),
       values);
 }
 
 DenseIntElementsAttr Builder::getI64VectorAttr(ArrayRef<int64_t> values) {
   return DenseIntElementsAttr::get(

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

 // 2. Rewrite as a load.
 //       CHECK: %[[loaded:.*]] = llvm.load %[[vecPtr]] : !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(true) : !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>">
+// CHECK: %[[C1:.*]] = llvm.mlir.constant(dense<[true, true, true, true, false, false, false, false]> : vector<8xi1>) : !llvm<"<8 x i1>">
+// CHECK: llvm.return %[[C1]] : !llvm<"<8 x i1>">
+
+func @genbool_2d() -> vector<4x4xi1> {
+  %v = vector.constant_mask [2, 2] : vector<4x4xi1>
+  return %v: vector<4x4xi1>
+}
+
+// CHECK-LABEL: func @genbool_2d
+// CHECK: %[[C1:.*]] = llvm.mlir.constant(dense<[true, true, false, false]> : vector<4xi1>) : !llvm<"<4 x i1>">
+// CHECK: %[[C2:.*]] = llvm.mlir.constant(dense<false> : vector<4x4xi1>) : !llvm<"[4 x <4 x i1>]">
+// CHECK: %[[T0:.*]] = llvm.insertvalue %[[C1]], %[[C2]][0] : !llvm<"[4 x <4 x i1>]">
+// CHECK: %[[T1:.*]] = llvm.insertvalue %[[C1]], %[[T0]][1] : !llvm<"[4 x <4 x i1>]">
+// CHECK: llvm.return %[[T1]] : !llvm<"[4 x <4 x i1>]">
+
+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:.*]]: !llvm<"<16 x float>">
 // CHECK:       %[[T:.*]] = llvm.intr.matrix.transpose %[[A]]
 // CHECK-SAME:      {columns = 4 : i32, rows = 4 : i32} :
 // CHECK-SAME:      !llvm<"<16 x float>"> into !llvm<"<16 x float>">
 // CHECK:       llvm.return %[[T]] : !llvm<"<16 x float>">
-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>
-}

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

 // 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 true
-// 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>
+// CHECK: %[[T0:.*]] = constant dense<[true, true, true, true, false, false, false, false]> : vector<8xi1>
+// CHECK: return %[[T0]] : 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 true
-// CHECK: %[[C1:.*]] = constant dense<false> : vector<4xi1>
+// CHECK: %[[C1:.*]] = constant dense<[true, true, false, 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>
+// CHECK: %[[T0:.*]] = vector.insert %[[C1]], %[[C2]] [0] : vector<4xi1> into vector<4x4xi1>
+// CHECK: %[[T1:.*]] = vector.insert %[[C1]], %[[T0]] [1] : vector<4xi1> into vector<4x4xi1>
+// CHECK: return %[[T1]] : 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 true
-// CHECK: %[[C1:.*]] = constant dense<false> : vector<4xi1>
+// CHECK: %[[C1:.*]] = constant dense<[true, true, true, 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>
+// CHECK: %[[T0:.*]] = vector.insert %[[C1]], %[[C2]] [0] : vector<4xi1> into vector<3x4xi1>
+// CHECK: %[[T1:.*]] = vector.insert %[[T0]], %[[C3]] [0] : vector<3x4xi1> into vector<2x3x4xi1>
+// CHECK: return %[[T1]] : 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