[mlir] Move SelectOp from Standard to Arithmetic

This is part of splitting up the standard dialect. See https://llvm.discourse.group/t/standard-dialect-the-final-chapter/ for discussion.

Differential Revision: https://reviews.llvm.org/D118648

Author: River707

flang/lib/Optimizer/Builder/Character.cpp

@@ -434,7 +434,7 @@ mlir::Value genMin(fir::FirOpBuilder &builder, mlir::Location loc,
                    mlir::Value a, mlir::Value b) {
   auto cmp =
       builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, a, b);
-  return builder.create<mlir::SelectOp>(loc, cmp, a, b);
+  return builder.create<mlir::arith::SelectOp>(loc, cmp, a, b);
 }
 
 void fir::factory::CharacterExprHelper::createAssign(
@@ -532,7 +532,8 @@ fir::CharBoxValue fir::factory::CharacterExprHelper::createSubstring(
   auto zero = builder.createIntegerConstant(loc, substringLen.getType(), 0);
   auto cdt = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt,
                                            substringLen, zero);
-  substringLen = builder.create<mlir::SelectOp>(loc, cdt, zero, substringLen);
+  substringLen =
+      builder.create<mlir::arith::SelectOp>(loc, cdt, zero, substringLen);
 
   return {substringRef, substringLen};
 }
@@ -570,8 +571,8 @@ fir::factory::CharacterExprHelper::createLenTrim(const fir::CharBoxValue &str) {
   // Compute length after iteration (zero if all blanks)
   mlir::Value newLen =
       builder.create<arith::AddIOp>(loc, iterWhile.getResult(1), one);
-  auto result =
-      builder.create<mlir::SelectOp>(loc, iterWhile.getResult(0), zero, newLen);
+  auto result = builder.create<mlir::arith::SelectOp>(
+      loc, iterWhile.getResult(0), zero, newLen);
   return builder.createConvert(loc, builder.getCharacterLengthType(), result);
 }
 

flang/lib/Optimizer/Builder/MutableBox.cpp

@@ -675,8 +675,8 @@ void fir::factory::genReallocIfNeeded(fir::FirOpBuilder &builder,
           // reallocate = reallocate || previous != required
           auto cmp = builder.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::ne, castPrevious, required);
-          mustReallocate =
-              builder.create<mlir::SelectOp>(loc, cmp, cmp, mustReallocate);
+          mustReallocate = builder.create<mlir::arith::SelectOp>(
+              loc, cmp, cmp, mustReallocate);
         };
         llvm::SmallVector<mlir::Value> previousLbounds;
         llvm::SmallVector<mlir::Value> previousExtents =

flang/lib/Optimizer/Builder/Runtime/Numeric.cpp

