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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: None (goldsteinn)

<details>
<summary>Changes</summary>

- **[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.

---
Full diff: https://github.com/llvm/llvm-project/pull/109347.diff


3 Files Affected:

- (modified) llvm/lib/Transforms/Utils/InlineFunction.cpp (+12) 
- (modified) llvm/test/Transforms/Inline/access-attributes-prop.ll (+21) 
- (modified) llvm/test/Transforms/Inline/ret_attr_align_and_noundef.ll (+40) 


``````````diff
diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 52da65ce01b82b..53e25e53c2a86f 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1383,6 +1383,12 @@ 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. Only propagate return attributes if we are in fact calling the
+      // same function.
+      if (InnerCB->getCalledFunction() != NewInnerCB->getCalledFunction())
+        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.
@@ -1477,6 +1483,12 @@ 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. Only propagate return attributes if we are in fact calling the
+    // same function.
+    if (RetVal->getCalledFunction() != NewRetVal->getCalledFunction())
+      continue;
     // Backward propagation of attributes to the returned value may be incorrect
     // if it is control flow dependent.
     // Consider:
diff --git a/llvm/test/Transforms/Inline/access-attributes-prop.ll b/llvm/test/Transforms/Inline/access-attributes-prop.ll
index 965f0335b63eab..2c55f5f3b1f6ca 100644
--- a/llvm/test/Transforms/Inline/access-attributes-prop.ll
+++ b/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
+}
diff --git a/llvm/test/Transforms/Inline/ret_attr_align_and_noundef.ll b/llvm/test/Transforms/Inline/ret_attr_align_and_noundef.ll
index dc685d2c4d1368..930bef43df1dba 100644
--- a/llvm/test/Transforms/Inline/ret_attr_align_and_noundef.ll
+++ b/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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/109347