[InstCombine] Fold icmp (select c,const,arg), null if icmp arg, null can be simplified

[aqjune]

Summary

This patch folds icmp (select c,const,arg), null if icmp arg, null can be simplified.

Resolves llvm.org/pr48975.

[xbolva00]

SimplifyCFG adds llvm.assume for “cond”, we should do it here as well.

[aqjune]

In D96663#2562250, @xbolva00 wrote:

SimplifyCFG adds llvm.assume for “cond”, we should do it here as well.

It seems SimplifyCFG is introducing llvm.assume to preserve information after removing a conditional branch.
In this case, the nonnull attribute is still there, so it might be okay to do the transformation without introducing assume IMO. What do you think?

[xbolva00]

In D96663#2571467, @aqjune wrote:

In D96663#2562250, @xbolva00 wrote:

SimplifyCFG adds llvm.assume for “cond”, we should do it here as well.

It seems SimplifyCFG is introducing llvm.assume to preserve information after removing a conditional branch.
In this case, the nonnull attribute is still there, so it might be okay to do the transformation without introducing assume IMO. What do you think?

But when we remove select, we can remove condition too (if trivially dead) - so maybe worth to preserve info that cond must be true/false?

[aqjune]

It sounds like a good idea, I updated the patch to insert assume(cond).

I think this is allowed only when nonnull is accompanied with noundef attribute:
llvm.assume(cond) immediately raises UB whereas passing null to nonnull is equivalent
to passing poison; noundef is needed for this according to LangRef.

[xbolva00]

Yes, thanks. LG.

Other reviewers?

[nikic]

I don't think this should be creating assumptions. It's pretty common for InstCombine transforms to lose some information while making folds, and we currently don't have any other InstCombine folds that try to preserve information through assumptions. Unless there is something specific to this transform that makes this particularly worthwhile, we should not deviate from the general pattern.

[aqjune]

I think we can say that this is conceptually equivalent to SimplifyCFG that removes a conditional branch, and SimplifyCFG inserts an assumption when removing a condition (test9_null_callsite at llvm/test/Transforms/SimplifyCFG/UnreachableEliminate.ll), so this transformation is also chosen to insert assume.
To speak generally, we can say that control flow information is preserved using assume in InstCombine, and the scheme should be consistent with what SimplifyCFG says.
But would it be a big change though?

[xbolva00]

^ Yes, this is exactly what I had in my mind too.

[spatel]

The code duplication suggests that we should have a common place for this logic (especially if it already exists in SimplifyCFG).
Could it be added to ValueTracking's isKnownNonZero() or wrap around that?
It might already be in there under isKnownNonNullFromDominatingCondition(), but we need to fix it up to deal with selects?

[aqjune]

In D96663#2577600, @spatel wrote:

The code duplication suggests that we should have a common place for this logic (especially if it already exists in SimplifyCFG).
Could it be added to ValueTracking's isKnownNonZero() or wrap around that?
It might already be in there under isKnownNonNullFromDominatingCondition(), but we need to fix it up to deal with selects?

Hmm, as you said, adding the "assume((select cond, null, p) != null)" pattern to isKnownNonZeroFromAssume seems enough?
Compiling the Rust code snippet from pr48975 already creates llvm.assume: https://godbolt.org/z/8KxWKP
That being said, it might be a good idea to update passingValueIsAlwaysUndefined to check CB->paramHasAttr(ArgIdx, Attribute::Dereferenceable) as well because dereferenceable is quite common. dereferenceable implies noundef, so it is valid to assume that passing null to nonnull dereferenceable is UB. What do you think?

[spatel]

In D96663#2578084, @aqjune wrote:

In D96663#2577600, @spatel wrote:

The code duplication suggests that we should have a common place for this logic (especially if it already exists in SimplifyCFG).
Could it be added to ValueTracking's isKnownNonZero() or wrap around that?
It might already be in there under isKnownNonNullFromDominatingCondition(), but we need to fix it up to deal with selects?

Hmm, as you said, adding the "assume((select cond, null, p) != null)" pattern to isKnownNonZeroFromAssume seems enough?
Compiling the Rust code snippet from pr48975 already creates llvm.assume: https://godbolt.org/z/8KxWKP
That being said, it might be a good idea to update passingValueIsAlwaysUndefined to check CB->paramHasAttr(ArgIdx, Attribute::Dereferenceable) as well because dereferenceable is quite common. dereferenceable implies noundef, so it is valid to assume that passing null to nonnull dereferenceable is UB. What do you think?

