[NFC] Introduce API to detect tokens penetrating LCSSA form

[mkazantsev]

Summary

Following discussion in PR56243, we need to somehow detect the situation
when token values penetrate LCSSA form for transforms that require that
it is maintained by all values (for example, to sustain use-def dominance
invarians). This patch introduces a parameter to LCSSA checkers to control
their ignorance about tokens.

[efriedma]

LGTM

llvm/include/llvm/Analysis/LoopInfo.h

bool isRotatedForm() const {
     BasicBlock *Latch = getLoopLatch();
     return Latch && isLoopExiting(Latch);
   }
 
   /// Return true if the loop induction variable starts at zero and increments
   /// by one each time through the loop.
   bool isCanonical(ScalarEvolution &SE) const;
 
-  /// Return true if the Loop is in LCSSA form.
-  bool isLCSSAForm(const DominatorTree &DT) const;
-
-  /// Return true if this Loop and all inner subloops are in LCSSA form.
-  bool isRecursivelyLCSSAForm(const DominatorTree &DT,
-                              const LoopInfo &LI) const;
+  /// Return true if the Loop is in LCSSA form. If \p IgnoreTokens is set to
+  /// true, token values defined inside loop are allowed to violate LCSSA form.
+  bool isLCSSAForm(const DominatorTree &DT, bool IgnoreTokens = true) const;
+
+  /// Return true if this Loop and all inner subloops are in LCSSA form. If \p
+  /// IgnoreTokens is set to true, token values defined inside loop are allowed
+  /// to violate LCSSA form.
+  bool isRecursivelyLCSSAForm(const DominatorTree &DT, const LoopInfo &LI,
+                              bool IgnoreTokens = true) const;
 
   /// Return true if the Loop is in the form that the LoopSimplify form
   /// transforms loops to, which is sometimes called normal form.
   bool isLoopSimplifyForm() const;
 
   /// Return true if the loop body is safe to clone in practice.
   bool isSafeToClone() const;
 

llvm/lib/Analysis/LoopInfo.cpp

bool Loop::isCanonical(ScalarEvolution &SE) const {
   if (!Step || !Step->isOne())
     return false;
 
   return true;
 }
 
 // Check that 'BB' doesn't have any uses outside of the 'L'
 static bool isBlockInLCSSAForm(const Loop &L, const BasicBlock &BB,
-                               const DominatorTree &DT) {
+                               const DominatorTree &DT, bool IgnoreTokens) {
   for (const Instruction &I : BB) {
     // Tokens can't be used in PHI nodes and live-out tokens prevent loop
     // optimizations, so for the purposes of considered LCSSA form, we
     // can ignore them.
-    if (I.getType()->isTokenTy())
+    if (IgnoreTokens && I.getType()->isTokenTy())
       continue;
 
     for (const Use &U : I.uses()) {
       const Instruction *UI = cast<Instruction>(U.getUser());
       const BasicBlock *UserBB = UI->getParent();
 
       // For practical purposes, we consider that the use in a PHI
       // occurs in the respective predecessor block. For more info,
       // see the `phi` doc in LangRef and the LCSSA doc.
       if (const PHINode *P = dyn_cast<PHINode>(UI))
         UserBB = P->getIncomingBlock(U);
 
       // Check the current block, as a fast-path, before checking whether
       // the use is anywhere in the loop.  Most values are used in the same
       // block they are defined in.  Also, blocks not reachable from the
       // entry are special; uses in them don't need to go through PHIs.
       if (UserBB != &BB && !L.contains(UserBB) &&
           DT.isReachableFromEntry(UserBB))
         return false;
     }
   }
   return true;
 }
 
