[llvm] [IndVarSimplify] Fix `IndVarSimplify` to skip on unfolding predicates when the loop contains control convergence operations. (PR #165643)

[Nikita Popov via llvm-commits]

================
@@ -1859,6 +1859,37 @@ bool IndVarSimplify::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) {
         }
       }
 
+  // If the loop body uses a convergence token defined within the loop, skip
+  // predication. This is to avoid changing the convergence behavior of the
+  // loop.
+  SmallVector<BasicBlock *, 16> blocks = ExitingBlocks;
+  SmallVector<Value *, 16> tokens = {};
+  size_t index = 0; // Assume Exiting Blocks are sorted.
+  while (index < blocks.size()) {
+    BasicBlock *BB = blocks[index];
+    index++;
+    const auto exitingBlockName = BB->getName();
+    for (Instruction &I : *BB) {
+      // Check if the instruction uses any convergence tokens.
+      if (auto *CB = dyn_cast<CallBase>(&I);
+          CB && !isa<ConvergenceControlInst>(&I)) {
+        auto token = CB->getConvergenceControlToken();
+        if (token && llvm::is_contained(tokens, token)) {
+          return false;
+        }
+      }
+      if (isa<ConvergenceControlInst>(&I)) {
+        tokens.push_back(cast<Value>(&I));
+      }
+    }
+
+    for (BasicBlock *Succ : successors(BB)) {
----------------
nikic wrote:

I don't think analyzeBasicBlock is appropriate here -- CodeMetrics is a utility for transforms that perform block cloning.

Looking a bit closer, I'm also not sure the problem here is really related to convergence control per se. This transform is generally not valid for cases where a value defined in the loop is used outside the loop. Normally, this is enforced by checking for the absence of LCSSA phis. However, token-returning instructions are a special case, because LCSSA cannot be formed for them.

So I think what you should be doing is looking for any token-returning instructions in the loop (not necessarily convergence control in particular) and bailing out if they have uses outside the loop.

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