From 8cb224b44394dee6ff8b764cd04ad736d647e44e Mon Sep 17 00:00:00 2001
From: Pavel Samolysov <samolisov@gmail.com>
Date: Wed, 22 May 2024 22:10:55 +0300
Subject: [PATCH] [PGO] Preserve analysis results when nothing was instrumented

The 'PGOInstrumentationGen' pass should preserve all analysis results
when nothing was actually instrumented. Currently, only modules that
contain at least a single function definition are instrumented. When a
module contains only function declarations and, optionally, global
variable definitions (a module for the regular-LTO phase for thin-LTO
when LTOUnit splitting is enabled, for example), such module is not
instrumented (yet?) and there is no reason to introduce the
'__llvm_profile_raw_version' variable into the module.
---
 .../Instrumentation/PGOInstrumentation.cpp    |  4 +-
 .../PGOInstrumentationTest.cpp                | 68 ++++++++++++++-----
 2 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index e6e474ed376069..dbe908bb5e72f3 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1916,6 +1916,7 @@ static bool InstrumentAllFunctions(
   std::unordered_multimap<Comdat *, GlobalValue *> ComdatMembers;
   collectComdatMembers(M, ComdatMembers);
 
+  bool AnythingInstrumented = false;
   for (auto &F : M) {
     if (skipPGOGen(F))
       continue;
@@ -1925,8 +1926,9 @@ static bool InstrumentAllFunctions(
     FunctionInstrumenter FI(M, F, TLI, ComdatMembers, BPI, BFI,
                             InstrumentationType);
     FI.instrument();
+    AnythingInstrumented = true;
   }
-  return true;
+  return AnythingInstrumented;
 }
 
 PreservedAnalyses
diff --git a/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp b/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp
index 9ccb13934cbd38..a4c076a8752fc3 100644
--- a/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp
+++ b/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp
@@ -103,9 +103,13 @@ class MockModuleAnalysisHandle
                ModuleAnalysisManager::Invalidator &));
 };
 
-struct PGOInstrumentationGenTest
-    : public Test,
-      WithParamInterface<std::tuple<StringRef, StringRef>> {
+template <typename ParamType> struct PGOTestName {
+  std::string operator()(const TestParamInfo<ParamType> &Info) const {
+    return std::get<1>(Info.param).str();
+  }
+};
+
+struct PGOInstrumentationGenTest : public Test {
   ModulePassManager MPM;
   PassBuilder PB;
   MockModuleAnalysisHandle MMAHandle;
@@ -141,12 +145,47 @@ struct PGOInstrumentationGenTest
   }
 };
 
+struct PGOInstrumentationGenInstrumentTest
+    : PGOInstrumentationGenTest,
+      WithParamInterface<std::tuple<StringRef, StringRef>> {};
+
 static constexpr StringRef CodeWithFuncDefs = R"(
   define i32 @f(i32 %n) {
   entry:
     ret i32 0
   })";
 
+INSTANTIATE_TEST_SUITE_P(
+    PGOInstrumetationGenTestSuite, PGOInstrumentationGenInstrumentTest,
+    Values(std::make_tuple(CodeWithFuncDefs, "instrument_function_defs")),
+    PGOTestName<PGOInstrumentationGenInstrumentTest::ParamType>());
+
+TEST_P(PGOInstrumentationGenInstrumentTest, Instrumented) {
+  const StringRef Code = std::get<0>(GetParam());
+  parseAssembly(Code);
+
+  ASSERT_THAT(M, NotNull());
+
+  Sequence PassSequence;
+  EXPECT_CALL(MMAHandle, run(Ref(*M), _))
+      .InSequence(PassSequence)
+      .WillOnce(DoDefault());
+  EXPECT_CALL(MMAHandle, invalidate(Ref(*M), _, _))
+      .InSequence(PassSequence)
+      .WillOnce(DoDefault());
+
+  MPM.run(*M, MAM);
+
+  const auto *IRInstrVar =
+      M->getNamedGlobal(INSTR_PROF_QUOTE(INSTR_PROF_RAW_VERSION_VAR));
+  ASSERT_THAT(IRInstrVar, NotNull());
+  EXPECT_FALSE(IRInstrVar->isDeclaration());
+}
+
+struct PGOInstrumentationGenIgnoreTest
+    : PGOInstrumentationGenTest,
+      WithParamInterface<std::tuple<StringRef, StringRef>> {};
+
 static constexpr StringRef CodeWithFuncDecls = R"(
   declare i32 @f(i32);
 )";
@@ -157,33 +196,26 @@ static constexpr StringRef CodeWithGlobals = R"(
 )";
 
 INSTANTIATE_TEST_SUITE_P(
-    PGOInstrumetationGenTestSuite, PGOInstrumentationGenTest,
-    Values(std::make_tuple(CodeWithFuncDefs, "instrument_function_defs"),
-           std::make_tuple(CodeWithFuncDecls, "instrument_function_decls"),
+    PGOInstrumetationGenIgnoreTestSuite, PGOInstrumentationGenIgnoreTest,
+    Values(std::make_tuple(CodeWithFuncDecls, "instrument_function_decls"),
            std::make_tuple(CodeWithGlobals, "instrument_globals")),
-    [](const TestParamInfo<PGOInstrumentationGenTest::ParamType> &Info) {
-      return std::get<1>(Info.param).str();
-    });
+    PGOTestName<PGOInstrumentationGenIgnoreTest::ParamType>());
 
-TEST_P(PGOInstrumentationGenTest, Instrumented) {
+TEST_P(PGOInstrumentationGenIgnoreTest, NotInstrumented) {
   const StringRef Code = std::get<0>(GetParam());
+
   parseAssembly(Code);
 
   ASSERT_THAT(M, NotNull());
 
-  Sequence PassSequence;
-  EXPECT_CALL(MMAHandle, run(Ref(*M), _))
-      .InSequence(PassSequence)
-      .WillOnce(DoDefault());
-  EXPECT_CALL(MMAHandle, invalidate(Ref(*M), _, _))
-      .InSequence(PassSequence)
-      .WillOnce(DoDefault());
+  EXPECT_CALL(MMAHandle, run(Ref(*M), _)).WillOnce(DoDefault());
+  EXPECT_CALL(MMAHandle, invalidate(Ref(*M), _, _)).Times(0);
 
   MPM.run(*M, MAM);
 
   const auto *IRInstrVar =
       M->getNamedGlobal(INSTR_PROF_QUOTE(INSTR_PROF_RAW_VERSION_VAR));
-  EXPECT_THAT(IRInstrVar, NotNull());
+  ASSERT_THAT(IRInstrVar, NotNull());
   EXPECT_FALSE(IRInstrVar->isDeclaration());
 }