Fixups

Author: MacDue

mlir/include/mlir/Dialect/Vector/IR/VectorOps.h

@@ -56,6 +56,9 @@ namespace detail {
 struct BitmaskEnumStorage;
 } // namespace detail
 
+/// Predefined constant_mask kinds.
+enum class ConstantMaskKind { AllFalse = 0, AllTrue };
+
 /// Default callback to build a region with a 'vector.yield' terminator with no
 /// arguments.
 void buildTerminatedBody(OpBuilder &builder, Location loc);

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

@@ -2362,6 +2362,11 @@ def Vector_ConstantMaskOp :
     ```
   }];
 
+  let builders = [
+    // Build with mixed static/dynamic operands.
+    OpBuilder<(ins "VectorType":$type, "ConstantMaskKind":$kind)>
+  ];
+
   let extraClassDeclaration = [{
     /// Return the result type of this op.
     VectorType getVectorType() {

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -5749,6 +5749,16 @@ void vector::TransposeOp::getCanonicalizationPatterns(
 // ConstantMaskOp
 //===----------------------------------------------------------------------===//
 
+void ConstantMaskOp::build(OpBuilder &builder, OperationState &result,
+                           VectorType type, ConstantMaskKind kind) {
+  assert(kind == ConstantMaskKind::AllTrue ||
+         kind == ConstantMaskKind::AllFalse);
+  build(builder, result, type,
+        kind == ConstantMaskKind::AllTrue
+            ? type.getShape()
+            : SmallVector<int64_t>(type.getRank(), 0));
+}
+
 LogicalResult ConstantMaskOp::verify() {
   auto resultType = llvm::cast<VectorType>(getResult().getType());
   // Check the corner case of 0-D vectors first.

mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp

@@ -17,7 +17,9 @@ using namespace mlir;
 using namespace mlir::vector;
 namespace {
 
-/// If `value` is a constant multiple of `vector.vscale` return the multiplier.
+/// If `value` is a constant multiple of `vector.vscale` (e.g. `%cst *
+/// vector.vscale`), return the multiplier (`%cst`). Otherwise, return
+/// `std::nullopt`.
 std::optional<int64_t> getConstantVscaleMultiplier(Value value) {
   if (value.getDefiningOp<vector::VectorScaleOp>())
     return 1;
@@ -78,17 +80,16 @@ LogicalResult resolveAllTrueCreateMaskOp(IRRewriter &rewriter,
     if (failed(dimLowerBoundSize))
       return failure();
     if (dimLowerBoundSize->scalable) {
-      // If the lower bound is scalable and < the mask dim size then this dim is
-      // not all-true.
+      // 1. The lower bound, LB, is scalable. If LB is < the mask dim size then
+      // this dim is not all-true.
       if (dimLowerBoundSize->baseSize < maskTypeDimSizes[i])
         return failure();
     } else {
-      // If the lower bound is a constant:
+      // 2. The lower bound, LB, is a constant.
       // - If the mask dim size is scalable then this dim is not all-true.
       if (maskTypeDimScalableFlags[i])
         return failure();
-      // - If the lower bound is < the _fixed-size_ mask dim size then this dim
-      // is not all-true.
+      // - If LB < the _fixed-size_ mask dim size then this dim is not all-true.
       if (dimLowerBoundSize->baseSize < maskTypeDimSizes[i])
         return failure();
     }
@@ -97,8 +98,8 @@ LogicalResult resolveAllTrueCreateMaskOp(IRRewriter &rewriter,
   // Replace createMaskOp with an all-true constant. This should result in the
   // mask being removed in most cases (as xfer ops + vector.mask have folds to
   // remove all-true masks).
-  auto allTrue = rewriter.create<arith::ConstantOp>(
-      createMaskOp.getLoc(), maskType, DenseElementsAttr::get(maskType, true));
+  auto allTrue = rewriter.create<vector::ConstantMaskOp>(
+      createMaskOp.getLoc(), maskType, ConstantMaskKind::AllTrue);
   rewriter.replaceAllUsesWith(createMaskOp, allTrue);
   return success();
 }

mlir/test/Dialect/Vector/eliminate-masks.mlir

@@ -1,9 +1,9 @@
-// RUN: mlir-opt %s -split-input-file -test-eliminate-vector-masks  | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -test-eliminate-vector-masks --split-input-file | FileCheck %s
 
 // This tests a general pattern the vectorizer tends to emit.
 
 // CHECK-LABEL: @eliminate_redundant_masks_through_insert_and_extracts
-// CHECK: %[[ALL_TRUE_MASK:.*]] = arith.constant dense<true> : vector<[4]xi1>
+// CHECK: %[[ALL_TRUE_MASK:.*]] = vector.constant_mask [4] : vector<[4]xi1>
 // CHECK: vector.transfer_read {{.*}} %[[ALL_TRUE_MASK]]
 // CHECK: vector.transfer_write {{.*}} %[[ALL_TRUE_MASK]]
 func.func @eliminate_redundant_masks_through_insert_and_extracts(%tensor: tensor<1x1000xf32>) {
@@ -40,7 +40,7 @@ func.func @eliminate_redundant_masks_through_insert_and_extracts(%tensor: tensor
 // -----
 
 // CHECK-LABEL: @negative_extract_slice_size_shrink
-// CHECK-NOT: arith.constant dense<true> : vector<[4]xi1>
+// CHECK-NOT: vector.constant_mask
 // CHECK: %[[MASK:.*]] = vector.create_mask
 // CHECK: "test.some_use"(%[[MASK]]) : (vector<[4]xi1>) -> ()
 func.func @negative_extract_slice_size_shrink(%tensor: tensor<1000xf32>) {
@@ -67,8 +67,25 @@ func.func @negative_extract_slice_size_shrink(%tensor: tensor<1000xf32>) {
 
 // -----
 
+// CHECK-LABEL: @trivially_all_true_case
+// CHECK: %[[ALL_TRUE_MASK:.*]] = vector.constant_mask [2, 4] : vector<2x[4]xi1>
+// CHECK: "test.some_use"(%[[ALL_TRUE_MASK]]) : (vector<2x[4]xi1>) -> ()
+func.func @trivially_all_true_case(%tensor: tensor<2x?xf32>)
+{
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+  %vscale = vector.vscale
+  %c4_vscale = arith.muli %vscale, %c4 : index
+  // Is found to be all true _without_ value bounds analysis.
+  %mask = vector.create_mask %c2, %c4_vscale : vector<2x[4]xi1>
+  "test.some_use"(%mask) : (vector<2x[4]xi1>) -> ()
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @negative_constant_dim_not_all_true
-// CHECK-NOT: arith.constant dense<true> : vector<2x[4]xi1>
+// CHECK-NOT: vector.constant_mask
 // CHECK: %[[MASK:.*]] = vector.create_mask
 // CHECK: "test.some_use"(%[[MASK]]) : (vector<2x[4]xi1>) -> ()
 func.func @negative_constant_dim_not_all_true()
@@ -87,7 +104,7 @@ func.func @negative_constant_dim_not_all_true()
 // -----
 
 // CHECK-LABEL: @negative_constant_vscale_multiple_not_all_true
-// CHECK-NOT: arith.constant dense<true> : vector<2x[4]xi1>
+// CHECK-NOT: vector.constant_mask
 // CHECK: %[[MASK:.*]] = vector.create_mask
 // CHECK: "test.some_use"(%[[MASK]]) : (vector<2x[4]xi1>) -> ()
 func.func @negative_constant_vscale_multiple_not_all_true() {
@@ -105,7 +122,7 @@ func.func @negative_constant_vscale_multiple_not_all_true() {
 // -----
 
 // CHECK-LABEL: @negative_value_bounds_fixed_dim_not_all_true
-// CHECK-NOT: arith.constant dense<true> : vector<3x[4]xi1>
+// CHECK-NOT: vector.constant_mask
 // CHECK: %[[MASK:.*]] = vector.create_mask
 // CHECK: "test.some_use"(%[[MASK]]) : (vector<3x[4]xi1>) -> ()
 func.func @negative_value_bounds_fixed_dim_not_all_true(%tensor: tensor<2x?xf32>)
@@ -114,7 +131,7 @@ func.func @negative_value_bounds_fixed_dim_not_all_true(%tensor: tensor<2x?xf32>
   %c4 = arith.constant 4 : index
   %vscale = vector.vscale
   %c4_vscale = arith.muli %vscale, %c4 : index
-  // This is _very_ simple but since tensor.dim is not a constant value bounds
+  // This is _very_ simple, but since tensor.dim is not a constant, value bounds
   // will be used to resolve it.
   %dim = tensor.dim %tensor, %c0 : tensor<2x?xf32>
   %mask = vector.create_mask %dim, %c4_vscale : vector<3x[4]xi1>
@@ -125,7 +142,7 @@ func.func @negative_value_bounds_fixed_dim_not_all_true(%tensor: tensor<2x?xf32>
 // -----
 
 // CHECK-LABEL: @negative_value_bounds_scalable_dim_not_all_true
-// CHECK-NOT: arith.constant dense<true> : vector<3x[4]xi1>
+// CHECK-NOT: vector.constant_mask
 // CHECK: %[[MASK:.*]] = vector.create_mask
 // CHECK: "test.some_use"(%[[MASK]]) : (vector<3x[4]xi1>) -> ()
 func.func @negative_value_bounds_scalable_dim_not_all_true(%tensor: tensor<2x100xf32>) {