@@ -295,7 +295,8 @@ mlir::Value fir::runtime::genNearest(fir::FirOpBuilder &builder,
   mlir::Value False = builder.createIntegerConstant(loc, boolTy, 0);
   mlir::Value True = builder.createIntegerConstant(loc, boolTy, 1);
 
-  mlir::Value positive = builder.create<mlir::SelectOp>(loc, cmp, True, False);
+  mlir::Value positive =
+      builder.create<mlir::arith::SelectOp>(loc, cmp, True, False);
   auto args = fir::runtime::createArguments(builder, loc, funcTy, x, positive);
 
   return builder.create<fir::CallOp>(loc, func, args).getResult(0);

flang/lib/Optimizer/Transforms/RewriteLoop.cpp

@@ -75,7 +75,7 @@ class CfgLoopConv : public mlir::OpRewritePattern<fir::DoLoopOp> {
       auto cond = rewriter.create<mlir::arith::CmpIOp>(
           loc, arith::CmpIPredicate::sle, iters, zero);
       auto one = rewriter.create<mlir::arith::ConstantIndexOp>(loc, 1);
-      iters = rewriter.create<mlir::SelectOp>(loc, cond, one, iters);
+      iters = rewriter.create<mlir::arith::SelectOp>(loc, cond, one, iters);
     }
 
     llvm::SmallVector<mlir::Value> loopOperands;

flang/test/Fir/loop02.fir

@@ -23,7 +23,7 @@ func private @y(%addr : !fir.ref<index>)
 // CHECK:           %[[VAL_6:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_7:.*]] = arith.cmpi sle, %[[VAL_5]], %[[VAL_6]] : index
 // CHECK:           %[[VAL_8:.*]] = arith.constant 1 : index
-// CHECK:           %[[VAL_9:.*]] = select %[[VAL_7]], %[[VAL_8]], %[[VAL_5]] : index
+// CHECK:           %[[VAL_9:.*]] = arith.select %[[VAL_7]], %[[VAL_8]], %[[VAL_5]] : index
 // CHECK:           br ^bb1(%[[VAL_1]], %[[VAL_9]] : index, index)
 // CHECK:         ^bb1(%[[VAL_10:.*]]: index, %[[VAL_11:.*]]: index):
 // CHECK:           %[[VAL_12:.*]] = arith.constant 0 : index

flang/unittests/Optimizer/Builder/Runtime/NumericTest.cpp

@@ -56,8 +56,8 @@ void testGenNearest(fir::FirOpBuilder &builder, mlir::Type xType,
   checkCallOp(nearest.getDefiningOp(), fctName, 2, /*addLocArg=*/false);
   auto callOp = mlir::dyn_cast<fir::CallOp>(nearest.getDefiningOp());
   mlir::Value select = callOp.getOperands()[1];
-  EXPECT_TRUE(mlir::isa<mlir::SelectOp>(select.getDefiningOp()));
-  auto selectOp = mlir::dyn_cast<mlir::SelectOp>(select.getDefiningOp());
+  EXPECT_TRUE(mlir::isa<mlir::arith::SelectOp>(select.getDefiningOp()));
+  auto selectOp = mlir::dyn_cast<mlir::arith::SelectOp>(select.getDefiningOp());
   mlir::Value cmp = selectOp.getCondition();
   EXPECT_TRUE(mlir::isa<mlir::arith::CmpFOp>(cmp.getDefiningOp()));
   auto cmpOp = mlir::dyn_cast<mlir::arith::CmpFOp>(cmp.getDefiningOp());

mlir/benchmark/python/common.py

@@ -29,7 +29,7 @@ def setup_passes(mlir_module):
         f"convert-scf-to-std,"
         f"func-bufferize,"
         f"arith-bufferize,"
-        f"builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),"
+        f"builtin.func(tensor-bufferize,finalizing-bufferize),"
         f"convert-vector-to-llvm"
         f"{{reassociate-fp-reductions=1 enable-index-optimizations=1}},"
         f"lower-affine,"

mlir/docs/Bufferization.md

@@ -87,7 +87,6 @@ The code, slightly simplified and annotated, is reproduced here:
   pm.addNestedPass<FuncOp>(createTCPBufferizePass()); // Bufferizes the downstream `tcp` dialect.
   pm.addNestedPass<FuncOp>(createSCFBufferizePass());
   pm.addNestedPass<FuncOp>(createLinalgBufferizePass());
-  pm.addNestedPass<FuncOp>(createStdBufferizePass());
   pm.addNestedPass<FuncOp>(createTensorBufferizePass());
   pm.addPass(createFuncBufferizePass());
 

mlir/include/mlir/Dialect/Arithmetic/IR/Arithmetic.h

@@ -12,6 +12,7 @@
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CastInterfaces.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

@@ -1088,7 +1088,7 @@ def Arith_CmpIOp : Arith_CompareOpOfAnyRank<"cmpi"> {
     %x = arith.cmpi "eq", %lhs, %rhs : vector<4xi64>
 
     // Generic form of the same operation.
-    %x = "std.arith.cmpi"(%lhs, %rhs) {predicate = 0 : i64}
+    %x = "arith.cmpi"(%lhs, %rhs) {predicate = 0 : i64}
         : (vector<4xi64>, vector<4xi64>) -> vector<4xi1>
     ```
   }];
@@ -1161,4 +1161,55 @@ def Arith_CmpFOp : Arith_CompareOp<"cmpf"> {
   let hasFolder = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// SelectOp
+//===----------------------------------------------------------------------===//
+
+def SelectOp : Arith_Op<"select", [
+    AllTypesMatch<["true_value", "false_value", "result"]>
+  ] # ElementwiseMappable.traits> {
+  let summary = "select operation";
+  let description = [{
+    The `arith.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. If an i1 is provided as the condition, the entire vector
+    or tensor is chosen.
+
+    Example:
+
+    ```mlir
+    // Custom form of scalar selection.
+    %x = arith.select %cond, %true, %false : i32
+
+    // Generic form of the same operation.
+    %x = "arith.select"(%cond, %true, %false) : (i1, i32, i32) -> i32
+
+    // Element-wise vector selection.
+    %vx = arith.select %vcond, %vtrue, %vfalse : vector<42xi1>, vector<42xf32>
+
+    // Full vector selection.
+    %vx = arith.select %cond, %vtrue, %vfalse : vector<42xf32>
+    ```
+  }];
+
+  let arguments = (ins BoolLike:$condition,
+                       AnyType:$true_value,
+                       AnyType:$false_value);
+  let results = (outs AnyType:$result);
+
+  let hasCanonicalizer = 1;
+  let hasFolder = 1;
+  let hasVerifier = 1;
+  
+  // FIXME: Switch this to use the declarative assembly format.
+  let printer = [{ return ::print(p, *this); }];
+  let parser = [{ return ::parse$cppClass(parser, result); }];
+}
+
 #endif // ARITHMETIC_OPS

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