If we can make the value tracking analysis better without adding or relying on llvm.assume, that would be best. I'm not sure about the current thinking on assume intrinsics, but we had found cases where the extra instructions/uses were interfering with optimization rather than helping. We've also seen cases like https://llvm.org/PR49171 where the cost of analyzing assumes is extreme. Not sure if those problems are overcome by assume bundles:
https://llvm.org/docs/LangRef.html#assume-operand-bundles

[aqjune]

ATM I don't have a clear idea about how to implement it in ValueTracking without relying on assume..

BTW, before diving into ValueTracking's update, I made D97244 to make SimplifyCFG slightly stronger. As @ojeda mentioned in the bug report, dereferenceable implies noundef, and the attribute is pretty common, so I think it is a good idea to update this first and see how it goes.
@ojeda Is it possible to compile Rust using a custom LLVM build? Then, maybe I can play with this a bit.

[ojeda]

Is it possible to compile Rust using a custom LLVM build?

Yeah, it is, although currently it is based on LLVM 11. For LLVM 12 support, take a look at https://github.com/rust-lang/rust/pull/81451. For LLVM 13, you probably need further tweaks.

When you configure Rust for building, use options like:

[llvm]
download-ci-llvm = false
skip-rebuild = true

[target.x86_64-unknown-linux-gnu]
llvm-config = "../your-custom-llvm-build/bin/llvm-config"

In general, take a look at config.toml.example, with src/bootstrap/defaults/config.codegen.toml as a start.

The guide is at https://rustc-dev-guide.rust-lang.org/building/how-to-build-and-run.html

[aqjune]

Make foldICmpInstWithConstantNotInt fold (select cond, null, arg) == null if arg has nonnullattr instead

[aqjune]

Add a negative test

[ojeda]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

3246

[aqjune]

Instead of adding two transformations, I tweaked InstCombinerImpl::foldICmpInstWithConstantNotInt so it fires if icmp is comparing to null and select's true/false operands are either a constant or nonnullattr argument.
It allows:

  %.0.i = select i1 %2, i8* null, i8* %x
  %4 = icmp ne i8* %.0.i, null
  call void @llvm.assume(i1 %4)
  ret i8* %.0.1
=>
  %.0.i = select i1 %2, i8* null, i8* %x
  %temp = xor i1 %2, 1
  call void @llvm.assume(i1 %temp)
  ret i8* %.0.1
=>  
  ...
  ret i8* %x

BTW, fixing SimplifyCFG to replace phi operand with null doesn't work for the Rust example because phi is already folded into select before inlining. It roughly looks like this: https://godbolt.org/z/c38qon

[aqjune]

Fix typo

[aqjune]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

3281

To support example_negative, this needed to be relaxed; I chose a safer direction that would less cause possible performance issues.

[xbolva00]

In D96663#2577600, @spatel wrote:

The code duplication suggests that we should have a common place for this logic (especially if it already exists in SimplifyCFG).
Could it be added to ValueTracking's isKnownNonZero() or wrap around that?
It might already be in there under isKnownNonNullFromDominatingCondition(), but we need to fix it up to deal with selects?

I dont think this could be unified so easily.

In D96663#2583945, @aqjune wrote:

Make foldICmpInstWithConstantNotInt fold (select cond, null, arg) == null if arg has nonnullattr instead

But original fold was more powerful and general IMHO.

[aqjune]

In D96663#2584000, @xbolva00 wrote:

In D96663#2583945, @aqjune wrote:

Make foldICmpInstWithConstantNotInt fold (select cond, null, arg) == null if arg has nonnullattr instead

But original fold was more powerful and general IMHO.

I think it is hard to tell which version is more general - the previous patch was explicitly relying on uses that have nonnull attributes whereas the new one relies on assume calls.
This patch has more power if there are many assume calls.

[aqjune]

Ping - I think this approach is a good direction to fix the bug without increasing the depth of uses to search.

[aqjune]

ping

[aqjune]

ping

[xbolva00]

In D96663#2594106, @aqjune wrote:

Ping - I think this approach is a good direction to fix the bug without increasing the depth of uses to search.

Yes. I think this is ready to go and already sits here for a long time. Wait a day for additional comments if any.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

3210

As this is already on InstCombinerImpl, is the Q argument necessary?

llvm/test/Transforms/InstCombine/assume-icmp-null-select.ll

20

I find the presence of assume in this example a bit confusing, as your transform does not actually use information from the assume. I'd replace it with a more generic "use".

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

3248

I don't really like this non-null attr specific case. I would suggest to simply always use SimplifyICmpInst for the isNullValue() case (and add a test for some non-attribute example).

I checked that doing so doesn't have any visible impact on CTMark (https://llvm-compile-time-tracker.com/compare.php?from=7baa2392fba0193c3e6631376c8e6a3c0318a743&to=45c227240929d8139da3dbee517e03d706555f23&stat=instructions).

