[LoopSimplify] don't separate nested loops with convergent calls

[sameerds]

Summary

When a loop has multiple backedges, loop simplification attempts to
separate them out into nested loops. This results in incorrect control
flow in the presence of some functions like a GPU barrier. This change
skips the transformation when such "convergent" function calls are
present in the loop body.

[sameerds]

Test update: Moved the convergent call to where it is actually relevant.

[nhaehnle]

Looks good to me.

[arsenm]

llvm/lib/Transforms/Utils/LoopSimplify.cpp

244

I thought we had a hasConvergentOp on loops that would avoid scanning through all the blocks again

[sameerds]

llvm/lib/Transforms/Utils/LoopSimplify.cpp

244

That is available when LoopAccessAnalysis is used. But LoopSimplify does not already use it, and it seems a bit expensive for just this one use.

llvm/lib/Transforms/Utils/LoopSimplify.cpp

 static Loop *separateNestedLoop(Loop *L, BasicBlock *Preheader,
                                 DominatorTree *DT, LoopInfo *LI,
                                 ScalarEvolution *SE, bool PreserveLCSSA,
                                 AssumptionCache *AC, MemorySSAUpdater *MSSAU) {
   // Don't try to separate loops without a preheader.
   if (!Preheader)
     return nullptr;
 
+  // Treat the presence of convergent functions conservatively. The
+  // transformation is invalid if calls to certain convergent
+  // functions (like an AMDGPU barrier) get included in the resulting
+  // inner loop. But blocks meant for the inner loop will be
+  // identified later at a point where it's too late to abort the
+  // transformation. Also, the convergent attribute is not really
+  // sufficient to express the semantics of functions that are
+  // affected by this transformation. So we choose to back off if such
+  // a function call is present until a better alternative becomes
+  // available. This is similar to the conservative treatment of
+  // convergent function calls in GVNHoist and JumpThreading.
+  for (auto BB : L->blocks()) {
+    for (auto &II : *BB) {
+      if (auto CI = dyn_cast<CallBase>(&II)) {
+        if (CI->isConvergent()) {
+          return nullptr;
+        }
+      }
+    }
+  }
+
   // The header is not a landing pad; preheader insertion should ensure this.
   BasicBlock *Header = L->getHeader();
   assert(!Header->isEHPad() && "Can't insert backedge to EH pad");
 
   PHINode *PN = findPHIToPartitionLoops(L, DT, AC);
   if (!PN) return nullptr;  // No known way to partition.
 
   // Pull out all predecessors that have varying values in the loop.  This

llvm/test/Transforms/LoopSimplify/convergent.ll

+; RUN: opt < %s -S -loop-simplify | FileCheck %s
+
+; Don't separate out nested loops if a convergent call is present
+
+; CHECK-NOT: BB1.outer
+; CHECK: BB1.backedge
+
+define i32 @test(i1 %loop_cond, i1 %exit_cond, i32 %init) {
+entry:
+  br label %BB1
+
+BB1:
+  %indvar = phi i32 [%indvar, %BB1], [%inc, %BB2], [%init, %entry]
+  call void @f() convergent
+  br i1 %loop_cond, label %BB1, label %BB2
+
+BB2:
+  %inc = add i32 %indvar, 1
+  br i1 %exit_cond, label %exit, label %BB1
+
+exit:
+  ret i32 %inc
+}
+
+declare void @f() convergent