[mlir][arith] Fold `select` with poison

[zero9178]

Summary

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.

[Hardcode84]

Which folding should applied first?

%val = select %true %poison %non_poison

or more complicated example

%1 = foldable_op %true
%2 = foldable_op %poison
%val = select %1 %2 %non_poison

Here we can end up with different code depending on order of folders.

[kuhar]

In D158986#4621310, @Hardcode84 wrote:

Which folding should applied first?

%val = select %true %poison %non_poison

or more complicated example

%1 = foldable_op %true
%2 = foldable_op %poison
%val = select %1 %2 %non_poison

Here we can end up with different code depending on order of folders.

Good question, we should have a test case for this.
I think either one is fine but select true, %poison, %nonpoison => %poison should probably allow for more speculation later. What does llvm do in this case?

mlir/test/Dialect/Arith/canonicalize.mlir

2579

nit: missing newline

[kuhar]

@zero9178

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.

This is super exciting! Do you have some immediate plan for growing this set of folds?

[zero9178]

Change order of fold application, doing constant folding prior to handling poison

[zero9178]

In D158986#4621310, @Hardcode84 wrote:

Which folding should applied first?

%val = select %true %poison %non_poison

or more complicated example

%1 = foldable_op %true
%2 = foldable_op %poison
%val = select %1 %2 %non_poison

Here we can end up with different code depending on order of folders.

Good point! I ran both patterns through alive2 (yes we can define different semantics for arith than LLVM, but for a lack of tooling lets go with that) and both transforms are luckily valid.
I think it makes a lot more sense to return the poison in this case as that will enable more transformations. This also matches what LLVM does.

[zero9178]

In D158986#4621475, @kuhar wrote:

@zero9178

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.

This is super exciting! Do you have some immediate plan for growing this set of folds?

I sadly do not have any concrete plans. The motivation for this fold in particular was that the lift-cfg-to-scf pattern creates ub.poison which after canonicalization end up in arith.select. So unless I see a need to implement more of these folds/patterns in arith I cannot guarantee that I'll be working on these anytime soon

[kuhar]

mlir/test/Dialect/Arith/canonicalize.mlir

2582

Can you also add a case when poison is the second alternative or with false, so that we can disambiguate returning poison value from returning the true/false value?

[zero9178]

Address review comments

[kuhar]

mlir/test/Dialect/Arith/canonicalize.mlir

2584

I thought you wanted this to be poison like in the original code?

I think it makes a lot more sense to return the poison in this case as that will enable more transformations. This also matches what LLVM does.

[zero9178]

mlir/test/Dialect/Arith/canonicalize.mlir

2584

That would be illegal if we are to follow LLVMs example: https://alive2.llvm.org/ce/z/3-KebL (this is also very desireable as this matches the semantics of condition control flow and would otherwise make a lot of the scf.if and cf.cond_br folds to arith.select illegal otherwise.
It was specifically select true, %poison, %other where I wanted it to return %poison rather than %other.

[kuhar]

LGTM

mlir/test/Dialect/Arith/canonicalize.mlir

2584

Sorry, I misread your testcase, I thought it was select false, arg, poison which can be folded: https://alive2.llvm.org/ce/z/unind0

[Hardcode84]

LGTM

[Hardcode84]

In D158986#4621475, @kuhar wrote:

@zero9178

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.

This is super exciting! Do you have some immediate plan for growing this set of folds?

Regarding more foldings, adding poison arg support to constFoldBinaryOpConditional and friends in CommonFolders.h seems straightforward and it should cover most of arith/math ops. I can create a patch.

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

 //===- ArithOps.cpp - MLIR Arith dialect ops implementation -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstdint>
 #include <functional>
 #include <utility>
 
 #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"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 
OpFoldResult arith::SelectOp::fold(FoldAdaptor adaptor) {
   // select true, %0, %1 => %0
   if (matchPattern(condition, m_One()))
     return trueVal;
 
   // select false, %0, %1 => %1
   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()))
     return condition;
 
   if (auto cmp = dyn_cast_or_null<arith::CmpIOp>(condition.getDefiningOp())) {
     auto pred = cmp.getPredicate();
     if (pred == arith::CmpIPredicate::eq || pred == arith::CmpIPredicate::ne) {

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

add_mlir_dialect_library(MLIRArithDialect
 
   LINK_LIBS PUBLIC
   MLIRCastInterfaces
   MLIRDialect
   MLIRInferIntRangeCommon
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRIR
+  MLIRUBDialect
   )
 
 add_mlir_dialect_library(MLIRArithValueBoundsOpInterfaceImpl
   ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRIR
   MLIRValueBoundsOpInterface
   )

mlir/test/Dialect/Arith/canonicalize.mlir

 //       CHECK:   %[[ONE:.*]] = arith.constant -1
 //       CHECK:   return %[[ONE]]
 func.func @foldOrXor6(%arg0: index) -> index {
   %0 = arith.constant -1 : index
   %1 = arith.xori %arg0, %0 : 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
+}