[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 ↗

(On Diff #325982)

[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 ↗

(On Diff #325985)

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 ↗

(On Diff #325985)

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

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

20 ↗

(On Diff #325985)

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 ↗

(On Diff #325985)

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 ↗

(On Diff #325985)

Removed it, thanks

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

20 ↗

(On Diff #325985)

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/InstCombineCalls.cpp

 //===- InstCombineCalls.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
 //
 //===----------------------------------------------------------------------===//
 //
if (!Call.use_empty() && !Call.isMustTailCall())
     if (Value *ReturnedArg = Call.getReturnedArgOperand()) {
       Type *CallTy = Call.getType();
       Type *RetArgTy = ReturnedArg->getType();
       if (RetArgTy->canLosslesslyBitCastTo(CallTy))
         return replaceInstUsesWith(
             Call, Builder.CreateBitOrPointerCast(ReturnedArg, CallTy));
     }
 
+  // Simplify nonnull arguments if possible.
+  //   ptr = select _, null, p
+  //   f(nonnull ptr)
+  // =>
+  //   f(nonnull p)
+  for (unsigned ArgNo = 0, N = Call.getNumArgOperands(); ArgNo != N; ++ArgNo) {
+    if (!Call.paramHasAttr(ArgNo, Attribute::NonNull))
+      continue;
+
+    Value *V = Call.getArgOperand(ArgNo);
+    Value *TrueVal, *FalseVal;
+    if (!match(V, m_Select(m_Value(), m_Value(TrueVal), m_Value(FalseVal))))
+      continue;
+
+    bool IsTrueNull = isa<ConstantPointerNull>(TrueVal);
+    if (IsTrueNull || isa<ConstantPointerNull>(FalseVal)) {
+      replaceUse(Call.getArgOperandUse(ArgNo), IsTrueNull ? FalseVal : TrueVal);
+      Changed = true;
+    }
+  }
+
   if (isAllocLikeFn(&Call, &TLI))
     return visitAllocSite(Call);
 
   return Changed ? &Call : nullptr;
 }
 
 /// If the callee is a constexpr cast of a function, attempt to move the cast to
 /// the arguments of the call/callbr/invoke.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

 //===- InstructionCombining.cpp - Combine multiple instructions -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
static bool isMustTailCall(Value *V) {
   return false;
 }
 
 Instruction *InstCombinerImpl::visitReturnInst(ReturnInst &RI) {
   if (RI.getNumOperands() == 0) // ret void
     return nullptr;
 
   Value *ResultOp = RI.getOperand(0);
+
+  auto Attrs = RI.getFunction()->getAttributes();
+  if (Attrs.hasAttribute(AttributeList::ReturnIndex, Attribute::NonNull)) {
+    //   ptr = select _, null, p
+    //   ret ptr
+    // =>
+    //   ret p
+    Value *TrueVal, *FalseVal;
+    if (match(ResultOp, m_Select(m_Value(), m_Value(TrueVal), m_Value(FalseVal)))) {
+      bool IsTrueNull = isa<ConstantPointerNull>(TrueVal);
+      if (IsTrueNull || isa<ConstantPointerNull>(FalseVal))
+        return replaceOperand(RI, 0, IsTrueNull ? FalseVal : TrueVal);
+    }
+  }
+
   Type *VTy = ResultOp->getType();
   if (!VTy->isIntegerTy() || isa<Constant>(ResultOp))
     return nullptr;
 
   // Don't replace result of musttail calls.
   if (isMustTailCall(ResultOp))
     return nullptr;
 

llvm/test/Transforms/InstCombine/nonnull-select.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; Passing select(cond, null, v) as nonnull should be optimized to passing v
 
 define nonnull i32* @pr48975(i32** %.0) {
 ; CHECK-LABEL: @pr48975(
-; CHECK-NEXT:    [[DOT1:%.*]] = load i32*, i32** [[DOT0:%.*]], align 8
-; CHECK-NEXT:    [[DOT2:%.*]] = icmp eq i32* [[DOT1]], null
-; CHECK-NEXT:    [[DOT3:%.*]] = bitcast i32** [[DOT0]] to i32*
-; CHECK-NEXT:    [[DOT4:%.*]] = select i1 [[DOT2]], i32* null, i32* [[DOT3]]
-; CHECK-NEXT:    ret i32* [[DOT4]]
+; CHECK-NEXT:    [[DOT3:%.*]] = bitcast i32** [[DOT0:%.*]] to i32*
+; CHECK-NEXT:    ret i32* [[DOT3]]
 ;
   %.1 = load i32*, i32** %.0, align 8
   %.2 = icmp eq i32* %.1, null
   %.3 = bitcast i32** %.0 to i32*
   %.4 = select i1 %.2, i32* null, i32* %.3
   ret i32* %.4
 }
 
 define nonnull i32* @nonnull_ret(i1 %cond, i32* %p) {
 ; CHECK-LABEL: @nonnull_ret(
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[COND:%.*]], i32* [[P:%.*]], i32* null
-; CHECK-NEXT:    ret i32* [[RES]]
+; CHECK-NEXT:    ret i32* [[P:%.*]]
 ;
   %res = select i1 %cond, i32* %p, i32* null
   ret i32* %res
 }
 
 define void @nonnull_call(i1 %cond, i32* %p) {
 ; CHECK-LABEL: @nonnull_call(
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[COND:%.*]], i32* null, i32* [[P:%.*]]
-; CHECK-NEXT:    call void @f(i32* nonnull [[RES]])
+; CHECK-NEXT:    call void @f(i32* nonnull [[P:%.*]])
 ; CHECK-NEXT:    ret void
 ;
   %res = select i1 %cond, i32* null, i32* %p
   call void @f(i32* nonnull %res)
   ret void
 }
 
 declare void @f(i32*)