[mlir][Standard] Allow select to use an i1 for vector and tensor values

[rriddle]

Summary

It currently requires that the condition match the shape of the selected value, but this is only really useful for things like masks. This revision allows for the use of i1 to mean that all of the vector/tensor is selected. This also matches the behavior of LLVM select. A benefit of this change is that transformations that want to generate selects, like those on the CFG, don't have to special case vector/tensor. Previously the only way to generate a select from an i1 was to use a splat, but that doesn't support dynamically shaped/unranked tensors.

Depends On D78683

[rriddle]

Remove unnecessary test

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1928

I would expect that you don't need the i1 here right? Or I misunderstand why you don't need the i1 in the "full vector selection" case below.

1934

It almost seems like an entirely different operation to me compared to the simple selection of one input or the other. Should we have a vselect or something?

[rriddle]

Resolve comments

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1928

Yep, thanks. Forgot to change this.

1934

vselect sounds weird given that it also supports tensors, but I agree that the semantics are kind of different. The optional predication is the only thing that removes this from solely being a "simple" select.

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1934

Do you plan to followup on looking into splitting this in two different ops?

[silvas]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1934

+1. maybe select_value for the one that is not describable as a pure elementwise vector/tensor op?

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
     floats. It has no standard attributes.
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 
-def SelectOp : Std_Op<"select", [NoSideEffect, SameOperandsAndResultShape,
-     AllTypesMatch<["true_value", "false_value", "result"]>,
-     TypesMatchWith<"condition type matches i1 equivalent of result type",
-                     "result", "condition",
-                     "getI1SameShape($_self)">]> {
+def SelectOp : Std_Op<"select", [NoSideEffect,
+     AllTypesMatch<["true_value", "false_value", "result"]>]> {
   let summary = "select operation";
   let description = [{
     The `select` operation chooses one value based on a binary condition
     supplied as its first operand. If the value of the first operand is `1`,
     the second operand is chosen, otherwise the third operand is chosen.
     The second and the third operand must have the same type.
 
     The operation applies to vectors and tensors elementwise given the _shape_
     of all operands is identical. The choice is made for each element
     individually based on the value at the same position as the element in the
-    condition operand.
+    condition operand. If an i1 is provided as the condition, the entire vector
+    or tensor is chosen.
 
     The `select` operation combined with [`cmpi`](#std-cmpi) can be used
     to implement `min` and `max` with signed or unsigned comparison semantics.
 
     Example:
 
     ```mlir
     // Custom form of scalar selection.
     %x = select %cond, %true, %false : i32
 
     // Generic form of the same operation.
     %x = "std.select"(%cond, %true, %false) : (i1, i32, i32) -> i32
 
-    // Vector selection is element-wise
-    %vx = "std.select"(%vcond, %vtrue, %vfalse)
-        : (vector<42xi1>, vector<42xf32>, vector<42xf32>) -> vector<42xf32>
+    // Element-wise vector selection.
+    %vx = std.select %vcond, %vtrue, %vfalse : vector<42xi1>, vector<42xf32>
+
+    // Full vector selection.
+    %vx = std.select %cond, %vtrue, %vfalse : vector<42xf32>
     ```
   }];
 
   let arguments = (ins BoolLike:$condition,
                        AnyType:$true_value,
                        AnyType:$false_value);
   let results = (outs AnyType:$result);
-  let verifier = ?;
 
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value condition,"
     "Value trueValue, Value falseValue", [{
       result.addOperands({condition, trueValue, falseValue});
       result.addTypes(trueValue.getType());
   }]>];
 
   let extraClassDeclaration = [{
       Value getCondition() { return condition(); }
       Value getTrueValue() { return true_value(); }
       Value getFalseValue() { return false_value(); }
   }];
 
   let hasFolder = 1;
-
-  let assemblyFormat = [{
-    $condition `,` $true_value `,` $false_value attr-dict `:` type($result)
-  }];
 }
 
 //===----------------------------------------------------------------------===//
 // ShiftLeftOp
 //===----------------------------------------------------------------------===//
 
 def ShiftLeftOp : IntArithmeticOp<"shift_left"> {
   let summary = "integer left-shift";

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

if (trueOperands == falseOperands) {
       return success();
     }
 
     // Otherwise, if the current block is the only predecessor insert selects
     // for any mismatched branch operands.
     if (trueDest->getUniquePredecessor() != condbr.getOperation()->getBlock())
       return failure();
 
-    // TODO: ATM Tensor/Vector SelectOp requires that the condition has the same
-    // shape as the operands. We should relax that to allow an i1 to signify
-    // that everything is selected.
-    auto doesntSupportsScalarI1 = [](Type type) {
-      return type.isa<TensorType>() || type.isa<VectorType>();
-    };
-    if (llvm::any_of(trueOperands.getTypes(), doesntSupportsScalarI1))
-      return failure();
-
     // Generate a select for any operands that differ between the two.
     SmallVector<Value, 8> mergedOperands;
     mergedOperands.reserve(trueOperands.size());
     Value condition = condbr.getCondition();
     for (auto it : llvm::zip(trueOperands, falseOperands)) {
       if (std::get<0>(it) == std::get<1>(it))
         mergedOperands.push_back(std::get<0>(it));
       else
if (matchPattern(condition, m_One()))
     return getTrueValue();
 
   // select false, %0, %1 => %1
   if (matchPattern(condition, m_Zero()))
     return getFalseValue();
   return nullptr;
 }
 
+static void print(OpAsmPrinter &p, SelectOp op) {
+  p << "select " << op.getOperands();
+  p.printOptionalAttrDict(op.getAttrs());
+  p << " : ";
+  if (ShapedType condType = op.getCondition().getType().dyn_cast<ShapedType>())
+    p << condType << ", ";
+  p << op.getType();
+}
+
+static ParseResult parseSelectOp(OpAsmParser &parser, OperationState &result) {
+  Type conditionType, resultType;
+  SmallVector<OpAsmParser::OperandType, 3> operands;
+  if (parser.parseOperandList(operands, /*requiredOperandCount=*/3) ||
+      parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseColonType(resultType))
+    return failure();
+
+  // Check for the explicit condition type if this is a masked tensor or vector.
+  if (succeeded(parser.parseOptionalComma())) {
+    conditionType = resultType;
+    if (parser.parseType(resultType))
+      return failure();
+  } else {
+    conditionType = parser.getBuilder().getI1Type();
+  }
+
+  result.addTypes(resultType);
+  return parser.resolveOperands(operands,
+                                {conditionType, resultType, resultType},
+                                parser.getNameLoc(), result.operands);
+}
+
+static LogicalResult verify(SelectOp op) {
+  Type conditionType = op.getCondition().getType();
+  if (conditionType.isSignlessInteger(1))
+    return success();
+
+  // If the result type is a vector or tensor, the type can be a mask with the
+  // same elements.
+  Type resultType = op.getType();
+  if (!resultType.isa<TensorType>() && !resultType.isa<VectorType>())
+    return op.emitOpError()
+           << "expected condition to be a signless i1, but got "
+           << conditionType;
+  Type shapedConditionType = getI1SameShape(resultType);
+  if (conditionType != shapedConditionType)
+    return op.emitOpError()
+           << "expected condition type to have the same shape "
+              "as the result type, expected "
+           << shapedConditionType << ", but got " << conditionType;
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // SignExtendIOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(SignExtendIOp op) {
   // Get the scalar type (which is either directly the type of the operand
   // or the vector's/tensor's element type.
   auto srcType = getElementTypeOrSelf(op.getOperand().getType());

mlir/test/Dialect/Standard/canonicalize-cf.mlir

^bb1(%result : i32):
   return
 }
 
 /// Test the folding of CondBranchOp when the successors are identical, but the
 /// arguments are different.
 
 // CHECK-LABEL: func @cond_br_same_successor_insert_select(
 // CHECK-SAME: %[[COND:.*]]: i1, %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
-func @cond_br_same_successor_insert_select(%cond : i1, %a : i32, %b : i32) -> i32 {
+// CHECK-SAME: %[[ARG2:.*]]: tensor<2xi32>, %[[ARG3:.*]]: tensor<2xi32>
+func @cond_br_same_successor_insert_select(
+      %cond : i1, %a : i32, %b : i32, %c : tensor<2xi32>, %d : tensor<2xi32>
+    ) -> (i32, tensor<2xi32>)  {
   // CHECK: %[[RES:.*]] = select %[[COND]], %[[ARG0]], %[[ARG1]]
-  // CHECK: return %[[RES]]
-
-  cond_br %cond, ^bb1(%a : i32), ^bb1(%b : i32)
-
-^bb1(%result : i32):
-  return %result : i32
-}
-
-/// Check that we don't generate a select if the type requires a splat.
-/// TODO: SelectOp should allow for matching a vector/tensor with i1.
-
-// CHECK-LABEL: func @cond_br_same_successor_no_select_tensor(
-func @cond_br_same_successor_no_select_tensor(%cond : i1, %a : tensor<2xi32>,
-                                              %b : tensor<2xi32>) -> tensor<2xi32>{
-  // CHECK: cond_br
-
-  cond_br %cond, ^bb1(%a : tensor<2xi32>), ^bb1(%b : tensor<2xi32>)
-
-^bb1(%result : tensor<2xi32>):
-  return %result : tensor<2xi32>
-}
-
-// CHECK-LABEL: func @cond_br_same_successor_no_select_vector(
-func @cond_br_same_successor_no_select_vector(%cond : i1, %a : vector<2xi32>,
-                                              %b : vector<2xi32>) -> vector<2xi32> {
-  // CHECK: cond_br
+  // CHECK: %[[RES2:.*]] = select %[[COND]], %[[ARG2]], %[[ARG3]]
+  // CHECK: return %[[RES]], %[[RES2]]
 
-  cond_br %cond, ^bb1(%a : vector<2xi32>), ^bb1(%b : vector<2xi32>)
+  cond_br %cond, ^bb1(%a, %c : i32, tensor<2xi32>), ^bb1(%b, %d : i32, tensor<2xi32>)
 
-^bb1(%result : vector<2xi32>):
-  return %result : vector<2xi32>
+^bb1(%result : i32, %result2 : tensor<2xi32>):
+  return %result, %result2 : i32, tensor<2xi32>
 }
 
 /// Test the compound folding of BranchOp and CondBranchOp.
 
 // CHECK-LABEL: func @cond_br_and_br_folding(
 func @cond_br_and_br_folding(%a : i32) {
   // CHECK-NEXT: return
 

mlir/test/IR/core-ops.mlir

^bb42(%t: tensor<4x4x?xf32>, %f: f32, %i: i32, %idx : index, %j: i64, %half: f16):
   %19 = cmpi "eq", %tci32, %tci32 : tensor<42 x i32>
 
   // CHECK: %{{[0-9]+}} = cmpi "eq", %cst_5, %cst_5 : vector<42xi32>
   %20 = cmpi "eq", %vci32, %vci32 : vector<42 x i32>
 
   // CHECK: %{{[0-9]+}} = select %{{[0-9]+}}, %arg3, %arg3 : index
   %21 = select %18, %idx, %idx : index
 
-  // CHECK: %{{[0-9]+}} = select %{{[0-9]+}}, %cst_4, %cst_4 : tensor<42xi32>
-  %22 = select %19, %tci32, %tci32 : tensor<42 x i32>
+  // CHECK: %{{[0-9]+}} = select %{{[0-9]+}}, %cst_4, %cst_4 : tensor<42xi1>, tensor<42xi32>
+  %22 = select %19, %tci32, %tci32 : tensor<42 x i1>, tensor<42 x i32>
 
-  // CHECK: %{{[0-9]+}} = select %{{[0-9]+}}, %cst_5, %cst_5 : vector<42xi32>
-  %23 = select %20, %vci32, %vci32 : vector<42 x i32>
+  // CHECK: %{{[0-9]+}} = select %{{[0-9]+}}, %cst_5, %cst_5 : vector<42xi1>, vector<42xi32>
+  %23 = select %20, %vci32, %vci32 : vector<42 x i1>, vector<42 x i32>
 
   // CHECK: %{{[0-9]+}} = select %{{[0-9]+}}, %arg3, %arg3 : index
   %24 = "std.select"(%18, %idx, %idx) : (i1, index, index) -> index
 
   // CHECK: %{{[0-9]+}} = select %{{[0-9]+}}, %cst_4, %cst_4 : tensor<42xi32>
-  %25 = "std.select"(%19, %tci32, %tci32) : (tensor<42 x i1>, tensor<42 x i32>, tensor<42 x i32>) -> tensor<42 x i32>
+  %25 = std.select %18, %tci32, %tci32 : tensor<42 x i32>
 
   // CHECK: %{{[0-9]+}} = divi_signed %arg2, %arg2 : i32
   %26 = divi_signed %i, %i : i32
 
   // CHECK: %{{[0-9]+}} = divi_signed %arg3, %arg3 : index
   %27 = divi_signed %idx, %idx : index
 
   // CHECK: %{{[0-9]+}} = divi_signed %cst_5, %cst_5 : vector<42xi32>

mlir/test/IR/invalid-ops.mlir

 func @func_with_ops(i1, i32, i64) {
 ^bb0(%cond : i1, %t : i32, %f : i64):
   // expected-error@+1 {{all of {true_value, false_value, result} have same type}}
   %r = "std.select"(%cond, %t, %f) : (i1, i32, i64) -> i32
 }
 
 // -----
 
-func @func_with_ops(i1, vector<42xi32>, vector<42xi32>) {
-^bb0(%cond : i1, %t : vector<42xi32>, %f : vector<42xi32>):
-  // expected-error@+1 {{requires the same shape for all operands and results}}
-  %r = "std.select"(%cond, %t, %f) : (i1, vector<42xi32>, vector<42xi32>) -> vector<42xi32>
+func @func_with_ops(vector<12xi1>, vector<42xi32>, vector<42xi32>) {
+^bb0(%cond : vector<12xi1>, %t : vector<42xi32>, %f : vector<42xi32>):
+  // expected-error@+1 {{expected condition type to have the same shape as the result type, expected 'vector<42xi1>', but got 'vector<12xi1>'}}
+  %r = "std.select"(%cond, %t, %f) : (vector<12xi1>, vector<42xi32>, vector<42xi32>) -> vector<42xi32>
 }
 
 // -----
 
-func @func_with_ops(i1, tensor<42xi32>, tensor<?xi32>) {
-^bb0(%cond : i1, %t : tensor<42xi32>, %f : tensor<?xi32>):
-  // expected-error@+1 {{ op requires the same shape for all operands and results}}
-  %r = "std.select"(%cond, %t, %f) : (i1, tensor<42xi32>, tensor<?xi32>) -> tensor<42xi32>
+func @func_with_ops(tensor<12xi1>, tensor<42xi32>, tensor<42xi32>) {
+^bb0(%cond : tensor<12xi1>, %t : tensor<42xi32>, %f : tensor<42xi32>):
+  // expected-error@+1 {{expected condition type to have the same shape as the result type, expected 'tensor<42xi1>', but got 'tensor<12xi1>'}}
+  %r = "std.select"(%cond, %t, %f) : (tensor<12xi1>, tensor<42xi32>, tensor<42xi32>) -> tensor<42xi32>
 }
 
 // -----
 
 func @invalid_cmp_shape(%idx : () -> ()) {
   // expected-error@+1 {{'lhs' must be signless-integer-like, but got '() -> ()'}}
   %cmp = cmpi "eq", %idx, %idx : () -> ()