@@ -146,7 +146,7 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
 
           Note: This method can return multiple OpOperands, indicating that the
           given OpResult may at runtime alias with any of the OpOperands. This
-          is useful for branches and for ops such as `std.select`.
+          is useful for branches and for ops such as `arith.select`.
         }],
         /*retType=*/"SmallVector<OpOperand *>",
         /*methodName=*/"getAliasingOpOperand",

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

@@ -456,57 +456,6 @@ def ReturnOp : Std_Op<"return", [NoSideEffect, HasParent<"FuncOp">,
   let hasVerifier = 1;
 }
 
-//===----------------------------------------------------------------------===//
-// SelectOp
-//===----------------------------------------------------------------------===//
-
-def SelectOp : Std_Op<"select", [NoSideEffect,
-    AllTypesMatch<["true_value", "false_value", "result"]>,
-    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
-    ElementwiseMappable.traits> {
-  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. If an i1 is provided as the condition, the entire vector
-    or tensor is chosen.
-
-    The `select` operation combined with [`cmpi`](#stdcmpi-cmpiop) 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
-
-    // 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 hasCanonicalizer = 1;
-  let hasFolder = 1;
-  let hasVerifier = 1;
-}
-
 //===----------------------------------------------------------------------===//
 // SwitchOp
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/StandardOps/Transforms/BufferizableOpInterfaceImpl.h

@@ -23,9 +23,6 @@ class BufferizeTypeConverter;
 
 class RewritePatternSet;
 
-/// Creates an instance of std bufferization pass.
-std::unique_ptr<Pass> createStdBufferizePass();
-
 /// Creates an instance of func bufferization pass.
 std::unique_ptr<Pass> createFuncBufferizePass();
 

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.h

@@ -11,11 +11,6 @@
 
 include "mlir/Pass/PassBase.td"
 
-def StdBufferize : Pass<"std-bufferize", "FuncOp"> {
-  let summary = "Bufferize the std dialect";
-  let constructor = "mlir::createStdBufferizePass()";
-}
-
 def FuncBufferize : Pass<"func-bufferize", "ModuleOp"> {
   let summary = "Bufferize func/call/return ops";
   let description = [{

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

@@ -30,13 +30,14 @@ SmallVector<Value> condenseValues(const SmallVector<Value> &values);
 
 // Takes the parameters for a clamp and turns it into a series of ops.
 template <typename T, typename P>
-mlir::SelectOp clampHelper(Location loc, Value arg, arith::ConstantOp min,
-                           arith::ConstantOp max, P pred, OpBuilder &rewriter) {
+arith::SelectOp clampHelper(Location loc, Value arg, arith::ConstantOp min,
+                            arith::ConstantOp max, P pred,
+                            OpBuilder &rewriter) {
   auto smallerThanMin = rewriter.create<T>(loc, pred, arg, min);
   auto minOrArg =
-      rewriter.create<mlir::SelectOp>(loc, smallerThanMin, min, arg);
+      rewriter.create<arith::SelectOp>(loc, smallerThanMin, min, arg);
   auto largerThanMax = rewriter.create<T>(loc, pred, max, arg);
-  return rewriter.create<mlir::SelectOp>(loc, largerThanMax, max, minOrArg);
+  return rewriter.create<arith::SelectOp>(loc, largerThanMax, max, minOrArg);
 }
 
 // Returns the values in an attribute as an array of values.

mlir/include/mlir/Dialect/Tosa/Utils/CoversionUtils.h

@@ -1368,14 +1368,14 @@ struct Elementwise : public TraitBase<ConcreteType, Elementwise> {
 ///
 /// Example:
 /// ```
-/// %tensor_select = "std.select"(%pred_tensor, %true_val, %false_val)
+/// %tensor_select = "arith.select"(%pred_tensor, %true_val, %false_val)
 ///                      : (tensor<?xi1>, tensor<?xf32>, tensor<?xf32>)
 ///                      -> tensor<?xf32>
 /// ```
 /// can be scalarized to
 ///
 /// ```
-/// %scalar_select = "std.select"(%pred, %true_val_scalar, %false_val_scalar)
+/// %scalar_select = "arith.select"(%pred, %true_val_scalar, %false_val_scalar)
 ///                      : (i1, f32, f32) -> f32
 /// ```
 template <typename ConcreteType>
@@ -1430,12 +1430,12 @@ struct Vectorizable : public TraitBase<ConcreteType, Vectorizable> {
 /// ```
 ///
 /// ```
-/// %scalar_pred = "std.select"(%pred, %true_val, %false_val)
+/// %scalar_pred = "arith.select"(%pred, %true_val, %false_val)
 ///                    : (i1, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
 /// ```
 /// can be tensorized to
 /// ```
-/// %tensor_pred = "std.select"(%pred, %true_val, %false_val)
+/// %tensor_pred = "arith.select"(%pred, %true_val, %false_val)
 ///                    : (tensor<?xi1>, tensor<?xf32>, tensor<?xf32>)
 ///                    -> tensor<?xf32>
 /// ```

mlir/include/mlir/IR/OpDefinition.h

@@ -96,8 +96,8 @@ class AffineApplyExpander
         loc, arith::CmpIPredicate::slt, remainder, zeroCst);
     Value correctedRemainder =
         builder.create<arith::AddIOp>(loc, remainder, rhs);
-    Value result = builder.create<SelectOp>(loc, isRemainderNegative,
-                                            correctedRemainder, remainder);
+    Value result = builder.create<arith::SelectOp>(
+        loc, isRemainderNegative, correctedRemainder, remainder);
     return result;
   }
 
@@ -134,12 +134,12 @@ class AffineApplyExpander
         loc, arith::CmpIPredicate::slt, lhs, zeroCst);
     Value negatedDecremented = builder.create<arith::SubIOp>(loc, noneCst, lhs);
     Value dividend =
-        builder.create<SelectOp>(loc, negative, negatedDecremented, lhs);
+        builder.create<arith::SelectOp>(loc, negative, negatedDecremented, lhs);
     Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
     Value correctedQuotient =
         builder.create<arith::SubIOp>(loc, noneCst, quotient);
-    Value result =
-        builder.create<SelectOp>(loc, negative, correctedQuotient, quotient);
+    Value result = builder.create<arith::SelectOp>(loc, negative,
+                                                   correctedQuotient, quotient);
     return result;
   }
 
@@ -175,14 +175,14 @@ class AffineApplyExpander
     Value negated = builder.create<arith::SubIOp>(loc, zeroCst, lhs);
     Value decremented = builder.create<arith::SubIOp>(loc, lhs, oneCst);
     Value dividend =
-        builder.create<SelectOp>(loc, nonPositive, negated, decremented);
+        builder.create<arith::SelectOp>(loc, nonPositive, negated, decremented);
     Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
     Value negatedQuotient =
         builder.create<arith::SubIOp>(loc, zeroCst, quotient);
     Value incrementedQuotient =
         builder.create<arith::AddIOp>(loc, quotient, oneCst);
-    Value result = builder.create<SelectOp>(loc, nonPositive, negatedQuotient,
-                                            incrementedQuotient);
+    Value result = builder.create<arith::SelectOp>(
+        loc, nonPositive, negatedQuotient, incrementedQuotient);
     return result;
   }
 
@@ -259,7 +259,8 @@ static Value buildMinMaxReductionSeq(Location loc,
   Value value = *valueIt++;
   for (; valueIt != values.end(); ++valueIt) {
     auto cmpOp = builder.create<arith::CmpIOp>(loc, predicate, value, *valueIt);
-    value = builder.create<SelectOp>(loc, cmpOp.getResult(), value, *valueIt);
+    value = builder.create<arith::SelectOp>(loc, cmpOp.getResult(), value,
+                                            *valueIt);
   }
 
   return value;