[Inliner] Preserve nonnull from callsites during inlining

[xbolva00]

Summary

Inliner fails to simplify inlined function based on known callsite attributes

declare void @h(ptr %p)

define void @f(ptr %p) {
	call void @g(ptr nonnull %p)
	ret void
}

define void @g(ptr %p) {
	call void @h(ptr %p)
	ret void
}

Inliner inlines g into f but we lose extra info that %p was nonnull ptr. Implemented suggested solution, so insert nonnull assumes in Inliner to preserve knowledge.

Fixes https://github.com/llvm/llvm-project/issues/61966

[aeubanks]

what's the compile time impact of this?

llvm/lib/Transforms/Utils/InlineFunction.cpp

1475

are there tests for isKnownNonZero?

llvm/test/Transforms/Inline/assumptions-from-callsite-attrs.ll

3

redundant RUN line

4

I wouldn't worry about the module inliner, it's still experimental

llvm/test/Transforms/PhaseOrdering/dae-dce.ll

22

isn't this now UB?

[nikic]

llvm/lib/Transforms/Utils/InlineFunction.cpp

1475

You need both nonnull and noundef to convert to assume.

[nikic]

llvm/test/Transforms/PhaseOrdering/dae-dce.ll

22

I checked, and this is also caused by not respecting noundef. The noundef on the call is dropped by the time of inlining.

[xbolva00]

Added check 'CB.getAttributes().hasParamAttr(ArgNo, Attribute::NoUndef)'. The impact of this PR is non-trivial:
https://llvm-compile-time-tracker.com/compare.php?from=06ddb7bfe24891eee43bcf19252708b0cf684562&to=d1559cc73ad38f7684a026f8c69c6c96af9246a3&stat=instructions:u

especially for tramp3d-v4, where we can also observe a bigger codesize change.

[aeubanks]

yeah that compile time increase is unacceptable

I've run into related issues where we lose the return value alignment after inlining, as pointed out in https://reviews.llvm.org/D123052

[xbolva00]

alignment info generally hits same issue

https://llvm.org/doxygen/InlineFunction_8cpp.html#a4224e0e010bb5413dc74243c276724f5

Disabled by default, because the added alignment assumptions may increase
compile-time and block optimizations. This option is not suitable for use
with frontends that emit comprehensive parameter alignment annotations.

So ideally we should have cheap and more general solution for this problem. Any ideas?

[efriedma]

Something "cheap" we could do is attach nonnull attributes to arguments of inlined calls, when appropriate.

llvm.assume is fundamentally pretty expensive; it bloats use-lists, and there's additional instructions feeding into the assume. There are some aspects we could make cheaper with an optimized form of llvm.assume, but it's not easy.

[xbolva00]

Something "cheap" we could do is attach nonnull attributes to arguments of inlined calls, when appropriate.

Lets assume following example:
https://godbolt.org/z/TY9h1PYqY

Clearly we fail to eliminate the "is not null" comparison and we have the attached nonnull attribute on the passed argument.
Can you please explain your approach more in details?

[nikic]

@xbolva00 Could you please explain what the larger context here is? Which frontend generates nonnull on call-site only?

[xbolva00]

In D147704#4322809, @nikic wrote:

@xbolva00 Could you please explain what the larger context here is? Which frontend generates nonnull on call-site only?

I was working on something else and I then found testcases in https://github.com/llvm/llvm-project/blob/ca42cd3e1264f8b44304019da07b6059a610fd24/llvm/test/Transforms/Inline/nonnull.ll#L75 and I was experimenting why some combinations of nonnull attribute placements work and some not. I am not aware of such frontend, indeed, but there are many LLVM-based compilers, so who knows..

I was wondering if a frontend or midend passes in some cases / contexts could add additional nonnull call site attribute to preserve knowledge. Currently I am only aware of added attributes on call sites with C libcalls by SimplifyLibCalls, so probably I have no real motivation case.

llvm/include/llvm/IR/IRBuilder.h

public:
   /// must be subtracted from the provided pointer to get the pointer with the
   /// specified alignment.
   ///
   /// This overload handles the condition where the Alignment is dependent
   /// on an existing value rather than a static value.
   CallInst *CreateAlignmentAssumption(const DataLayout &DL, Value *PtrValue,
                                       Value *Alignment,
                                       Value *OffsetValue = nullptr);
