xxx

llvm · Apr 6, 2023 · f605924 · f605924
1 parent 3b5ff3a
commit f605924
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 0 deletions.
diff --git a/llvm/include/llvm/IR/IRBuilder.h b/llvm/include/llvm/IR/IRBuilder.h
@@ -2543,6 +2543,9 @@ class IRBuilderBase {
   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

diff --git a/llvm/lib/IR/IRBuilder.cpp b/llvm/lib/IR/IRBuilder.cpp
@@ -1410,6 +1410,12 @@ CallInst *IRBuilderBase::CreateAlignmentAssumption(const DataLayout &DL,
   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;

diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1460,6 +1460,23 @@ static void AddAlignmentAssumptions(CallBase &CB, InlineFunctionInfo &IFI) {
   }
 }
 
+static void AddAssumptionsFromCallSiteAttrs(CallBase &CB, InlineFunctionInfo &IFI) {
+  if (!IFI.GetAssumptionCache)
+    return;
+
+  AssumptionCache *AC = &IFI.GetAssumptionCache(*CB.getCaller());
+  Function *CalledFunc = CB.getCalledFunction();
+  IRBuilder<> Builder(&CB);
+
+  for (Argument &Arg : CalledFunc->args()) {
+    unsigned ArgNo = Arg.getArgNo();
+    if (CB.paramHasAttr(ArgNo, Attribute::NonNull)) {
+      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,
@@ -2130,6 +2147,8 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
       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

diff --git a/llvm/test/Transforms/Inline/assumptions-from-callsite-attrs.ll b/llvm/test/Transforms/Inline/assumptions-from-callsite-attrs.ll
@@ -8,6 +8,8 @@ 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
 ;

diff --git a/llvm/test/Transforms/Inline/nonnull.ll b/llvm/test/Transforms/Inline/nonnull.ll
@@ -54,6 +54,7 @@ done:
 define void @caller(ptr nonnull %arg) {
 ; CHECK-LABEL: define void @caller
 ; CHECK-SAME: (ptr nonnull [[ARG:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[ARG]]) ]
 ; CHECK-NEXT:    call void @bar()
 ; CHECK-NEXT:    ret void
 ;
@@ -75,6 +76,7 @@ define void @caller2(ptr %arg) {
 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: