[IPO] Prevent removal of some convergent attr

When a callee is marked as `convergent`, some targets like HLSL/SPIR-V
add a convergent token to the call.
This is valid if both functions are marked as `convergent`.

ADCE/BDCE and other DCE passes were allowed to remove convergence
intrinsics when their token were unused.
This meant a leaf function could lose all its convergence intrinsics.
This would allow further optimization to remove the `convergent`
attribute from the callee.
Issue was the caller was not updated, and we now had a convergence token
attached to a call function calling a non-convergent function.

This commit limits the removal of the `convergent` attribute when only
users are calls without convergence intrinsics.
If any other use is found (either a call to a convergent, or something
else like taking the function address), the removal is prevented.

Author: Keenuts

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

@@ -1864,6 +1864,19 @@ 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;
+
+    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 +1980,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,46 @@
+; RUN: opt %s -passes=adce -S | FileCheck %s
+
+; CHECK:      Function Attrs: convergent
+; CHECK-NEXT: define i32 @foo(i32 %a) #0 {
+define i32 @foo(i32 %a) #0 {
+entry:
+; CHECK-NOT: %0 = call token @llvm.experimental.convergence.entry()
+  %0 = call token @llvm.experimental.convergence.entry()
+  ret i32 %a
+}
+
+; CHECK:      Function Attrs: convergent
+; CHECK-NEXT: define void @bar() #0 {
+define void @bar() #0 {
+entry:
+; CHECK-NOT: %0 = call token @llvm.experimental.convergence.entry()
+  %0 = call token @llvm.experimental.convergence.anchor()
+  ret void
+}
+
+; CHECK:      Function Attrs: convergent
+; CHECK-NEXT: define void @baz() #0 {
+define void @baz() #0 {
+entry:
+; CHECK-NOT: %0 = call token @llvm.experimental.convergence.entry()
+  %0 = call token @llvm.experimental.convergence.entry()
+  br label %header
+
+header:
+; CHECK-NOT: %1 = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %0) ]
+  %1 = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %0) ]
+  br i1 true, label %body, label %exit
+
+body:
+  br label %header
+
+exit:
+  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) }

llvm/test/Transforms/BDCE/convergence.ll

@@ -0,0 +1,47 @@
+; RUN: opt %s -passes=bdce -S | FileCheck %s
+
+; CHECK:      Function Attrs: convergent
+; CHECK-NEXT: define i32 @foo(i32 %a) #0 {
+define i32 @foo(i32 %a) #0 {
+entry:
+; CHECK-NOT: %0 = call token @llvm.experimental.convergence.entry()
+  %0 = call token @llvm.experimental.convergence.entry()
+  ret i32 %a
+}
+
+; CHECK:      Function Attrs: convergent
+; CHECK-NEXT: define void @bar() #0 {
+define void @bar() #0 {
+entry:
+; CHECK-NOT: %0 = call token @llvm.experimental.convergence.entry()
+  %0 = call token @llvm.experimental.convergence.anchor()
+  ret void
+}
+
+; CHECK:      Function Attrs: convergent
+; CHECK-NEXT: define void @baz() #0 {
+define void @baz() #0 {
+entry:
+; CHECK-NOT: %0 = call token @llvm.experimental.convergence.entry()
+  %0 = call token @llvm.experimental.convergence.entry()
+  br label %header
+
+header:
+; CHECK-NOT: %1 = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %0) ]
+  %1 = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %0) ]
+  br i1 true, label %body, label %exit
+
+body:
+  br label %header
+
+exit:
+  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) }
+

llvm/test/Transforms/FunctionAttrs/convergent.ll

@@ -129,3 +129,35 @@ define i32 @noopt_friend() convergent {
   %a = call i32 @noopt()
   ret i32 0
 }
+
+; A function which should normally be stripped of its convergent attribute,
+; but because it's used in a controlled convergence call, the attribute
+; remains.
+; This could be improved by propagating the non-convergence outside of the
+; function, but for the time being, we stop when we encounter this scenario.
+define i32 @leaf_noconvergent_used() convergent {
+; CHECK: Function Attrs: convergent mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define {{[^@]+}}@leaf_noconvergent_used
+; CHECK-SAME: () #[[ATTR7:[0-9]+]] {
+; CHECK-NEXT:    ret i32 0
+;
+  ret i32 0
+}
+
+define i32 @nonleaf_convergent() convergent {
+; CHECK: Function Attrs: convergent mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define {{[^@]+}}@nonleaf_convergent
+; CHECK-SAME: () #[[ATTR7]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = call token @llvm.experimental.convergence.entry()
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @leaf_noconvergent_used() [ "convergencectrl"(token [[TMP1]]) ]
+; CHECK-NEXT:    ret i32 0
+;
+  %1 = call token @llvm.experimental.convergence.entry()
+  %2 = call i32 @leaf_noconvergent_used() [ "convergencectrl"(token %1) ]
+  ret i32 0
+}
+
+
+declare token @llvm.experimental.convergence.entry() #1
+
+attributes #1 = { convergent nocallback nofree nosync nounwind willreturn memory(none) }