+  /// Create an assume intrinsic call that represents a nonnull
+  /// assumption on the provided pointer.
+  CallInst *CreateNonNullAssumption(Value *PtrValue);
 };
 
 /// This provides a uniform API for creating instructions and inserting
 /// them into a basic block: either at the end of a BasicBlock, or at a specific
 /// iterator location in a block.
 ///
 /// Note that the builder does not expose the full generality of LLVM
 /// instructions.  For access to extra instruction properties, use the mutators

llvm/lib/IR/IRBuilder.cpp

CallInst *IRBuilderBase::CreateAlignmentAssumption(const DataLayout &DL,
                                                    Value *PtrValue,
                                                    Value *Alignment,
                                                    Value *OffsetValue) {
   assert(isa<PointerType>(PtrValue->getType()) &&
          "trying to create an alignment assumption on a non-pointer?");
   return CreateAlignmentAssumptionHelper(DL, PtrValue, Alignment, OffsetValue);
 }
 
+CallInst *IRBuilderBase::CreateNonNullAssumption(Value *PtrValue) {
+  Value *Vals[] = { PtrValue };
+  OperandBundleDefT<Value *> NonNullOpB("nonnull", Vals);
+  return CreateAssumption(ConstantInt::getTrue(getContext()), {NonNullOpB});
+}
+
 IRBuilderDefaultInserter::~IRBuilderDefaultInserter() = default;
 IRBuilderCallbackInserter::~IRBuilderCallbackInserter() = default;
 IRBuilderFolder::~IRBuilderFolder() = default;
 void ConstantFolder::anchor() {}
 void NoFolder::anchor() {}

llvm/lib/Transforms/Utils/InlineFunction.cpp

if (getKnownAlignment(ArgVal, DL, &CB, AC, &DT) >= *Alignment)
       continue;
 
     CallInst *NewAsmp = IRBuilder<>(&CB).CreateAlignmentAssumption(
         DL, ArgVal, Alignment->value());
     AC->registerAssumption(cast<AssumeInst>(NewAsmp));
   }
 }
 
+static void AddAssumptionsFromCallSiteAttrs(CallBase &CB, InlineFunctionInfo &IFI) {
+  if (!IFI.GetAssumptionCache)
+    return;
+
+  AssumptionCache *AC = &IFI.GetAssumptionCache(*CB.getCaller());
+  auto &DL = CB.getCaller()->getParent()->getDataLayout();
+  Function *CalledFunc = CB.getCalledFunction();
+  IRBuilder<> Builder(&CB);
+
+  for (Argument &Arg : CalledFunc->args()) {
+    unsigned ArgNo = Arg.getArgNo();
+    auto *ArgVal = CB.getArgOperand(ArgNo);
+    if (CB.getAttributes().hasParamAttr(ArgNo, Attribute::NonNull) && !isKnownNonZero(ArgVal, DL, 0, AC)) {
+      CallInst *NewAsmp = Builder.CreateNonNullAssumption(CB.getArgOperand(ArgNo));
+      AC->registerAssumption(cast<AssumeInst>(NewAsmp));
+    }
+  }
+}
+
 static void HandleByValArgumentInit(Type *ByValType, Value *Dst, Value *Src,
                                     Module *M, BasicBlock *InsertBlock,
                                     InlineFunctionInfo &IFI,
                                     Function *CalledFunc) {
   IRBuilder<> Builder(InsertBlock, InsertBlock->begin());
 
   Value *Size =
       Builder.getInt64(M->getDataLayout().getTypeStoreSize(ByValType));
for (Function::arg_iterator I = CalledFunc->arg_begin(),
         if (ActualArg != *AI)
           ByValInits.push_back(
               {ActualArg, (Value *)*AI, CB.getParamByValType(ArgNo)});
       }
 
       VMap[&*I] = ActualArg;
     }
 
+    AddAssumptionsFromCallSiteAttrs(CB, IFI);
+
     // TODO: Remove this when users have been updated to the assume bundles.
     // Add alignment assumptions if necessary. We do this before the inlined
     // instructions are actually cloned into the caller so that we can easily
     // check what will be known at the start of the inlined code.
     AddAlignmentAssumptions(CB, IFI);
 
     AssumptionCache *AC =
         IFI.GetAssumptionCache ? &IFI.GetAssumptionCache(*Caller) : nullptr;