[aqjune]

Address comments

[aqjune]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

3210

Removed it, thanks

llvm/test/Transforms/InstCombine/assume-icmp-null-select.ll

20

llvm.assume is necessary because it tells that %y_is_null is false, helping ret i8* %res to be transformed into ret i8* %x.

The series of transformations looks like this:

(1) InstCombine visits %nonnull = icmp ne i8* %res, null
(2) foldICmpInstWithConstantNotInt looks into the definition of %res which is select i1 %y_is_null, i8* null, i8* %x and folds %nonnull into select i1 %y_is_null, false, true.
During this folding, %x's dereferenceable tag is used by ValueTracking to answer isNonNull query.
(3) The llvm.assume now tells that %y_is_null is false, allowing the ret i8* %res to be transformed into ret i8* %x.

[aqjune]

ping

As suggested, the argument's nonnull check is fixed to use SimplifyICmpInst instead
May I land this patch?

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

 //===- InstCombineCompares.cpp --------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
Instruction *InstCombinerImpl::foldICmpIntrinsicWithConstant(ICmpInst &Cmp,
   default:
     break;
   }
 
   return nullptr;
 }
 
 /// Handle icmp with constant (but not simple integer constant) RHS.
-Instruction *InstCombinerImpl::foldICmpInstWithConstantNotInt(ICmpInst &I) {
+Instruction *
+InstCombinerImpl::foldICmpInstWithConstantNotInt(ICmpInst &I,
+                                                 const SimplifyQuery &Q) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   Constant *RHSC = dyn_cast<Constant>(Op1);
   Instruction *LHSI = dyn_cast<Instruction>(Op0);
   if (!RHSC || !LHSI)
     return nullptr;
 
   switch (LHSI->getOpcode()) {
   case Instruction::GetElementPtr:
if (LHSI->getParent() == I.getParent())
         return NV;
     break;
   case Instruction::Select: {
     // If either operand of the select is a constant, we can fold the
     // comparison into the select arms, which will cause one to be
     // constant folded and the select turned into a bitwise or.
     Value *Op1 = nullptr, *Op2 = nullptr;
     ConstantInt *CI = nullptr;
-    if (Constant *C = dyn_cast<Constant>(LHSI->getOperand(1))) {
-      Op1 = ConstantExpr::getICmp(I.getPredicate(), C, RHSC);
-      CI = dyn_cast<ConstantInt>(Op1);
+
+    auto SimplifyOp = [&](Value *V) {
+      Value *Op = nullptr;
+      if (Constant *C = dyn_cast<Constant>(V)) {
+        Op = ConstantExpr::getICmp(I.getPredicate(), C, RHSC);
+      } else if (RHSC->isNullValue()) {
+        // If null is being compared, check if V is a nonnull argument.
+        Argument *A = dyn_cast<Argument>(V);
+        if (A && A->hasNonNullAttr(true))
+          Op = SimplifyICmpInst(I.getPredicate(), A, RHSC, Q);
       }
-    if (Constant *C = dyn_cast<Constant>(LHSI->getOperand(2))) {
-      Op2 = ConstantExpr::getICmp(I.getPredicate(), C, RHSC);
+      return Op;
+    };
+    Op1 = SimplifyOp(LHSI->getOperand(1));
+    if (Op1)
+      CI = dyn_cast<ConstantInt>(Op1);
+
+    Op2 = SimplifyOp(LHSI->getOperand(2));
+    if (Op2)
       CI = dyn_cast<ConstantInt>(Op2);
-    }
 
     // We only want to perform this transformation if it will not lead to
     // additional code. This is true if either both sides of the select
     // fold to a constant (in which case the icmp is replaced with a select
     // which will usually simplify) or this is the only user of the
     // select (in which case we are trading a select+icmp for a simpler
     // select+icmp) or all uses of the select can be replaced based on
     // dominance information ("Global cases").
     bool Transform = false;
     if (Op1 && Op2)
       Transform = true;
     else if (Op1 || Op2) {
       // Local case
       if (LHSI->hasOneUse())
         Transform = true;
       // Global cases
       else if (CI && !CI->isZero())
         // When Op1 is constant try replacing select with second operand.
         // Otherwise Op2 is constant and try replacing select with first
         // operand.
         Transform =
             replacedSelectWithOperand(cast<SelectInst>(LHSI), &I, Op1 ? 2 : 1);
     }
     if (Transform) {
       if (!Op1)
         Op1 = Builder.CreateICmp(I.getPredicate(), LHSI->getOperand(1), RHSC,
                                  I.getName());
       if (!Op2)
         Op2 = Builder.CreateICmp(I.getPredicate(), LHSI->getOperand(2), RHSC,
                                  I.getName());
Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
   if (Instruction *Res = foldICmpInstWithConstant(I))
     return Res;
 
   // Try to match comparison as a sign bit test. Intentionally do this after
   // foldICmpInstWithConstant() to potentially let other folds to happen first.
   if (Instruction *New = foldSignBitTest(I))
     return New;
 
-  if (Instruction *Res = foldICmpInstWithConstantNotInt(I))
+  if (Instruction *Res = foldICmpInstWithConstantNotInt(I, Q))
     return Res;
 
   // If we can optimize a 'icmp GEP, P' or 'icmp P, GEP', do so now.
   if (GEPOperator *GEP = dyn_cast<GEPOperator>(Op0))
     if (Instruction *NI = foldGEPICmp(GEP, Op1, I.getPredicate(), I))
       return NI;
   if (GEPOperator *GEP = dyn_cast<GEPOperator>(Op1))
     if (Instruction *NI = foldGEPICmp(GEP, Op0,

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

 //===- InstCombineInternal.h - InstCombine pass internals -------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
public:
   Instruction *foldICmpAddOpConst(Value *X, const APInt &C,
                                   ICmpInst::Predicate Pred);
   Instruction *foldICmpWithCastOp(ICmpInst &ICI);
 
   Instruction *foldICmpUsingKnownBits(ICmpInst &Cmp);
   Instruction *foldICmpWithDominatingICmp(ICmpInst &Cmp);
   Instruction *foldICmpWithConstant(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstant(ICmpInst &Cmp);
-  Instruction *foldICmpInstWithConstantNotInt(ICmpInst &Cmp);
+  Instruction *foldICmpInstWithConstantNotInt(ICmpInst &Cmp,
+                                              const SimplifyQuery &SQ);
   Instruction *foldICmpBinOp(ICmpInst &Cmp, const SimplifyQuery &SQ);
   Instruction *foldICmpEquality(ICmpInst &Cmp);
   Instruction *foldIRemByPowerOfTwoToBitTest(ICmpInst &I);
   Instruction *foldSignBitTest(ICmpInst &I);
   Instruction *foldICmpWithZero(ICmpInst &Cmp);
 
   Value *foldUnsignedMultiplicationOverflowCheck(ICmpInst &Cmp);
 

llvm/test/Transforms/InstCombine/assume-icmp-null-select.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define align 8 dereferenceable(24) i8* @example(i8* readonly align 8 dereferenceable(24) %x) {
+; CHECK-LABEL: @example(
+; CHECK-NEXT:    [[DOT0:%.*]] = bitcast i8* [[X:%.*]] to {}**
+; CHECK-NEXT:    [[TMP1:%.*]] = load {}*, {}** [[DOT0]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne {}* [[TMP1]], null
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP2]])
+; CHECK-NEXT:    ret i8* [[X]]
+;
+  %.0 = bitcast i8* %x to {}**
+  %1 = load {}*, {}** %.0, align 8
+  %2 = icmp eq {}* %1, null
+  %3 = getelementptr inbounds i8, i8* %x, i64 0
+  %.0.i = select i1 %2, i8* null, i8* %3
+  %4 = icmp ne i8* %.0.i, null
+  call void @llvm.assume(i1 %4)
+  ret i8* %.0.i
+}
+
+; TODO: this should be optimized to 'ret %x' as well.
+define align 8 dereferenceable(24) i8* @example_negative(i8* readonly align 8 %x) {
+; CHECK-LABEL: @example_negative(
+; CHECK-NEXT:    [[DOT0:%.*]] = bitcast i8* [[X:%.*]] to {}**
+; CHECK-NEXT:    [[TMP1:%.*]] = load {}*, {}** [[DOT0]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq {}* [[TMP1]], null
+; CHECK-NEXT:    [[DOT0_I:%.*]] = select i1 [[TMP2]], i8* null, i8* [[X]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne i8* [[DOT0_I]], null
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP3]])
+; CHECK-NEXT:    ret i8* [[DOT0_I]]
+;
+  %.0 = bitcast i8* %x to {}**
+  %1 = load {}*, {}** %.0, align 8
+  %2 = icmp eq {}* %1, null
+  %3 = getelementptr inbounds i8, i8* %x, i64 0
+  %.0.i = select i1 %2, i8* null, i8* %3
+  %4 = icmp ne i8* %.0.i, null
+  call void @llvm.assume(i1 %4)
+  ret i8* %.0.i
+}
+
+declare void @llvm.assume(i1)