Handle convergence intrinsics

Author: justinfargnoli

llvm/lib/Transforms/Utils/IRNormalizer.cpp

@@ -468,8 +468,13 @@ void IRNormalizer::reorderInstructions(Function &F) const {
     // Reorder based on the topological sort.
     while (!TopologicalSort.empty()) {
       auto *Instruction = TopologicalSort.top();
-      auto *FirstNonPHIOrDbgOrAlloca = &*BB.getFirstNonPHIOrDbgOrAlloca();
-      Instruction->moveBefore(FirstNonPHIOrDbgOrAlloca);
+      auto FirstNonPHIOrDbgOrAlloca = BB.getFirstNonPHIOrDbgOrAlloca();
+      if (auto *Call = dyn_cast<CallInst>(&*FirstNonPHIOrDbgOrAlloca)) {
+        if (Call->getIntrinsicID() == Intrinsic::experimental_convergence_entry || 
+            Call->getIntrinsicID() == Intrinsic::experimental_convergence_loop)
+          FirstNonPHIOrDbgOrAlloca++;
+      }
+      Instruction->moveBefore(&*FirstNonPHIOrDbgOrAlloca);
       TopologicalSort.pop();
     }
   }
@@ -508,6 +513,10 @@ void IRNormalizer::reorderDefinition(
       return;
     if (Call->getIntrinsicID() == Intrinsic::experimental_deoptimize)
       return;
+    if (Call->getIntrinsicID() == Intrinsic::experimental_convergence_entry)
+      return;
+    if (Call->getIntrinsicID() == Intrinsic::experimental_convergence_loop)
+      return;
   }
   if (auto *BitCast = dyn_cast<BitCastInst>(Definition)) {
     if (auto *Call = dyn_cast<CallInst>(BitCast->getOperand(0))) {

llvm/test/Transforms/IRNormalizer/regression-convergence-tokens.ll

@@ -0,0 +1,27 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=normalize < %s | FileCheck %s
+
+; Function Attrs: convergent nounwind readnone
+define i32 @nested(i32 %src) #0 {
+; CHECK-LABEL: define i32 @nested(
+; CHECK-SAME: i32 [[A0:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[BB15160:.*:]]
+; CHECK-NEXT:    [[T1:%.*]] = call token @llvm.experimental.convergence.entry()
+; CHECK-NEXT:    %"vl15001llvm.experimental.convergence.anchor()" = call token @llvm.experimental.convergence.anchor()
+; CHECK-NEXT:    %"op68297llvm.amdgcn.readfirstlane.i32([[A0]], vl15001llvm.experimental.convergence.anchor())" = call i32 @llvm.amdgcn.readfirstlane.i32(i32 [[A0]]) [ "convergencectrl"(token %"vl15001llvm.experimental.convergence.anchor()") ]
+; CHECK-NEXT:    ret i32 undef
+;
+  %t1 = call token @llvm.experimental.convergence.entry()
+  %t2 = call token @llvm.experimental.convergence.anchor()
+  %r2 = call i32 @llvm.amdgcn.readfirstlane(i32 %src) [ "convergencectrl"(token %t2) ]
+  ret i32 undef
+}
+
+; Function Attrs: convergent nounwind readnone
+declare i32 @llvm.amdgcn.readfirstlane(i32) #0
+
+declare token @llvm.experimental.convergence.entry()
+
+declare token @llvm.experimental.convergence.anchor()
+
+attributes #0 = { convergent nounwind readnone }