-bool Loop::isLCSSAForm(const DominatorTree &DT) const {
+bool Loop::isLCSSAForm(const DominatorTree &DT, bool IgnoreTokens) const {
   // For each block we check that it doesn't have any uses outside of this loop.
   return all_of(this->blocks(), [&](const BasicBlock *BB) {
-    return isBlockInLCSSAForm(*this, *BB, DT);
+    return isBlockInLCSSAForm(*this, *BB, DT, IgnoreTokens);
   });
 }
 
-bool Loop::isRecursivelyLCSSAForm(const DominatorTree &DT,
-                                  const LoopInfo &LI) const {
+bool Loop::isRecursivelyLCSSAForm(const DominatorTree &DT, const LoopInfo &LI,
+                                  bool IgnoreTokens) const {
   // For each block we check that it doesn't have any uses outside of its
   // innermost loop. This process will transitively guarantee that the current
   // loop and all of the nested loops are in LCSSA form.
   return all_of(this->blocks(), [&](const BasicBlock *BB) {
-    return isBlockInLCSSAForm(*LI.getLoopFor(BB), *BB, DT);
+    return isBlockInLCSSAForm(*LI.getLoopFor(BB), *BB, DT, IgnoreTokens);
   });
 }
 
 bool Loop::isLoopSimplifyForm() const {
   // Normal-form loops have a preheader, a single backedge, and all of their
   // exits have all their predecessors inside the loop.
   return getLoopPreheader() && getLoopLatch() && hasDedicatedExits();
 }

llvm/unittests/Analysis/LoopInfoTest.cpp

runWithLoopInfoPlus(
       *M, "foo", [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
         Function::iterator FI = F.begin();
         BasicBlock *Header = &*(++FI);
         Loop *L = LI.getLoopFor(Header);
         EXPECT_NE(L, nullptr);
         EXPECT_EQ(L->getInductionVariable(SE)->getName(), "count.07");
       });
 }
+
+// Test that we correctly identify tokens breaching LCSSA form.
+TEST(LoopInfoTest, TokenLCSSA) {
+  const char *ModuleStr =
+      "define void @test() gc \"statepoint-example\" {\n"
+      "entry:\n"
+      "  br label %outer_loop\n"
+      "outer_loop:\n"
+      "  br label %inner_loop\n"
+      "inner_loop:\n"
+      "  %token = call token (i64, i32, i8 addrspace(1)* (i64, i32, i32, "
+      "i32)*, i32, i32, ...) "
+      "@llvm.experimental.gc.statepoint.p0f_p1i8i64i32i32i32f(i64 2882400000, "
+      "i32 0, i8 addrspace(1)* (i64, i32, i32, i32)* nonnull elementtype(i8 "
+      "addrspace(1)* (i64, i32, i32, i32)) @foo, i32 4, i32 0, i64 undef, i32 "
+      "5, i32 5, i32 undef, i32 0, i32 0) [ \"deopt\"(), \"gc-live\"(i8 "
+      "addrspace(1)* undef) ]\n"
+      "  br i1 undef, label %inner_loop, label %outer_backedge\n"
+      "outer_backedge:\n"
+      "  br i1 undef, label %outer_loop, label %exit\n"
+      "exit:\n"
+      "  %tmp35 = call coldcc i8 addrspace(1)* "
+      "@llvm.experimental.gc.relocate.p1i8(token %token, i32 0, i32 0) ; "
+      "(undef, undef)\n"
+      "  ret void\n"
+      "}\n"
+      "declare i8 addrspace(1)* @foo(i64, i32, i32, i32)\n"
+      "declare i8 addrspace(1)* @llvm.experimental.gc.relocate.p1i8(token, i32 "
+      "immarg, i32 immarg) #0\n"
+      "declare token "
+      "@llvm.experimental.gc.statepoint.p0f_p1i8i64i32i32i32f(i64 immarg, i32 "
+      "immarg, i8 addrspace(1)* (i64, i32, i32, i32)*, i32 immarg, i32 immarg, "
+      "...)\n"
+      "attributes #0 = { nounwind readnone }\n";
+
+  // Parse the module.
+  LLVMContext Context;
+  std::unique_ptr<Module> M = makeLLVMModule(Context, ModuleStr);
+
+  runWithLoopInfoPlus(*M, "test",
+                      [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+    Function::iterator FI = F.begin();
+    BasicBlock *OuterHeader = &*(++FI);
+    Loop *OuterLoop = LI.getLoopFor(OuterHeader);
+    BasicBlock *InnerHeader = &*(++FI);
+    Loop *InnerLoop = LI.getLoopFor(InnerHeader);
+    EXPECT_NE(OuterLoop, nullptr);
+    EXPECT_NE(InnerLoop, nullptr);
+    DominatorTree DT(F);
+    EXPECT_TRUE(OuterLoop->isLCSSAForm(DT, /*IgnoreTokens*/ true));
+    EXPECT_FALSE(OuterLoop->isLCSSAForm(DT, /*IgnoreTokens*/ false));
+    EXPECT_TRUE(InnerLoop->isLCSSAForm(DT, /*IgnoreTokens*/ true));
+    EXPECT_FALSE(InnerLoop->isLCSSAForm(DT, /*IgnoreTokens*/ false));
+    EXPECT_TRUE(
+        OuterLoop->isRecursivelyLCSSAForm(DT, LI, /*IgnoreTokens*/ true));
+    EXPECT_FALSE(
+        OuterLoop->isRecursivelyLCSSAForm(DT, LI, /*IgnoreTokens*/ false));
+    EXPECT_TRUE(
+        InnerLoop->isRecursivelyLCSSAForm(DT, LI, /*IgnoreTokens*/ true));
+    EXPECT_FALSE(
+        InnerLoop->isRecursivelyLCSSAForm(DT, LI, /*IgnoreTokens*/ false));
+  });
+}