[OPT] Search whole BB for convergence token.

[s-perron]

The spec for llvm.experimental.convergence.entry says that is must be in
the entry block for a function, and must preceed any other convergent
operation. It does not have to be the first instruction in the entry
block.

Inlining assumes that the call to llvm.experimental.convergence.entry
will be the first instruction after any phi instructions. This commit
modifies inlining to search the entire block for the call.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Steven Perron (s-perron)

Changes

The spec for llvm.experimental.convergence.entry says that is must be in
the entry block for a function, and must preceed any other convergent
operation. It does not have to be the first instruction in the entry
block.

Inlining assumes that the call to llvm.experimental.convergence.entry
will be the first instruction after any phi instructions. This commit
modifies inlining to search the entire block for the call.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/InlineFunction.cpp (+21-16)
• (modified) llvm/test/Transforms/Inline/convergence-inline.ll (+24)

diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 55ad2b6d620003..3386aadd8092f0 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -180,6 +180,19 @@ namespace {
     }
   };
 
+  IntrinsicInst *getConevrgenceEntryIfAny(BasicBlock &BB) {
+    auto *I = BB.getFirstNonPHI();
+    while (I) {
+      if (auto *IntrinsicCall = dyn_cast<IntrinsicInst>(I)) {
+        if (IntrinsicCall->getIntrinsicID() ==
+            Intrinsic::experimental_convergence_entry) {
+          return IntrinsicCall;
+        }
+      }
+      I = I->getNextNode();
+    }
+    return nullptr;
+  }
 } // end anonymous namespace
 
 /// Get or create a target for the branch from ResumeInsts.
@@ -2443,15 +2456,10 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
   // fully implements convergence control tokens, there is no mixing of
   // controlled and uncontrolled convergent operations in the whole program.
   if (CB.isConvergent()) {
-    auto *I = CalledFunc->getEntryBlock().getFirstNonPHI();
-    if (auto *IntrinsicCall = dyn_cast<IntrinsicInst>(I)) {
-      if (IntrinsicCall->getIntrinsicID() ==
-          Intrinsic::experimental_convergence_entry) {
-        if (!ConvergenceControlToken) {
-          return InlineResult::failure(
-              "convergent call needs convergencectrl operand");
-        }
-      }
+    auto *I = getConevrgenceEntryIfAny(CalledFunc->getEntryBlock());
+    if (I && !ConvergenceControlToken) {
+      return InlineResult::failure(
+          "convergent call needs convergencectrl operand");
     }
   }
 
@@ -2742,13 +2750,10 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
   }
 
   if (ConvergenceControlToken) {
-    auto *I = FirstNewBlock->getFirstNonPHI();
-    if (auto *IntrinsicCall = dyn_cast<IntrinsicInst>(I)) {
-      if (IntrinsicCall->getIntrinsicID() ==
-          Intrinsic::experimental_convergence_entry) {
-        IntrinsicCall->replaceAllUsesWith(ConvergenceControlToken);
-        IntrinsicCall->eraseFromParent();
-      }
+    auto *IntrinsicCall = getConevrgenceEntryIfAny(*FirstNewBlock);
+    if (IntrinsicCall) {
+      IntrinsicCall->replaceAllUsesWith(ConvergenceControlToken);
+      IntrinsicCall->eraseFromParent();
     }
   }
 
diff --git a/llvm/test/Transforms/Inline/convergence-inline.ll b/llvm/test/Transforms/Inline/convergence-inline.ll
index 8c67e6a59b7db1..4996a2376be638 100644
--- a/llvm/test/Transforms/Inline/convergence-inline.ll
+++ b/llvm/test/Transforms/Inline/convergence-inline.ll
@@ -185,6 +185,30 @@ define void @test_two_calls() convergent {
   ret void
 }
 
+define i32 @token_not_first(i32 %x) convergent alwaysinline {
+; CHECK-LABEL: @token_not_first(
+; CHECK-NEXT:    {{%.*}} = alloca ptr, align 8
+; CHECK-NEXT:    [[TOKEN:%.*]] = call token @llvm.experimental.convergence.entry()
+; CHECK-NEXT:    [[Y:%.*]] = call i32 @g(i32 [[X:%.*]]) [ "convergencectrl"(token [[TOKEN]]) ]
+; CHECK-NEXT:    ret i32 [[Y]]
+;
+  %p = alloca ptr, align 8
+  %token = call token @llvm.experimental.convergence.entry()
+  %y = call i32 @g(i32 %x) [ "convergencectrl"(token %token) ]
+  ret i32 %y
+}
+
+define void @test_token_not_first() convergent {
+; CHECK-LABEL: @test_token_not_first(
+; CHECK-NEXT:    [[TOKEN:%.*]] = call token @llvm.experimental.convergence.entry()
+; CHECK-NEXT:    {{%.*}} = call i32 @g(i32 23) [ "convergencectrl"(token [[TOKEN]]) ]
+; CHECK-NEXT:    ret void
+;
+  %token = call token @llvm.experimental.convergence.entry()
+  %x = call i32 @token_not_first(i32 23) [ "convergencectrl"(token %token) ]
+  ret void
+}
+
 declare void @f(i32) convergent
 declare i32 @g(i32) convergent
 

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[ssahasra]

LGTM, but just one nit needs fixing.