[mlir][vector] Add canonicalzation pattern for shape_cast(create_mask)

[awarzynski]

Summary

This is primarily to avoid trailing unit dims:

%1 = vector.create_mask %c1, %dim_0, %c1, %c1 : vector<1x4x1x1xi1>
%2 = vector.shape_cast %1 : vector<1x4x1x1xi1> to vector<1x4xi1>

becomes:

%1 = vector.create_mask %c1, %dim_0 : vector<1x4xi1>

[benmxwl-arm]

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

5287–5288

This would trim vector<1x[1]x1xi128> to vector<1xi128> (same issue as in trimLeadingOneDims())

[awarzynski]

Fix (and test) how scalable vectors ar treated, add support for constant masks

[dcaballe]

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

5283

Could you check if we have something like this already? We are doing something similar in quite a few places. If so, having a single utility function would be great! You can take a look at "DropUnitDims" and other patterns...

[awarzynski]

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

5283

You are probably referring to trimLeadingOneDims (that I am updating in https://reviews.llvm.org/D157993).

Note that file is called "VectorDropLeadUnitDim.cpp" (_lead unit dim_) and in here I am getting rid of the _trailing_ unit dim, hence not sharing things over there. Having said that, I am implementing more patterns like this one, so we may want to group them together somewhere. I guess no need to decide just now?

5287–5288

Thanks, good catch! That's one very good reason to always test "scalable" variants too.

[awarzynski]

Switch to more consistent initialisation

[benmxwl-arm]

LGTM! Just a minor nit/question:

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

5336–5337

If any of the trailing 1 dims are false, the whole mask will be false (which is a slightly different replacement).

5344–5346

nit: semantically, should this be a success if it drops zero dims? I believe right now this will successfully replace the create_mask with an identical one.

[dcaballe]

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

5283

SG. I was hoping for a utility function that we can use for any combination of trailing, leading and interleaved cases...

5289

shouldn't this be part of ShapeCastOp::getCanonicalizationPatterns and not part of the transpose op?

5336

Not sure we should aim for supporting the zero case... That should probably be optimized in a different way as it means that the mask is always false...

5336–5337

I think that should never happen because all the dimensions should be at least one, right?

5344–5346

Good catch! Wouldn't this create an infinite loop? :)

5355

nit: preincrement per coding standard

[c-rhodes]

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

5283

nit: a comment describing what this does particularly w.r.t. scalable would be helpful

[c-rhodes]

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

5287–5288

nit: separate variables for old/new doesn't seem necessary

5299

nit: oldScalableDims, although I think it would be clearer if this is just false, and newShape = 1 above.

5345

nit: use shapeOpResTy

5362–5364

nit: I think you can just use maskDimSizes?

5368

nit: use shapeOpResTy

mlir/test/Dialect/Vector/canonicalize.mlir

2234

unused (and below)

2235

dim_1? (and below)

[awarzynski]

Thanks for all the review, an update is coming shortly :)

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

5289

Sorry, forgot to move it when pushing to Phab.

5336–5337

Hm, we might be referring to different things here. Looking at this example:

%1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x[4]x1x1xi1>

As we are only looking at _the trailing unit dimensions_ (so the trailing 1x1 in the case above), the dimension will be exactly:

• 1 in the 1st iteration,
• 1 in the 2nd iteration.

However, the corresponding mask value could be either 0 or 1. As Diego points out ^^^, zero case should be handled differently elsewhere, so I will just remove this comment.

5344–5346

Good point, though I am struggling to create an example that would trigger this.

mlir/test/Dialect/Vector/canonicalize.mlir

2235

Let me use dim - that will be even less confusing.

[awarzynski]

• Rename FoldShapeCastCreateMask as ShapeCastCreateMaskFolderTrailingOneDim
• Update tests
• Add documetation
• Only support cases when the mask input value is 1

[c-rhodes]

LGTM

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

4882

nit: shapeOpResTy

[dcaballe]

Thanks for addressing the feedback :)

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

LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
     auto newAttr =
         DenseElementsAttr::get(llvm::cast<VectorType>(shapeCastOp.getType()),
                                dense.getSplatValue<Attribute>());
     rewriter.replaceOpWithNewOp<arith::ConstantOp>(shapeCastOp, newAttr);
     return success();
   }
 };
 
+/// Helper function that computes a new vector type based on the input vector
+/// type by removing the trailing one dims:
+///
+///   vector<4x1x1xi1> --> vector<4x1>
+///
+static VectorType trimTrailingOneDims(VectorType oldType) {
+  ArrayRef<int64_t> oldShape = oldType.getShape();
+  ArrayRef<int64_t> newShape = oldShape;
+
+  ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
+  ArrayRef<bool> newScalableDims = oldScalableDims;
+
+  while (!newShape.empty() && newShape.back() == 1 && !newScalableDims.back()) {
+    newShape = newShape.drop_back(1);
+    newScalableDims = newScalableDims.drop_back(1);
+  }
+
+  // Make sure we have at least 1 dimension.
+  // TODO: Add support for 0-D vectors.
+  if (newShape.empty()) {
+    newShape = oldShape.take_back();
+    newScalableDims = oldScalableDims.take_back();
+  }
+
+  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
+}
+
+/// Folds qualifying shape_cast(create_mask) into a new create_mask
+///
+/// Looks at `vector.shape_cast` Ops that simply "drop" the trailing unit
+/// dimension. If the input vector comes from `vector.create_mask` for which
+/// the corresponding mask input value is 1 (e.g. `%c1` below), then it is safe
+/// to fold shape_cast into create_mask.
+///
+/// BEFORE:
+///    %1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x[4]x1x1xi1>
+///    %2 = vector.shape_cast %1 : vector<1x[4]x1x1xi1> to vector<1x[4]xi1>
+/// AFTER:
+///    %0 = vector.create_mask %c1, %dim : vector<1x[4]xi1>
+class ShapeCastCreateMaskFolderTrailingOneDim final
+    : public OpRewritePattern<ShapeCastOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ShapeCastOp shapeOp,
+                                PatternRewriter &rewriter) const override {
+    Value shapeOpSrc = shapeOp->getOperand(0);
+    auto createMaskOp = shapeOpSrc.getDefiningOp<vector::CreateMaskOp>();
+    auto constantMaskOp = shapeOpSrc.getDefiningOp<vector::ConstantMaskOp>();
+    if (!createMaskOp && !constantMaskOp)
+      return failure();
+
+    VectorType shapeOpResTy = shapeOp.getResultVectorType();
+    VectorType shapeOpSrcTy = shapeOp.getSourceVectorType();
+
+    VectorType newVecType = trimTrailingOneDims(shapeOpSrcTy);
+    if (newVecType != shapeOpResTy)
+      return failure();
+
+    auto numDimsToDrop =
+        shapeOpSrcTy.getShape().size() - shapeOpResTy.getShape().size();
+
+    // No unit dims to drop
+    if (!numDimsToDrop)
+      return failure();
+
+    if (createMaskOp) {
+      auto maskOperands = createMaskOp.getOperands();
+      auto numMaskOperands = maskOperands.size();
+
+      // Check every mask dim size to see whether it can be dropped
+      for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
+           --i) {
+        auto constant = maskOperands[i].getDefiningOp<arith::ConstantIndexOp>();
+        if (!constant || (constant.value() != 1))
+          return failure();
+      }
+      SmallVector<Value> newMaskOperands =
+          maskOperands.drop_back(numDimsToDrop);
+
+      rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(shapeOp, shapeOpResTy,
+                                                        newMaskOperands);
+      return success();
+    }
+
+    if (constantMaskOp) {
+      auto maskDimSizes = constantMaskOp.getMaskDimSizes().getValue();
+      auto numMaskOperands = maskDimSizes.size();
+
+      // Check every mask dim size to see whether it can be dropped
+      for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
+           --i) {
+        if (cast<IntegerAttr>(maskDimSizes[i]).getValue() != 1)
+          return failure();
+      }
+
+      auto newMaskOperands = maskDimSizes.drop_back(numDimsToDrop);
+      ArrayAttr newMaskOperandsAttr = rewriter.getArrayAttr(newMaskOperands);
+
+      rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(shapeOp, shapeOpResTy,
+                                                          newMaskOperandsAttr);
+      return success();
+    }
+
+    return failure();
+  }
+};
+
 /// Pattern to rewrite a ShapeCast(Broadcast) -> Broadcast.
 /// This only applies when the shape of the broadcast source
 /// 1. is a suffix of the shape of the result (i.e. when broadcast without
 ///    reshape is expressive enough to capture the result in a single op), or
 /// 2. has the same element count as the shape cast result.
 class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
     return failure();
   }
 };
 
 } // namespace
 
 void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
-  results.add<ShapeCastConstantFolder, ShapeCastBroadcastFolder>(context);
+  results.add<ShapeCastConstantFolder, ShapeCastCreateMaskFolderTrailingOneDim,
+              ShapeCastBroadcastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // VectorBitCastOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult BitCastOp::verify() {
   auto sourceVectorType = getSourceVectorType();
LogicalResult matchAndRewrite(TransposeOp transpOp,
 
     rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(
         transpOp, transpOp.getResultVectorType(),
         ArrayAttr::get(transpOp.getContext(), newMaskDimSizes));
     return success();
   }
 };
 
 } // namespace
 
 void vector::TransposeOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results.add<FoldTransposeCreateMask, FoldTransposedScalarBroadcast,
               TransposeFolder, FoldTransposeSplat>(context);
 }
 
 void vector::TransposeOp::getTransp(SmallVectorImpl<int64_t> &results) {
   populateFromInt64AttrArray(getTransp(), results);
 }
 
 //===----------------------------------------------------------------------===//
 // ConstantMaskOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ConstantMaskOp::verify() {
   auto resultType = llvm::cast<VectorType>(getResult().getType());
   // Check the corner case of 0-D vectors first.
   if (resultType.getRank() == 0) {
     if (getMaskDimSizes().size() != 1)
       return emitError("array attr must have length 1 for 0-D vectors");
     auto dim = llvm::cast<IntegerAttr>(getMaskDimSizes()[0]).getInt();
     if (dim != 0 && dim != 1)
       return emitError("mask dim size must be either 0 or 1 for 0-D vectors");
LogicalResult ConstantMaskOp::verify() {
   bool anyZeros = llvm::is_contained(maskDimSizes, 0);
   bool allZeros = llvm::all_of(maskDimSizes, [](int64_t s) { return s == 0; });
   if (anyZeros && !allZeros)
     return emitOpError("expected all mask dim sizes to be zeros, "
                        "as a result of conjunction with zero mask dim");
   // Verify that if the mask type is scalable, dimensions should be zero because
   // constant scalable masks can only be defined for the "none set" or "all set"
   // cases, and there is no VLA way to define an "all set" case for
   // `vector.constant_mask`. In the future, a convention could be established
   // to decide if a specific dimension value could be considered as "all set".
   if (resultType.isScalable() &&
       llvm::cast<IntegerAttr>(getMaskDimSizes()[0]).getInt() != 0)
     return emitOpError("expected mask dim sizes for scalable masks to be 0");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // CreateMaskOp
 //===----------------------------------------------------------------------===//
 
 void CreateMaskOp::build(OpBuilder &builder, OperationState &result,
                          VectorType type,
                          ArrayRef<OpFoldResult> mixedOperands) {
   SmallVector<Value> operands =
       getValueOrCreateConstantIndexOp(builder, result.location, mixedOperands);
   build(builder, result, type, operands);
 }
 
 LogicalResult CreateMaskOp::verify() {
   auto vectorType = llvm::cast<VectorType>(getResult().getType());
   // Verify that an operand was specified for each result vector each dimension.
   if (vectorType.getRank() == 0) {
     if (getNumOperands() != 1)
       return emitOpError(
           "must specify exactly one operand for 0-D create_mask");
   } else if (getNumOperands() !=
              llvm::cast<VectorType>(getResult().getType()).getRank()) {
     return 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:
   using OpRewritePattern::OpRewritePattern;

mlir/test/Dialect/Vector/canonicalize.mlir

 func.func @all_true_vector_mask(%a : vector<3x4xf32>) -> vector<3x4xf32> {
 //   CHECK-NOT:   vector.mask
 //       CHECK:   %[[ADD:.*]] = arith.addf %[[IN]], %[[IN]] : vector<3x4xf32>
 //       CHECK:   return %[[ADD]] : vector<3x4xf32>
   %all_true = vector.constant_mask [3, 4] : vector<3x4xi1>
   %0 = vector.mask %all_true { arith.addf %a, %a : vector<3x4xf32> } : vector<3x4xi1> -> vector<3x4xf32>
   return %0 : vector<3x4xf32>
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask(
+// CHECK-SAME:     %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x4xi1> {
+func.func @fold_shape_cast_with_mask(%arg0: tensor<1x?xf32>) -> vector<1x4xi1> {
+// CHECK-NOT: vector.shape_cast
+// CHECK:     %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK:     %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
+// CHECK:     %[[VAL_3:.*]] = vector.create_mask %[[VAL_1]], %[[VAL_2]] : vector<1x4xi1>
+// CHECK:     return %[[VAL_3]] : vector<1x4xi1>
+  %c1 = arith.constant 1 : index
+  %dim = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+  %1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x4x1x1xi1>
+  %2 = vector.shape_cast %1 : vector<1x4x1x1xi1> to vector<1x4xi1>
+  return %2 : vector<1x4xi1>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask_scalable(
+// CHECK-SAME:    %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x[4]xi1> {
+func.func @fold_shape_cast_with_mask_scalable(%arg0: tensor<1x?xf32>) -> vector<1x[4]xi1> {
+// CHECK-NOT: vector.shape_cast
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
+// CHECK:           %[[VAL_3:.*]] = vector.create_mask %[[VAL_1]], %[[VAL_2]] : vector<1x[4]xi1>
+// CHECK:           return %[[VAL_3]] : vector<1x[4]xi1>
+  %c1 = arith.constant 1 : index
+  %dim = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+  %1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x[4]x1x1xi1>
+  %2 = vector.shape_cast %1 : vector<1x[4]x1x1xi1> to vector<1x[4]xi1>
+  return %2 : vector<1x[4]xi1>
+}
+
+// -----
+
+// Check that scalable "1" (i.e. [1]) is not folded
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask_scalable_one(
+// CHECK-SAME:    %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x[1]xi1> {
+func.func @fold_shape_cast_with_mask_scalable_one(%arg0: tensor<1x?xf32>) -> vector<1x[1]xi1>{
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
+// CHECK:           %[[VAL_3:.*]] = vector.create_mask %[[VAL_1]], %[[VAL_2]] : vector<1x[1]xi1>
+// CHECK:           return %[[VAL_3]] : vector<1x[1]xi1>
+  %c1 = arith.constant 1 : index
+  %dim = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+  %1 = vector.create_mask %c1, %dim, %c1 : vector<1x[1]x1xi1>
+  %2 = vector.shape_cast %1 : vector<1x[1]x1xi1> to vector<1x[1]xi1>
+  return %2 : vector<1x[1]xi1>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @fold_shape_cast_with_constant_mask() -> vector<4xi1> {
+func.func @fold_shape_cast_with_constant_mask() -> vector<4xi1>{
+// CHECK-NOT: vector.shape_cast
+// CHECK:           %[[VAL_0:.*]] = vector.constant_mask [1] : vector<4xi1>
+// CHECK:           return %[[VAL_0]] : vector<4xi1>
+  %1 = vector.constant_mask [1, 1, 1] : vector<4x1x1xi1>
+  %2 = vector.shape_cast %1 : vector<4x1x1xi1> to vector<4xi1>
+  return %2 : vector<4xi1>
+}