llvm/test/Transforms/Inline/assumptions-from-callsite-attrs.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -S -passes=inline %s | FileCheck %s
 ; RUN: opt -S -passes='cgscc(inline)' %s | FileCheck %s
 ; RUN: opt -S -passes='module-inline' %s | FileCheck %s
 
 declare void @h(ptr %p, ptr %q, ptr %z)
 
 define void @f(ptr %p, ptr %q, ptr %z) {
 ; CHECK-LABEL: define void @f
 ; CHECK-SAME: (ptr [[P:%.*]], ptr [[Q:%.*]], ptr [[Z:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[P]]) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[Z]]) ]
 ; CHECK-NEXT:    call void @h(ptr [[P]], ptr [[Q]], ptr [[Z]])
 ; CHECK-NEXT:    ret void
 ;
   call void @g(ptr nonnull %p, ptr %q, ptr nonnull %z)
   ret void
 }
 
 define void @g(ptr %p, ptr %q, ptr %z) {
 ; CHECK-LABEL: define void @g
 ; CHECK-SAME: (ptr [[P:%.*]], ptr [[Q:%.*]], ptr [[Z:%.*]]) {
 ; CHECK-NEXT:    call void @h(ptr [[P]], ptr [[Q]], ptr [[Z]])
 ; CHECK-NEXT:    ret void
 ;
   call void @h(ptr %p, ptr %q, ptr %z)
   ret void
 }

llvm/test/Transforms/Inline/nonnull.ll

 ;
   call void @callee(ptr %arg)
   ret void
 }
 
 define void @caller3(ptr %arg) {
 ; CHECK-LABEL: define void @caller3
 ; CHECK-SAME: (ptr [[ARG:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[ARG]]) ]
 ; CHECK-NEXT:    [[CMP_I:%.*]] = icmp eq ptr [[ARG]], null
 ; CHECK-NEXT:    br i1 [[CMP_I]], label [[EXPENSIVE_I:%.*]], label [[DONE_I:%.*]]
 ; CHECK:       expensive.i:
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    br label [[CALLEE_EXIT:%.*]]
 ; CHECK:       done.i:
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    br label [[CALLEE_EXIT]]
 ; CHECK:       callee.exit:
 ; CHECK-NEXT:    ret void
 ;
   call void @callee(ptr nonnull %arg)
   ret void
 }
 
+; Positive test - arg is known non null
 define void @caller4(ptr dereferenceable(8) %arg) {
 ; CHECK-LABEL: define void @caller4
 ; CHECK-SAME: (ptr dereferenceable(8) [[ARG:%.*]]) {
 ; CHECK-NEXT:    call void @callee(ptr dereferenceable(8) [[ARG]])
 ; CHECK-NEXT:    ret void
 ;
   call void @callee(ptr dereferenceable(8) %arg)
   ret void

llvm/test/Transforms/PhaseOrdering/dae-dce.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes='default<O3>' < %s | FileCheck %s --check-prefixes=CHECK,DEFAULT
 ; RUN: opt -S -passes='lto<O3>' < %s | FileCheck %s --check-prefixes=CHECK,LTO
 
 declare void @llvm.trap()
 
 define void @do_trap(ptr %ptr) {
 ; CHECK-LABEL: @do_trap(
 ; CHECK-NEXT:    tail call void @llvm.trap()
 ; CHECK-NEXT:    unreachable
 ;
   call void @llvm.trap()
   unreachable
 }
 
 define internal void @capture_and_trap(ptr %ptr) noinline {
-; CHECK-LABEL: @capture_and_trap(
-; CHECK-NEXT:    tail call void @llvm.trap()
-; CHECK-NEXT:    unreachable
+; DEFAULT-LABEL: @capture_and_trap(
+; DEFAULT-NEXT:    tail call void @llvm.trap()
+; DEFAULT-NEXT:    unreachable
+;
+; LTO-LABEL: @capture_and_trap(
+; LTO-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr poison) ]
+; LTO-NEXT:    tail call void @llvm.trap()
+; LTO-NEXT:    unreachable
 ;
   %alloca = alloca ptr, align 4
   store ptr %ptr, ptr %alloca, align 4
   call void @do_trap(ptr noundef nonnull %alloca)
   unreachable
 }
 
 define internal void @dead_fn1() {

 if:
   call void @capture_and_trap(ptr @dead_fn1)
   unreachable
 
 else:
   call void @capture_and_trap(ptr @dead_fn2)
   unreachable
 }
-;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
-; DEFAULT: {{.*}}
-; LTO: {{.*}}