[PGO][PGSO] Distinguish queries from unit tests and explicitly enable for the existing IR passes only. NFC.

[hjyamauchi]

Summary

This is one more prep step necessary before the code gen pass instrumentation
code could go in.

[davidxl]

lgtm

[hjyamauchi]

Rebase.

llvm/include/llvm/Transforms/Utils/SizeOpts.h

 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
 
 extern cl::opt<bool> EnablePGSO;
 extern cl::opt<bool> PGSOLargeWorkingSetSizeOnly;
+extern cl::opt<bool> PGSOIRPassOrTestOnly;
 extern cl::opt<bool> PGSOColdCodeOnly;
 extern cl::opt<bool> ForcePGSO;
 extern cl::opt<int> PgsoCutoffInstrProf;
 extern cl::opt<int> PgsoCutoffSampleProf;
 
 namespace llvm {
 
 class BasicBlock;
 class BlockFrequencyInfo;
 class Function;
 class ProfileSummaryInfo;
 
 enum class PGSOQueryType {
-  IRPass,  // A query call from an IR-level transform pass.
-  Other,   // Others.
+  IRPass, // A query call from an IR-level transform pass.
+  Test,   // A query call from a unit test.
+  Other,  // Others.
 };
 
 template<typename AdapterT, typename FuncT, typename BFIT>
 bool shouldFuncOptimizeForSizeImpl(const FuncT *F, ProfileSummaryInfo *PSI,
                                    BFIT *BFI, PGSOQueryType QueryType) {
   assert(F);
   if (!PSI || !BFI || !PSI->hasProfileSummary())
     return false;
   if (ForcePGSO)
     return true;
   if (!EnablePGSO)
     return false;
+  // Temporarily enable size optimizations only for the IR pass or test query
+  // sites for gradual commit/rollout. This is to be removed later.
+  if (PGSOIRPassOrTestOnly && !(QueryType == PGSOQueryType::IRPass ||
+                                QueryType == PGSOQueryType::Test))
+    return false;
   if (PGSOColdCodeOnly ||
       (PGSOLargeWorkingSetSizeOnly && !PSI->hasLargeWorkingSetSize())) {
     // Even if the working set size isn't large, size-optimize cold code.
     return AdapterT::isFunctionColdInCallGraph(F, PSI, *BFI);
   }
   return !AdapterT::isFunctionHotInCallGraphNthPercentile(
       PSI->hasSampleProfile() ? PgsoCutoffSampleProf : PgsoCutoffInstrProf,
       F, PSI, *BFI);
 }
 
 template<typename AdapterT, typename BlockT, typename BFIT>
 bool shouldOptimizeForSizeImpl(const BlockT *BB, ProfileSummaryInfo *PSI,
                                BFIT *BFI, PGSOQueryType QueryType) {
   assert(BB);
   if (!PSI || !BFI || !PSI->hasProfileSummary())
     return false;
   if (ForcePGSO)
     return true;
   if (!EnablePGSO)
     return false;
+  // Temporarily enable size optimizations only for the IR pass or test query
+  // sites for gradual commit/rollout. This is to be removed later.
+  if (PGSOIRPassOrTestOnly && !(QueryType == PGSOQueryType::IRPass ||
+                                QueryType == PGSOQueryType::Test))
+    return false;
   if (PGSOColdCodeOnly ||
       (PGSOLargeWorkingSetSizeOnly && !PSI->hasLargeWorkingSetSize())) {
     // Even if the working set size isn't large, size-optimize cold code.
     return AdapterT::isColdBlock(BB, PSI, BFI);
   }
   return !AdapterT::isHotBlockNthPercentile(
       PSI->hasSampleProfile() ? PgsoCutoffSampleProf : PgsoCutoffInstrProf,
       BB, PSI, BFI);

llvm/lib/Transforms/Utils/SizeOpts.cpp

cl::opt<bool> PGSOLargeWorkingSetSizeOnly(
     cl::desc("Apply the profile guided size optimizations only "
              "if the working set size is large (except for cold code.)"));
 
 cl::opt<bool> PGSOColdCodeOnly(
     "pgso-cold-code-only", cl::Hidden, cl::init(true),
     cl::desc("Apply the profile guided size optimizations only "
              "to cold code."));
 
+cl::opt<bool> PGSOIRPassOrTestOnly(
+    "pgso-ir-pass-or-test-only", cl::Hidden, cl::init(true),
+    cl::desc("Apply the profile guided size optimizations only"
+             "to the IR passes or tests."));
+
 cl::opt<bool> ForcePGSO(
     "force-pgso", cl::Hidden, cl::init(false),
     cl::desc("Force the (profiled-guided) size optimizations. "));
 
 cl::opt<int> PgsoCutoffInstrProf(
     "pgso-cutoff-instr-prof", cl::Hidden, cl::init(250000), cl::ZeroOrMore,
     cl::desc("The profile guided size optimization profile summary cutoff "
              "for instrumentation profile."));

llvm/unittests/Target/X86/MachineSizeOptsTest.cpp

TEST_F(MachineSizeOptsTest, Test) {
   auto iter = BB0.succ_begin();
   MachineBasicBlock *BB1 = *iter;
   iter++;
   MachineBasicBlock *BB2 = *iter;
   iter++;
   ASSERT_TRUE(iter == BB0.succ_end());
   MachineBasicBlock *BB3 = *BB1->succ_begin();
   ASSERT_TRUE(BB3 == *BB2->succ_begin());
-  EXPECT_FALSE(shouldOptimizeForSize(F, &PSI, MBFI_F));
-  EXPECT_TRUE(shouldOptimizeForSize(G, &PSI, MBFI_G));
-  EXPECT_FALSE(shouldOptimizeForSize(H, &PSI, MBFI_H));
-  EXPECT_FALSE(shouldOptimizeForSize(&BB0, &PSI, MBFI_F));
-  EXPECT_FALSE(shouldOptimizeForSize(BB1, &PSI, MBFI_F));
-  EXPECT_TRUE(shouldOptimizeForSize(BB2, &PSI, MBFI_F));
-  EXPECT_FALSE(shouldOptimizeForSize(BB3, &PSI, MBFI_F));
+  EXPECT_FALSE(shouldOptimizeForSize(F, &PSI, MBFI_F, PGSOQueryType::Test));
+  EXPECT_TRUE(shouldOptimizeForSize(G, &PSI, MBFI_G, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(H, &PSI, MBFI_H, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(&BB0, &PSI, MBFI_F, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(BB1, &PSI, MBFI_F, PGSOQueryType::Test));
+  EXPECT_TRUE(shouldOptimizeForSize(BB2, &PSI, MBFI_F, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(BB3, &PSI, MBFI_F, PGSOQueryType::Test));
 }
 
 const char* MachineSizeOptsTest::MIRString = R"MIR(
 --- |
   define i32 @g(i32 %x) !prof !14 {
     ret i32 0
   }
 

llvm/unittests/Transforms/Utils/SizeOptsTest.cpp

TEST_F(SizeOptsTest, Test) {
   BlockFrequencyInfo *BFI_G = BFID_G.get();
   BlockFrequencyInfo *BFI_H = BFID_H.get();
   BasicBlock &BB0 = F->getEntryBlock();
   BasicBlock *BB1 = BB0.getTerminator()->getSuccessor(0);
   BasicBlock *BB2 = BB0.getTerminator()->getSuccessor(1);
   BasicBlock *BB3 = BB1->getSingleSuccessor();
 
   EXPECT_TRUE(PSI.hasProfileSummary());
-  EXPECT_FALSE(shouldOptimizeForSize(F, &PSI, BFI_F));
-  EXPECT_TRUE(shouldOptimizeForSize(G, &PSI, BFI_G));
-  EXPECT_FALSE(shouldOptimizeForSize(H, &PSI, BFI_H));
-  EXPECT_FALSE(shouldOptimizeForSize(&BB0, &PSI, BFI_F));
-  EXPECT_FALSE(shouldOptimizeForSize(BB1, &PSI, BFI_F));
-  EXPECT_TRUE(shouldOptimizeForSize(BB2, &PSI, BFI_F));
-  EXPECT_FALSE(shouldOptimizeForSize(BB3, &PSI, BFI_F));
+  EXPECT_FALSE(shouldOptimizeForSize(F, &PSI, BFI_F, PGSOQueryType::Test));
+  EXPECT_TRUE(shouldOptimizeForSize(G, &PSI, BFI_G, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(H, &PSI, BFI_H, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(&BB0, &PSI, BFI_F, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(BB1, &PSI, BFI_F, PGSOQueryType::Test));
+  EXPECT_TRUE(shouldOptimizeForSize(BB2, &PSI, BFI_F, PGSOQueryType::Test));
+  EXPECT_FALSE(shouldOptimizeForSize(BB3, &PSI, BFI_F, PGSOQueryType::Test));
 }
 
 const char* SizeOptsTest::IRString = R"IR(
   define i32 @g(i32 %x) !prof !14 {
     ret i32 0
   }
 
   define i32 @h(i32 %x) !prof !15 {