From b8c39eb2756f140c506a5c957fb0f3eaed657316 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin@google.com>
Date: Wed, 8 Jun 2022 14:30:43 -0700
Subject: [PATCH] Fix FunctionPropertiesAnalysis updating callsite in 1-BB loop

If the callsite is in a single BB loop, we need to exclude the BB from
the successor set (in which it'd be a member), because that set forms a
boundary at which we stop traversing the CFG, when re-ingesting BBs
after inlining; but after inlining, the callsite BB's new successors
should be visited.

Reviewed By: kazu

Differential Revision: https://reviews.llvm.org/D127178
---
 llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp   |  8 +++
 .../Analysis/FunctionPropertiesAnalysisTest.cpp    | 67 ++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp b/llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp
index 66498f4..1bac19a 100644
--- a/llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp
+++ b/llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp
@@ -134,6 +134,14 @@ FunctionPropertiesUpdater::FunctionPropertiesUpdater(
   // We track successors separately, too, because they form a boundary, together
   // with the CB BB ('Entry') between which the inlined callee will be pasted.
   Successors.insert(succ_begin(&CallSiteBB), succ_end(&CallSiteBB));
+
+  // Exclude the CallSiteBB, if it happens to be its own successor (1-BB loop).
+  // We are only interested in BBs the graph moves past the callsite BB to
+  // define the frontier past which we don't want to re-process BBs. Including
+  // the callsite BB in this case would prematurely stop the traversal in
+  // finish().
+  Successors.erase(&CallSiteBB);
+
   for (const auto *BB : Successors)
     LikelyToChangeBBs.insert(BB);
 
diff --git a/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp b/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp
index 76a63c0..872e37c 100644
--- a/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp
+++ b/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp
@@ -543,4 +543,71 @@ lpad:
   EXPECT_EQ(FPI, FunctionPropertiesInfo::getFunctionPropertiesInfo(*F1, LINew));
 }
 
+TEST_F(FunctionPropertiesAnalysisTest, InlineSameLoopBB) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = makeLLVMModule(C,
+                                             R"IR(
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+declare i32 @a()
+declare i32 @b()
+
+define i32 @f1(i32 %a) {
+entry:
+  br label %loop
+loop:
+  %i = call i32 @f2(i32 %a)
+  %c = icmp slt i32 %i, %a
+  br i1 %c, label %loop, label %end
+end:
+  %r = phi i32 [%i, %loop], [%a, %entry]
+  ret i32 %r
+}
+
+define i32 @f2(i32 %a) {
+  %cnd = icmp slt i32 %a, 0
+  br i1 %cnd, label %then, label %else
+then:
+  %r1 = call i32 @a()
+  br label %end
+else:
+  %r2 = call i32 @b()
+  br label %end
+end:
+  %r = phi i32 [%r1, %then], [%r2, %else]
+  ret i32 %r
+}
+)IR");
+
+  Function *F1 = M->getFunction("f1");
+  CallBase *CB = findCall(*F1);
+  EXPECT_NE(CB, nullptr);
+
+  FunctionPropertiesInfo ExpectedInitial;
+  ExpectedInitial.BasicBlockCount = 3;
+  ExpectedInitial.TotalInstructionCount = 6;
+  ExpectedInitial.BlocksReachedFromConditionalInstruction = 2;
+  ExpectedInitial.Uses = 1;
+  ExpectedInitial.DirectCallsToDefinedFunctions = 1;
+  ExpectedInitial.MaxLoopDepth = 1;
+  ExpectedInitial.TopLevelLoopCount = 1;
+
+  FunctionPropertiesInfo ExpectedFinal = ExpectedInitial;
+  ExpectedFinal.BasicBlockCount = 6;
+  ExpectedFinal.DirectCallsToDefinedFunctions = 0;
+  ExpectedFinal.BlocksReachedFromConditionalInstruction = 4;
+  ExpectedFinal.TotalInstructionCount = 12;
+
+  auto FPI = buildFPI(*F1);
+  EXPECT_EQ(FPI, ExpectedInitial);
+
+  FunctionPropertiesUpdater FPU(FPI, *CB);
+  InlineFunctionInfo IFI;
+  auto IR = llvm::InlineFunction(*CB, IFI);
+  EXPECT_TRUE(IR.isSuccess());
+  FPU.finish(*LI);
+  EXPECT_EQ(FPI, ExpectedFinal);
+}
+
 } // end anonymous namespace
-- 
2.7.4