[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