[IPO] Prevent removal of some convergent attr

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -2902,6 +2902,22 @@ struct AANonConvergentImpl : public AANonConvergent {
   }
 };
 
+static bool FunctionCalledWithConvergenceToken(const Function *F) {
+  for (auto &Use : F->uses()) {
+    CallBase *CB = dyn_cast<CallBase>(Use.getUser());
+    if (!CB)
+      continue;
+
+    // We are not called, just used as an argument.
+    if (CB->getCalledFunction() != F)
+      continue;
+
+    if (CB->getConvergenceControlToken())
+      return true;
+  }
+  return false;
+}
+
 struct AANonConvergentFunction final : AANonConvergentImpl {
   AANonConvergentFunction(const IRPosition &IRP, Attributor &A)
       : AANonConvergentImpl(IRP, A) {}
@@ -2929,6 +2945,11 @@ struct AANonConvergentFunction final : AANonConvergentImpl {
                                            UsedAssumedInformation)) {
       return indicatePessimisticFixpoint();
     }
+
+    const Function *F = this->getIRPosition().getAssociatedFunction();
+    if (FunctionCalledWithConvergenceToken(F))
+      return indicatePessimisticFixpoint();
+
     return ChangeStatus::UNCHANGED;
   }
 

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

@@ -1864,6 +1864,23 @@ static bool InstrBreaksNonConvergent(Instruction &I,
          !SCCNodes.contains(CB->getCalledFunction());
 }
 
+static bool FunctionRequiresConvergence(const Function *F) {
+  for (auto &Use : F->uses()) {
+    CallBase *CB = dyn_cast<CallBase>(Use.getUser());
+    if (!CB)
+      return true;
+
+    // We are not called, just used as an argument.
+    if (CB->getCalledFunction() != F)
+      continue;
+
+    if (CB->getConvergenceControlToken())
+      return true;
+  }
+
+  return false;
+}
+
 /// Helper for NoUnwind inference predicate InstrBreaksAttribute.
 static bool InstrBreaksNonThrowing(Instruction &I, const SCCNodeSet &SCCNodes) {
   if (!I.mayThrow(/* IncludePhaseOneUnwind */ true))
@@ -1967,7 +1984,9 @@ static void inferConvergent(const SCCNodeSet &SCCNodes,
   AI.registerAttrInference(AttributeInferer::InferenceDescriptor{
       Attribute::Convergent,
       // Skip non-convergent functions.
-      [](const Function &F) { return !F.isConvergent(); },
+      [](const Function &F) {
+        return !F.isConvergent() || FunctionRequiresConvergence(&F);
+      },
       // Instructions that break non-convergent assumption.
       [SCCNodes](Instruction &I) {
         return InstrBreaksNonConvergent(I, SCCNodes);

llvm/test/Transforms/ADCE/convergence.ll

@@ -0,0 +1,92 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt %s -passes=adce -S | FileCheck %s
+
+; CHECK:      Function Attrs: convergent
+define i32 @foo(i32 %a) #0 {
+; CHECK-LABEL: define i32 @foo(
+; CHECK-SAME: i32 [[A:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret i32 [[A]]
+;
+entry:
+  %tk = call token @llvm.experimental.convergence.entry()
+  ret i32 %a
+}
+
+; CHECK:      Function Attrs: convergent
+define void @bar() #0 {
+; CHECK-LABEL: define void @bar(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tk = call token @llvm.experimental.convergence.anchor()
+  ret void
+}
+
+; CHECK:      Function Attrs: convergent
+define void @baz() #0 {
+; CHECK-LABEL: define void @baz(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[HEADER:.*]]
+; CHECK:       [[HEADER]]:
+; CHECK-NEXT:    br i1 true, label %[[BODY:.*]], label %[[EXIT:.*]]
+; CHECK:       [[BODY]]:
+; CHECK-NEXT:    br label %[[HEADER]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tk0 = call token @llvm.experimental.convergence.entry()
+  br label %header
+
+header:
+  %tk1 = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %tk0) ]
+  br i1 true, label %body, label %exit
+
+body:
+  br label %header
+
+exit:
+  ret void
+}
+
+define void @indirect_inner() #0 {
+; CHECK-LABEL: define void @indirect_inner(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tk0 = call token @llvm.experimental.convergence.entry()
+  ret void
+}
+
+define void @indirect() #0 {
+; CHECK-LABEL: define void @indirect(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TK0:%.*]] = call token @llvm.experimental.convergence.entry()
+; CHECK-NEXT:    [[VAR:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    store ptr @indirect_inner, ptr [[VAR]], align 8
+; CHECK-NEXT:    [[PTR:%.*]] = load ptr, ptr [[VAR]], align 8
+; CHECK-NEXT:    call void [[PTR]]() #[[ATTR0]] [ "convergencectrl"(token [[TK0]]) ]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tk0 = call token @llvm.experimental.convergence.entry()
+  %var = alloca ptr, align 8
+  store ptr @indirect_inner, ptr %var, align 8
+  %ptr = load ptr, ptr %var, align 8
+  call void %ptr() convergent [ "convergencectrl"(token %tk0) ]
+  ret void
+}
+
+declare token @llvm.experimental.convergence.entry() #1
+declare token @llvm.experimental.convergence.anchor() #1
+declare token @llvm.experimental.convergence.loop() #1
+
+attributes #0 = { convergent }
+attributes #1 = { convergent nocallback nofree nosync nounwind willreturn memory(none) }