[NFC] Introduce API to detect tokens penetrating LCSSA form

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.

Differential Revision: https://reviews.llvm.org/D129983
Reviewed By: efriedma

llvm/include/llvm/Analysis/LoopInfo.h

@@ -814,12 +814,15 @@ public:
   /// 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 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.
-  bool isRecursivelyLCSSAForm(const DominatorTree &DT,
-                              const LoopInfo &LI) 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.

llvm/lib/Analysis/LoopInfo.cpp

@@ -425,12 +425,12 @@ bool Loop::isCanonical(ScalarEvolution &SE) const {
 
 // 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()) {
@@ -455,20 +455,20 @@ static bool isBlockInLCSSAForm(const Loop &L, const BasicBlock &BB,
   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);
   });
 }
 

llvm/unittests/Analysis/LoopInfoTest.cpp

@@ -1584,3 +1584,66 @@ TEST(LoopInfoTest, LoopInductionVariable) {
         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));
+  });
+}