[mlir][arith] Fold select with poison

If either of the operands of `select` is fully poisoned we can simply return the other.
This PR implements this optimization inside the `fold` method.

Note that this patch is the first to add a dependency on the UB dialect within Arith. Given this was inevitable (and part of the motivation) it should be fine I believe.

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

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

@@ -13,6 +13,7 @@
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/CommonFolders.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -2194,6 +2195,13 @@ OpFoldResult arith::SelectOp::fold(FoldAdaptor adaptor) {
   if (matchPattern(condition, m_Zero()))
     return falseVal;
 
+  // If either operand is fully poisoned, return the other.
+  if (isa_and_nonnull<ub::PoisonAttr>(adaptor.getTrueValue()))
+    return falseVal;
+
+  if (isa_and_nonnull<ub::PoisonAttr>(adaptor.getFalseValue()))
+    return trueVal;
+
   // select %x, true, false => %x
   if (getType().isInteger(1) && matchPattern(getTrueValue(), m_One()) &&
       matchPattern(getFalseValue(), m_Zero()))

mlir/lib/Dialect/Arith/IR/CMakeLists.txt

@@ -28,6 +28,7 @@ add_mlir_dialect_library(MLIRArithDialect
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRIR
+  MLIRUBDialect
   )
 
 add_mlir_dialect_library(MLIRArithValueBoundsOpInterfaceImpl

mlir/test/Dialect/Arith/canonicalize.mlir

@@ -2567,3 +2567,20 @@ func.func @foldOrXor6(%arg0: index) -> index {
   %2 = arith.ori %arg0, %1 : index
   return %2 : index
 }
+
+// CHECK-LABEL: @selectOfPoison
+// CHECK-SAME: %[[ARG:[[:alnum:]]+]]: i32
+// CHECK: %[[UB:.*]] = ub.poison : i32
+// CHECK: return %[[ARG]], %[[ARG]], %[[UB]], %[[ARG]]
+func.func @selectOfPoison(%cond : i1, %arg: i32) -> (i32, i32, i32, i32) {
+  %poison = ub.poison : i32
+  %select1 = arith.select %cond, %poison, %arg : i32
+  %select2 = arith.select %cond, %arg, %poison : i32
+
+  // Check that constant folding is applied prior to poison handling.
+  %true = arith.constant true
+  %false = arith.constant false
+  %select3 = arith.select %true, %poison, %arg : i32
+  %select4 = arith.select %false, %poison, %arg : i32
+  return %select1, %select2, %select3, %select4 : i32, i32, i32, i32
+}