[mlir][Vector] Fold selects of single-element i1 vectors

[dcaballe]

Summary

This patch adds a folding to select operation between an all-true and all-false vector.
For now, only single element vectors (i.e., vector<1xi1>) are supported. Multi-element
cases are caught by InstCombine.

[ThomasRaoux]

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

1079–1082

it is probably better to do:

return denseAttr.isSplat() && denseAttr.getSplatValue<bool>() == value;

looping through all the elements can be expensive for large shapes.

1135

should this be a vector.broadcast? That would generalize to any vector size?

[awarzynski]

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

1107

There's not sign extension taking place here, is there?

[dcaballe]

Addressed feedback

[dcaballe]

Sorry for the long delay. I just forgot about this! :) I addressed all the comments.

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

1107

Not for the scalar case, good catch :)

1135

Good point!

[dcaballe]

Fix naming in test

[awarzynski]

LGTM, thanks!

I see that you have also addressed comments from Thomas, so IMHO it's OK to land this without waiting for further feedback.

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

rewriter.replaceOp(
                                   op.getOperand(0)));
     return success();
   }
 
 private:
   const bool force32BitVectorIndices;
 };
 
+/// Returns true if all the `i1` elements of `constantOp` are set to `value`.
+static bool allI1ConstantValuesSetTo(arith::ConstantOp constantOp, bool value) {
+  auto denseAttr = dyn_cast<DenseIntElementsAttr>(constantOp.getValue());
+  // TODO: Support non-dense constant.
+  if (!denseAttr)
+    return false;
+
+  assert(denseAttr.getElementType().isInteger(1) && "Unexpected type");
+  return denseAttr.isSplat() && denseAttr.getSplatValue<bool>() == value;
+}
+
+/// Folds a select operation between an all-true and all-false vector. For now,
+/// only single element vectors (i.e., vector<1xi1>) are supported. That is:
+///
+///   %true = arith.constant dense<true> : vector<1xi1>
+///   %false = arith.constant dense<false> : vector<1xi1>
+///   %result = arith.select %cond, %true, %false : i1, vector<1xi1>
+///   =>
+///   %result = vector.broadcast %cond : i1 to vector<1xi1>
+///
+/// InstCombine seems to handle vectors with multiple elements but not the
+/// single element ones.
+struct FoldI1Select : public OpRewritePattern<arith::SelectOp> {
+  using OpRewritePattern<arith::SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::SelectOp selectOp,
+                                PatternRewriter &rewriter) const override {
+    auto vecType = dyn_cast<VectorType>(selectOp.getType());
+    if (!vecType || !vecType.getElementType().isInteger(1))
+      return failure();
+
+    // Only scalar conditions can be folded.
+    Value cond = selectOp.getCondition();
+    if (isa<VectorType>(cond.getType()))
+      return failure();
+
+    // TODO: Support n-D and scalable vectors.
+    if (vecType.getRank() != 1 || vecType.isScalable())
+      return failure();
+
+    // TODO: Support vectors with multiple elements.
+    if (vecType.getShape()[0] != 1)
+      return failure();
+
+    auto trueConst = selectOp.getTrueValue().getDefiningOp<arith::ConstantOp>();
+    if (!trueConst || !allI1ConstantValuesSetTo(trueConst, true))
+      return failure();
+
+    auto falseConst =
+        selectOp.getFalseValue().getDefiningOp<arith::ConstantOp>();
+    if (!falseConst || !allI1ConstantValuesSetTo(falseConst, false))
+      return failure();
+
+    // Replace select with its condition broadcasted to single element vector.
+    auto elemType = rewriter.getIntegerType(vecType.getNumElements());
+    auto bcastType = VectorType::get(/*shape=*/{1}, elemType);
+    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(selectOp, bcastType, cond);
+    return success();
+  }
+};
+
 // Drop inner most contiguous unit dimensions from transfer_read operand.
 class DropInnerMostUnitDims : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                 PatternRewriter &rewriter) const override {
     // TODO: support 0-d corner case.
     if (readOp.getTransferRank() == 0)
       return failure();
 
     // TODO: support mask.
     if (readOp.getMask())
       return failure();

 
 void mlir::vector::populateVectorMaskMaterializationPatterns(
     RewritePatternSet &patterns, bool force32BitVectorIndices,
     PatternBenefit benefit) {
   patterns.add<VectorCreateMaskOpConversion,
                MaterializeTransferMask<vector::TransferReadOp>,
                MaterializeTransferMask<vector::TransferWriteOp>>(
       patterns.getContext(), force32BitVectorIndices, benefit);
+  patterns.add<FoldI1Select>(patterns.getContext(), benefit);
 }
 
 void mlir::vector::populateShapeCastFoldingPatterns(RewritePatternSet &patterns,
                                                     PatternBenefit benefit) {
   patterns.add<ShapeCastOpFolder>(patterns.getContext(), benefit);
 }
 
 void mlir::vector::populateBubbleVectorBitCastOpPatterns(

mlir/test/Dialect/Vector/vector-materialize-mask.mlir

+// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
+
+func.func @select_single_i1_vector(%cond : i1) -> vector<1xi1> {
+  %true = arith.constant dense<true> : vector<1xi1>
+  %false = arith.constant dense<false> : vector<1xi1>
+  %select = arith.select %cond, %true, %false : i1, vector<1xi1>
+  return %select : vector<1xi1>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%func_op: !transform.op<"func.func">):
+  transform.apply_patterns to %func_op {
+    transform.apply_patterns.vector.materialize_masks
+  } : !transform.op<"func.func">
+}
+
+// CHECK-LABEL: func @select_single_i1_vector
+// CHECK-SAME: %[[COND:.*]]: i1
+// CHECK:      %[[BCAST:.*]] = vector.broadcast %[[COND]] : i1 to vector<1xi1>
+// CHECK:      return %[[BCAST]] : vector<1xi1>