From af58684f272046f293a9f469f03d23bd2b138349 Mon Sep 17 00:00:00 2001
From: Ellis Hoag <ellis.sparky.hoag@gmail.com>
Date: Thu, 14 Jul 2022 11:40:53 -0700
Subject: [PATCH] [InstrProf] Add options to profile function groups

Add two options, `-fprofile-function-groups=N` and `-fprofile-selected-function-group=i` used to partition functions into `N` groups and only instrument the functions in group `i`. Similar options were added to xray in https://reviews.llvm.org/D87953 and the goal is the same; to reduce instrumented size overhead by spreading the overhead across multiple builds. Raw profiles from different groups can be added like normal using the `llvm-profdata merge` command.

Reviewed By: ianlevesque

Differential Revision: https://reviews.llvm.org/D129594
---
 clang/docs/ClangCommandLineReference.rst     |  4 +++
 clang/docs/UsersManual.rst                   | 26 ++++++++++++++++++
 clang/include/clang/Basic/CodeGenOptions.def |  4 +++
 clang/include/clang/Driver/Options.td        |  9 +++++++
 clang/lib/CodeGen/CodeGenFunction.cpp        |  2 +-
 clang/lib/CodeGen/CodeGenModule.cpp          | 19 +++++++++++--
 clang/lib/CodeGen/CodeGenModule.h            | 12 ++++++---
 clang/lib/Driver/ToolChains/Clang.cpp        | 21 +++++++++++++++
 clang/test/CodeGen/profile-function-groups.c | 24 +++++++++++++++++
 compiler-rt/test/profile/instrprof-groups.c  | 28 ++++++++++++++++++++
 10 files changed, 143 insertions(+), 6 deletions(-)
 create mode 100644 clang/test/CodeGen/profile-function-groups.c
 create mode 100644 compiler-rt/test/profile/instrprof-groups.c

diff --git a/clang/docs/ClangCommandLineReference.rst b/clang/docs/ClangCommandLineReference.rst
index 776b84da9657..216872b60cdc 100644
--- a/clang/docs/ClangCommandLineReference.rst
+++ b/clang/docs/ClangCommandLineReference.rst
@@ -2329,6 +2329,10 @@ Use instrumentation data for profile-guided optimization
 
 Filename defining the list of functions/files to instrument
 
+.. option:: -fprofile-function-groups=<N>, -fprofile-selected-function-group=<i>
+
+Partition functions into <N> groups and select only functions in group <i> to be instrumented
+
 .. option:: -fprofile-remapping-file=<file>
 
 Use the remappings described in <file> to match the profile data against names in the program
diff --git a/clang/docs/UsersManual.rst b/clang/docs/UsersManual.rst
index e12dc72407b1..c2767d65adbf 100644
--- a/clang/docs/UsersManual.rst
+++ b/clang/docs/UsersManual.rst
@@ -2513,6 +2513,32 @@ When the file contains only excludes, all files and functions except for the
 excluded ones will be instrumented. Otherwise, only the files and functions
 specified will be instrumented.
 
+Instrument function groups
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Sometimes it is desirable to minimize the size overhead of instrumented
+binaries. One way to do this is to partition functions into groups and only
+instrument functions in a specified group. This can be done using the
+`-fprofile-function-groups` and `-fprofile-selected-function-group` options.
+
+.. option:: -fprofile-function-groups=<N>, -fprofile-selected-function-group=<i>
+
+  The following uses 3 groups
+
+  .. code-block:: console
+
+    $ clang++ -Oz -fprofile-generate=group_0/ -fprofile-function-groups=3 -fprofile-selected-function-group=0 code.cc -o code.0
+    $ clang++ -Oz -fprofile-generate=group_1/ -fprofile-function-groups=3 -fprofile-selected-function-group=1 code.cc -o code.1
+    $ clang++ -Oz -fprofile-generate=group_2/ -fprofile-function-groups=3 -fprofile-selected-function-group=2 code.cc -o code.2
+
+  After collecting raw profiles from the three binaries, they can be merged into
+  a single profile like normal.
+
+  .. code-block:: console
+
+    $ llvm-profdata merge -output=code.profdata group_*/*.profraw
+
+
 Profile remapping
 ^^^^^^^^^^^^^^^^^
 
diff --git a/clang/include/clang/Basic/CodeGenOptions.def b/clang/include/clang/Basic/CodeGenOptions.def
index b1d394edd04a..ef7957979dcc 100644
--- a/clang/include/clang/Basic/CodeGenOptions.def
+++ b/clang/include/clang/Basic/CodeGenOptions.def
@@ -213,6 +213,10 @@ CODEGENOPT(AtomicProfileUpdate , 1, 0) ///< Set -fprofile-update=atomic
 ENUM_CODEGENOPT(ProfileInstr, ProfileInstrKind, 2, ProfileNone)
 /// Choose profile kind for PGO use compilation.
 ENUM_CODEGENOPT(ProfileUse, ProfileInstrKind, 2, ProfileNone)
+/// Partition functions into N groups and select only functions in group i to be
+/// instrumented. Selected group numbers can be 0 to N-1 inclusive.
+VALUE_CODEGENOPT(ProfileTotalFunctionGroups, 32, 1)
+VALUE_CODEGENOPT(ProfileSelectedFunctionGroup, 32, 0)
 CODEGENOPT(CoverageMapping , 1, 0) ///< Generate coverage mapping regions to
                                    ///< enable code coverage analysis.
 CODEGENOPT(DumpCoverageMapping , 1, 0) ///< Dump the generated coverage mapping
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
index 532d7780c529..404effb4e1de 100644
--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@@ -1333,6 +1333,15 @@ def fprofile_list_EQ : Joined<["-"], "fprofile-list=">,
     Group<f_Group>, Flags<[CC1Option, CoreOption]>,
     HelpText<"Filename defining the list of functions/files to instrument">,
     MarshallingInfoStringVector<LangOpts<"ProfileListFiles">>;
+def fprofile_function_groups : Joined<["-"], "fprofile-function-groups=">,
+  Group<f_Group>, Flags<[CC1Option]>, MetaVarName<"<N>">,
+  HelpText<"Partition functions into N groups and select only functions in group i to be instrumented using -fprofile-selected-function-group">,
+  MarshallingInfoInt<CodeGenOpts<"ProfileTotalFunctionGroups">, "1">;
+def fprofile_selected_function_group :
+  Joined<["-"], "fprofile-selected-function-group=">, Group<f_Group>,
+  Flags<[CC1Option]>, MetaVarName<"<i>">,
+  HelpText<"Partition functions into N groups using -fprofile-function-groups and select only functions in group i to be instrumented. The valid range is 0 to N-1 inclusive">,
+  MarshallingInfoInt<CodeGenOpts<"ProfileSelectedFunctionGroup">>;
 def fswift_async_fp_EQ : Joined<["-"], "fswift-async-fp=">,
     Group<f_Group>, Flags<[CC1Option, CC1AsOption, CoreOption]>, MetaVarName<"<option>">,
     HelpText<"Control emission of Swift async extended frame info">,
diff --git a/clang/lib/CodeGen/CodeGenFunction.cpp b/clang/lib/CodeGen/CodeGenFunction.cpp
index 17c1c91c7e8f..5012bd822bd3 100644
--- a/clang/lib/CodeGen/CodeGenFunction.cpp
+++ b/clang/lib/CodeGen/CodeGenFunction.cpp
@@ -852,7 +852,7 @@ void CodeGenFunction::StartFunction(GlobalDecl GD, QualType RetTy,
   }
 
   if (CGM.getCodeGenOpts().getProfileInstr() != CodeGenOptions::ProfileNone)
-    if (CGM.isProfileInstrExcluded(Fn, Loc))
+    if (CGM.isFunctionBlockedFromProfileInstr(Fn, Loc))
       Fn->addFnAttr(llvm::Attribute::NoProfile);
 
   unsigned Count, Offset;
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index bdee31b504ae..1e90b0914538 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -58,6 +58,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/InstrProfReader.h"
+#include "llvm/Support/CRC.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ConvertUTF.h"
@@ -2843,8 +2844,8 @@ bool CodeGenModule::imbueXRayAttrs(llvm::Function *Fn, SourceLocation Loc,
   return true;
 }
 
-bool CodeGenModule::isProfileInstrExcluded(llvm::Function *Fn,
-                                           SourceLocation Loc) const {
+bool CodeGenModule::isFunctionBlockedByProfileList(llvm::Function *Fn,
+                                                   SourceLocation Loc) const {
   const auto &ProfileList = getContext().getProfileList();
   // If the profile list is empty, then instrument everything.
   if (ProfileList.isEmpty())
@@ -2871,6 +2872,20 @@ bool CodeGenModule::isProfileInstrExcluded(llvm::Function *Fn,
   return ProfileList.getDefault();
 }
 
+bool CodeGenModule::isFunctionBlockedFromProfileInstr(
+    llvm::Function *Fn, SourceLocation Loc) const {
+  if (isFunctionBlockedByProfileList(Fn, Loc))
+    return true;
+
+  auto NumGroups = getCodeGenOpts().ProfileTotalFunctionGroups;
+  if (NumGroups > 1) {
+    auto Group = llvm::crc32(arrayRefFromStringRef(Fn->getName())) % NumGroups;
+    if (Group != getCodeGenOpts().ProfileSelectedFunctionGroup)
+      return true;
+  }
+  return false;
+}
+
 bool CodeGenModule::MustBeEmitted(const ValueDecl *Global) {
   // Never defer when EmitAllDecls is specified.
   if (LangOpts.EmitAllDecls)
diff --git a/clang/lib/CodeGen/CodeGenModule.h b/clang/lib/CodeGen/CodeGenModule.h
index 10b49da27dab..1123a7cfce15 100644
--- a/clang/lib/CodeGen/CodeGenModule.h
+++ b/clang/lib/CodeGen/CodeGenModule.h
@@ -1340,9 +1340,15 @@ public:
   bool imbueXRayAttrs(llvm::Function *Fn, SourceLocation Loc,
                       StringRef Category = StringRef()) const;
 
-  /// Returns true if function at the given location should be excluded from
-  /// profile instrumentation.
-  bool isProfileInstrExcluded(llvm::Function *Fn, SourceLocation Loc) const;
+  /// \returns true if \p Fn at \p Loc should be excluded from profile
+  /// instrumentation by the SCL passed by \p -fprofile-list.
+  bool isFunctionBlockedByProfileList(llvm::Function *Fn,
+                                      SourceLocation Loc) const;
+
+  /// \returns true if \p Fn at \p Loc should be excluded from profile
+  /// instrumentation.
+  bool isFunctionBlockedFromProfileInstr(llvm::Function *Fn,
+                                         SourceLocation Loc) const;
 
   SanitizerMetadata *getSanitizerMetadata() {
     return SanitizerMD.get();
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
index 97435f1a73de..3e36d4def2f3 100644
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -956,6 +956,27 @@ static void addPGOAndCoverageFlags(const ToolChain &TC, Compilation &C,
     CmdArgs.push_back("-fprofile-update=atomic");
   }
 
+  int FunctionGroups = 1;
+  int SelectedFunctionGroup = 0;
+  if (const auto *A = Args.getLastArg(options::OPT_fprofile_function_groups)) {
+    StringRef Val = A->getValue();
+    if (Val.getAsInteger(0, FunctionGroups) || FunctionGroups < 1)
+      D.Diag(diag::err_drv_invalid_int_value) << A->getAsString(Args) << Val;
+  }
+  if (const auto *A =
+          Args.getLastArg(options::OPT_fprofile_selected_function_group)) {
+    StringRef Val = A->getValue();
+    if (Val.getAsInteger(0, SelectedFunctionGroup) ||
+        SelectedFunctionGroup < 0 || SelectedFunctionGroup >= FunctionGroups)
+      D.Diag(diag::err_drv_invalid_int_value) << A->getAsString(Args) << Val;
+  }
+  if (FunctionGroups != 1)
+    CmdArgs.push_back(Args.MakeArgString("-fprofile-function-groups=" +
+                                         Twine(FunctionGroups)));
+  if (SelectedFunctionGroup != 0)
+    CmdArgs.push_back(Args.MakeArgString("-fprofile-selected-function-group=" +
+                                         Twine(SelectedFunctionGroup)));
+
   // Leave -fprofile-dir= an unused argument unless .gcda emission is
   // enabled. To be polite, with '-fprofile-arcs -fno-profile-arcs' consider
   // the flag used. There is no -fno-profile-dir, so the user has no
diff --git a/clang/test/CodeGen/profile-function-groups.c b/clang/test/CodeGen/profile-function-groups.c
new file mode 100644
index 000000000000..9e04c9060df4
--- /dev/null
+++ b/clang/test/CodeGen/profile-function-groups.c
@@ -0,0 +1,24 @@
+// RUN: %clang -fprofile-generate -fprofile-function-groups=3 -fprofile-selected-function-group=0 -emit-llvm -S %s -o - | FileCheck %s --check-prefixes=CHECK,SELECT0
+// RUN: %clang -fprofile-generate -fprofile-function-groups=3 -fprofile-selected-function-group=1 -emit-llvm -S %s -o - | FileCheck %s --check-prefixes=CHECK,SELECT1
+// RUN: %clang -fprofile-generate -fprofile-function-groups=3 -fprofile-selected-function-group=2 -emit-llvm -S %s -o - | FileCheck %s --check-prefixes=CHECK,SELECT2
+
+// Group 0
+// SELECT0-NOT: noprofile
+// SELECT1: noprofile
+// SELECT2: noprofile
+// CHECK: define {{.*}} @hoo()
+void hoo() {}
+
+// Group 1
+// SELECT0: noprofile
+// SELECT1-NOT: noprofile
+// SELECT2: noprofile
+// CHECK: define {{.*}} @goo()
+void goo() {}
+
+// Group 2
+// SELECT0: noprofile
+// SELECT1: noprofile
+// SELECT2-NOT: noprofile
+// CHECK: define {{.*}} @boo()
+void boo() {}
diff --git a/compiler-rt/test/profile/instrprof-groups.c b/compiler-rt/test/profile/instrprof-groups.c
new file mode 100644
index 000000000000..1a56453da914
--- /dev/null
+++ b/compiler-rt/test/profile/instrprof-groups.c
@@ -0,0 +1,28 @@
+// RUN: %clang_pgogen -fprofile-function-groups=3 -fprofile-selected-function-group=0 %s -o %t.0.out
+// RUN: %clang_pgogen -fprofile-function-groups=3 -fprofile-selected-function-group=1 %s -o %t.1.out
+// RUN: %clang_pgogen -fprofile-function-groups=3 -fprofile-selected-function-group=2 %s -o %t.2.out
+// RUN: env LLVM_PROFILE_FILE=%t.0.profraw %run %t.0.out
+// RUN: env LLVM_PROFILE_FILE=%t.1.profraw %run %t.1.out
+// RUN: env LLVM_PROFILE_FILE=%t.2.profraw %run %t.2.out
+// RUN: llvm-profdata merge -o %t.profdata %t.*.profraw
+// RUN: llvm-profdata show %t.profdata --all-functions | FileCheck %s
+
+int foo(int i) { return 4 * i + 1; }
+int bar(int i) { return 4 * i + 2; }
+int goo(int i) { return 4 * i + 3; }
+
+int main(int argc, char *argv[]) {
+  foo(5);
+  bar(6);
+  goo(7);
+  return 0;
+}
+
+// Even though we ran this code three times, we expect all counts to be one if
+// functions were partitioned into groups correctly.
+
+// CHECK: Counters: 1
+// CHECK: Counters: 1
+// CHECK: Counters: 1
+// CHECK: Counters: 1
+// CHECK: Total functions: 4