[Inliner] Fix bug where attributes are propagated incorrectly (#109347)

- **[Inliner] Add tests for incorrect propagation of return attrs; NFC**
- **[Inliner] Fix bug where attributes are propagated incorrectly**

The bug stems from the fact that we assume the new (inlined) callsite
is calling the same function as the original (callee) callsite. While
this is typically the case, since `VMap` simplifies the new
instructions, callee intrinsics callsites can end up not corresponding
with the same function.

This can lead to buggy propagation.

Author: goldsteinn

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -1352,7 +1352,8 @@ static bool MayContainThrowingOrExitingCallAfterCB(CallBase *Begin,
 // Add attributes from CB params and Fn attributes that can always be propagated
 // to the corresponding argument / inner callbases.
 static void AddParamAndFnBasicAttributes(const CallBase &CB,
-                                         ValueToValueMapTy &VMap) {
+                                         ValueToValueMapTy &VMap,
+                                         ClonedCodeInfo &InlinedFunctionInfo) {
   auto *CalledFunction = CB.getCalledFunction();
   auto &Context = CalledFunction->getContext();
 
@@ -1383,6 +1384,11 @@ static void AddParamAndFnBasicAttributes(const CallBase &CB,
       auto *NewInnerCB = dyn_cast_or_null<CallBase>(VMap.lookup(InnerCB));
       if (!NewInnerCB)
         continue;
+      // The InnerCB might have be simplified during the inlining
+      // process which can make propagation incorrect.
+      if (InlinedFunctionInfo.isSimplified(InnerCB, NewInnerCB))
+        continue;
+
       AttributeList AL = NewInnerCB->getAttributes();
       for (unsigned I = 0, E = InnerCB->arg_size(); I < E; ++I) {
         // Check if the underlying value for the parameter is an argument.
@@ -1458,7 +1464,8 @@ static AttrBuilder IdentifyValidPoisonGeneratingAttributes(CallBase &CB) {
   return Valid;
 }
 
-static void AddReturnAttributes(CallBase &CB, ValueToValueMapTy &VMap) {
+static void AddReturnAttributes(CallBase &CB, ValueToValueMapTy &VMap,
+                                ClonedCodeInfo &InlinedFunctionInfo) {
   AttrBuilder ValidUB = IdentifyValidUBGeneratingAttributes(CB);
   AttrBuilder ValidPG = IdentifyValidPoisonGeneratingAttributes(CB);
   if (!ValidUB.hasAttributes() && !ValidPG.hasAttributes())
@@ -1477,6 +1484,11 @@ static void AddReturnAttributes(CallBase &CB, ValueToValueMapTy &VMap) {
     auto *NewRetVal = dyn_cast_or_null<CallBase>(VMap.lookup(RetVal));
     if (!NewRetVal)
       continue;
+
+    // The RetVal might have be simplified during the inlining
+    // process which can make propagation incorrect.
+    if (InlinedFunctionInfo.isSimplified(RetVal, NewRetVal))
+      continue;
     // Backward propagation of attributes to the returned value may be incorrect
     // if it is control flow dependent.
     // Consider:
@@ -2693,11 +2705,11 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
 
     // Clone return attributes on the callsite into the calls within the inlined
     // function which feed into its return value.
-    AddReturnAttributes(CB, VMap);
+    AddReturnAttributes(CB, VMap, InlinedFunctionInfo);
 
     // Clone attributes on the params of the callsite to calls within the
     // inlined function which use the same param.
-    AddParamAndFnBasicAttributes(CB, VMap);
+    AddParamAndFnBasicAttributes(CB, VMap, InlinedFunctionInfo);
 
     propagateMemProfMetadata(CalledFunc, CB,
                              InlinedFunctionInfo.ContainsMemProfMetadata, VMap);

llvm/test/Transforms/Inline/access-attributes-prop.ll

@@ -559,3 +559,24 @@ define void @prop_byval_readonly(ptr %p) {
   call void @foo_byval_readonly(ptr %p)
   ret void
 }
+
+define ptr @caller_bad_param_prop(ptr %p1, ptr %p2, i64 %x) {
+; CHECK-LABEL: define {{[^@]+}}@caller_bad_param_prop
+; CHECK-SAME: (ptr [[P1:%.*]], ptr [[P2:%.*]], i64 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call ptr [[P1]](i64 [[X]], ptr [[P2]])
+; CHECK-NEXT:    ret ptr [[TMP1]]
+;
+  %1 = call ptr %p1(i64 %x, ptr %p2)
+  %2 = call ptr @callee_bad_param_prop(ptr %1)
+  ret ptr %2
+}
+
+define ptr @callee_bad_param_prop(ptr readonly %x) {
+; CHECK-LABEL: define {{[^@]+}}@callee_bad_param_prop
+; CHECK-SAME: (ptr readonly [[X:%.*]]) {
+; CHECK-NEXT:    [[R:%.*]] = tail call ptr @llvm.ptrmask.p0.i64(ptr [[X]], i64 -1)
+; CHECK-NEXT:    ret ptr [[R]]
+;
+  %r = tail call ptr @llvm.ptrmask(ptr %x, i64 -1)
+  ret ptr %r
+}

llvm/test/Transforms/Inline/ret_attr_align_and_noundef.ll

@@ -421,3 +421,43 @@ define i8 @caller16_not_intersecting_ranges() {
   call void @use.val(i8 %r)
   ret i8 %r
 }
+
+
+define ptr @caller_bad_ret_prop(ptr %p1, ptr %p2, i64 %x, ptr %other) {
+; CHECK-LABEL: define ptr @caller_bad_ret_prop
+; CHECK-SAME: (ptr [[P1:%.*]], ptr [[P2:%.*]], i64 [[X:%.*]], ptr [[OTHER:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call noundef ptr [[P1]](i64 [[X]], ptr [[P2]])
+; CHECK-NEXT:    [[CMP_I:%.*]] = icmp eq ptr [[TMP1]], null
+; CHECK-NEXT:    br i1 [[CMP_I]], label [[T_I:%.*]], label [[F_I:%.*]]
+; CHECK:       T.i:
+; CHECK-NEXT:    br label [[CALLEE_BAD_RET_PROP_EXIT:%.*]]
+; CHECK:       F.i:
+; CHECK-NEXT:    br label [[CALLEE_BAD_RET_PROP_EXIT]]
+; CHECK:       callee_bad_ret_prop.exit:
+; CHECK-NEXT:    [[TMP2:%.*]] = phi ptr [ [[OTHER]], [[T_I]] ], [ [[TMP1]], [[F_I]] ]
+; CHECK-NEXT:    ret ptr [[TMP2]]
+;
+  %1 = call noundef ptr %p1(i64 %x, ptr %p2)
+  %2 = call nonnull ptr @callee_bad_ret_prop(ptr %1, ptr %other)
+  ret ptr %2
+}
+
+define ptr @callee_bad_ret_prop(ptr %x, ptr %other) {
+; CHECK-LABEL: define ptr @callee_bad_ret_prop
+; CHECK-SAME: (ptr [[X:%.*]], ptr [[OTHER:%.*]]) {
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[X]], null
+; CHECK-NEXT:    br i1 [[CMP]], label [[T:%.*]], label [[F:%.*]]
+; CHECK:       T:
+; CHECK-NEXT:    ret ptr [[OTHER]]
+; CHECK:       F:
+; CHECK-NEXT:    [[R:%.*]] = tail call ptr @llvm.ptrmask.p0.i64(ptr [[X]], i64 -1)
+; CHECK-NEXT:    ret ptr [[R]]
+;
+  %cmp = icmp eq ptr %x, null
+  br i1 %cmp, label %T, label %F
+T:
+  ret ptr %other
+F:
+  %r = tail call ptr @llvm.ptrmask(ptr %x, i64 -1)
+  ret ptr %r
+}