[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

5143–5144

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

5139

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

5139

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?

5143–5144

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

5192–5193

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

5200–5202

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

5139

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

5192

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...

5192–5193

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

5200–5202

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

5211

nit: preincrement per coding standard

5237

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

[c-rhodes]

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

5139

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

[c-rhodes]

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

5143–5144

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

5155

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

5201

nit: use shapeOpResTy

5218–5220

nit: I think you can just use maskDimSizes?

5224

nit: use shapeOpResTy

mlir/test/Dialect/Vector/canonicalize.mlir

2211

unused (and below)

2212

dim_1? (and below)

[awarzynski]

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

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

5192–5193

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.

5200–5202

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

5237

Sorry, forgot to move it when pushing to Phab.

mlir/test/Dialect/Vector/canonicalize.mlir

2212

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

4860

nit: shapeOpResTy

[dcaballe]

Thanks for addressing the feedback :)

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

OpFoldResult BitCastOp::fold(FoldAdaptor adaptor) {
   Type srcElemType = getSourceVectorType().getElementType();
   Type dstElemType = getResultVectorType().getElementType();
 
   if (auto floatPack = llvm::dyn_cast<DenseFPElementsAttr>(sourceConstant)) {
     if (floatPack.isSplat()) {
       auto splat = floatPack.getSplatValue<FloatAttr>();
 
       // Casting fp16 into fp32.
       if (srcElemType.isF16() && dstElemType.isF32()) {
         uint32_t bits = static_cast<uint32_t>(
             splat.getValue().bitcastToAPInt().getZExtValue());
         // Duplicate the 16-bit pattern.
         bits = (bits << 16) | (bits & 0xffff);
         APInt intBits(32, bits);
         APFloat floatBits(llvm::APFloat::IEEEsingle(), intBits);
         return DenseElementsAttr::get(getResultVectorType(), floatBits);
       }
LogicalResult matchAndRewrite(TransposeOp transpOp,
 
     rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(
         transpOp, transpOp.getResultVectorType(),
         ArrayAttr::get(transpOp.getContext(), newMaskDimSizes));
     return success();
   }
 };
 
+static VectorType trimTrailingOneDims(VectorType oldType) {
+  ArrayRef<int64_t> oldShape = oldType.getShape();
+  ArrayRef<int64_t> newShape(oldShape.begin(), oldShape.end());
+
+  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 = oldType.getScalableDims().take_back();
+  }
+
+  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
+}
+
+/// Folds shape_cast(create_mask) into a new create_mask
+class FoldShapeCastCreateMask 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();
+
+    if (createMaskOp) {
+      auto maskOperands = createMaskOp.getOperands();
+      auto numMaskOperands = maskOperands.size();
+
+      // Check every mask dim size whether it can be dropped
+      for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
+           i--) {
+        auto constant = maskOperands[i].getDefiningOp<arith::ConstantIndexOp>();
+        // TODO: Support mask values other than 1 (0 might be the only other
+        // viable option)
+        if (!constant || (constant.value() != 1))
+          return failure();
+      }
+      SmallVector<Value> newMaskOperands =
+          maskOperands.drop_back(numDimsToDrop);
+
+      rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(
+          shapeOp, shapeOp.getResultVectorType(), newMaskOperands);
+      return success();
+    }
+
+    if (constantMaskOp) {
+      auto maskDimSizes = constantMaskOp.getMaskDimSizes().getValue();
+      auto numMaskOperands = maskDimSizes.size();
+
+      // Check every mask dim size whether it can be dropped
+      for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
+           i--) {
+        // TODO: Support mask values other than 1 (0 might be the only other
+        // viable option)
+        if (cast<IntegerAttr>(maskDimSizes[i]).getValue() != 1)
+          return failure();
+      }
+
+      auto newMaskOperands =
+          constantMaskOp.getMaskDimSizesAttr().getValue().drop_back(
+              numDimsToDrop);
+      ArrayAttr newMaskOperandsAttr = rewriter.getArrayAttr(newMaskOperands);
+
+      rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(
+          shapeOp, shapeOp.getResultVectorType(), newMaskOperandsAttr);
+      return success();
+    }
+
+    return failure();
+  }
+};
+
 } // namespace
 
 void vector::TransposeOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
-  results.add<FoldTransposeCreateMask, FoldTransposedScalarBroadcast,
-              TransposeFolder, FoldTransposeSplat>(context);
+  results
+      .add<FoldShapeCastCreateMask, FoldTransposeCreateMask,
+           FoldTransposedScalarBroadcast, TransposeFolder, FoldTransposeSplat>(
+          context);
 }
 
 void vector::TransposeOp::getTransp(SmallVectorImpl<int64_t> &results) {
   populateFromInt64AttrArray(getTransp(), results);
 }
 
 //===----------------------------------------------------------------------===//
 // ConstantMaskOp

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
+  %c0 = arith.constant 0 : index
+  %dim_0 = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+  %1 = vector.create_mask %c1, %dim_0, %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
+  %c0 = arith.constant 0 : index
+  %dim_0 = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+  %1 = vector.create_mask %c1, %dim_0, %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
+  %c0 = arith.constant 0 : index
+  %dim_0 = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+  %1 = vector.create_mask %c1, %dim_0, %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>
+}