[ThinLTO] Add module names to ThinLTO final objects

lld/ELF/Config.h

@@ -323,6 +323,7 @@ struct Config {
   bool zText;
   bool zRetpolineplt;
   bool zWxneeded;
+  bool ltoOutputModuleName;
   DiscardPolicy discard;
   GnuStackKind zGnustack;
   ICFLevel icf;

lld/ELF/Driver.cpp

@@ -1317,6 +1317,8 @@ static void readConfigs(opt::InputArgList &args) {
   else
     error("invalid codegen optimization level for LTO: " + Twine(ltoCgo));
   config->ltoObjPath = args.getLastArgValue(OPT_lto_obj_path_eq);
+  config->ltoOutputModuleName = args.hasFlag(
+      OPT_lto_output_module_name, OPT_no_lto_output_module_name, false);
   config->ltoPartitions = args::getInteger(args, OPT_lto_partitions, 1);
   config->ltoSampleProfile = args.getLastArgValue(OPT_lto_sample_profile);
   config->ltoBasicBlockSections =

lld/ELF/LTO.cpp

@@ -351,23 +351,65 @@ std::vector<InputFile *> BitcodeCompiler::compile() {
   if (!config->thinLTOCacheDir.empty())
     pruneCache(config->thinLTOCacheDir, config->thinLTOCachePolicy, files);
 
-  if (!config->ltoObjPath.empty()) {
-    saveBuffer(buf[0], config->ltoObjPath);
-    for (unsigned i = 1; i != maxTasks; ++i)
-      saveBuffer(buf[i], config->ltoObjPath + Twine(i));
-  }
+  auto doSaveBuffer = [&](const StringRef Arg, const StringRef Suffix = "") {
+    // There are a few cases:
+    // (1) path/test.o (using current directory)
+    // (2) /tmp/test-a7a1e4.o (using tmp directory)
+    // (3) if the input obj is in a archive. the module name is like
+    // "arch/x86/built-in.a(procfs.o at 11368)"
+    //
+    // This function replaces '/' and '(' with '-' and terminates at the
+    // last '.'.  it returns the following for the above cases, respectively,
+    // (1) path_test
+    // (2) tmp_test-a7a1e4 (remove the first /).
+    // (3) arch_x86_build-in.a_procfs
+    //
+    auto getFileNameString = [](const StringRef Str) {
+      if (Str.empty())
+        return std::string();
+      size_t End = Str.find_last_of(".");
+      size_t Begin = 0;
+      if (Str[0] == '/' || Str[0] == '\\')
+        Begin = 1;
+      std::string Ret = Str.substr(Begin, End - Begin).str();
+      auto position = std::string::npos;
+      while ((position = Ret.find_first_of("/\\(")) != std::string::npos) {
+        Ret.replace(position, 1, 1, '_');
+      }
+      return Ret;
+    };
+
+    auto saveBufferOrFile = [](const StringRef &Buf, const MemoryBuffer *File,
+                               const Twine &Path) {
+      if (Buf.empty() && File)
+        return saveBuffer(File->getBuffer(), Path);
+      saveBuffer(Buf, Path);
+    };
 
-  if (config->saveTempsArgs.contains("prelink")) {
     if (!buf[0].empty())
-      saveBuffer(buf[0], config->outputFile + ".lto.o");
-    for (unsigned i = 1; i != maxTasks; ++i)
-      saveBuffer(buf[i], config->outputFile + Twine(i) + ".lto.o");
-  }
+      saveBufferOrFile(buf[0], files[0].get(), Arg + Suffix);
+    for (unsigned i = 1; i != maxTasks; ++i) {
+      if (!config->ltoOutputModuleName) {
+        saveBufferOrFile(buf[i], files[i].get(), Arg + Twine(i) + Suffix);
+      } else {
+        const std::string Name =
+            getFileNameString(ltoObj->getModuleName(i - 1));
+        saveBufferOrFile(buf[i], files[i].get(),
+                         Arg + "_" + Twine(i) + "_" + Name + Suffix);
+      }
+    }
+  };
+
+  if (!config->ltoObjPath.empty())
+    doSaveBuffer(config->ltoObjPath,
+                 config->ltoOutputModuleName ? ".lto.o" : "");
+
+  if (config->saveTempsArgs.contains("prelink"))
+    doSaveBuffer(config->outputFile, ".lto.o");
 
   if (config->ltoEmitAsm) {
-    saveBuffer(buf[0], config->outputFile);
-    for (unsigned i = 1; i != maxTasks; ++i)
-      saveBuffer(buf[i], config->outputFile + Twine(i));
+    doSaveBuffer(config->outputFile,
+                 config->ltoOutputModuleName ? ".lto.s" : "");
     return {};
   }
 

lld/ELF/Options.td

@@ -610,6 +610,9 @@ defm lto_pgo_warn_mismatch: BB<"lto-pgo-warn-mismatch",
 defm lto_known_safe_vtables : EEq<"lto-known-safe-vtables",
   "When --lto-validate-all-vtables-have-type-infos is enabled, skip validation on these vtables (_ZTV symbols)">;
 def lto_obj_path_eq: JJ<"lto-obj-path=">;
+defm lto_output_module_name: BB<"lto-output-module-name",
+  "In ThinLTO, using the module name in the final native objects or asm files",
+  "Do not include the module names in final objects and asm files">;
 def lto_sample_profile: JJ<"lto-sample-profile=">,
   HelpText<"Sample profile file path">;
 defm lto_validate_all_vtables_have_type_infos: BB<"lto-validate-all-vtables-have-type-infos",

lld/test/ELF/lto/thinlto_finallink_output.ll

@@ -0,0 +1,59 @@
+; REQUIRES: x86
+;
+; RUN: cd %T
+; RUN: opt -module-summary %s -o obj1.o
+; RUN: opt -module-summary %p/Inputs/thinlto.ll -o obj2.o
+; RUN: opt -module-summary %s -o %t_obj3.o
+; RUN: opt -module-summary %p/Inputs/thinlto.ll -o %t_obj4.o
+;
+; Objects with a relative path.
+; RUN: rm -f *.lto.o *.s
+; RUN: ld.lld --lto-output-module-name --save-temps=prelink --lto-obj-path=aaa --lto-emit-asm --thinlto-jobs=1 --entry=f obj1.o obj2.o -o bin1
+; RUN: ls -1 *.lto.o *.s | FileCheck %s --check-prefixes=OBJPATHOUT1,PRELINKOUT1
+; With thinlto-jobs=all.
+; RUN: rm -f *.lto.o *.s
+; RUN: ld.lld --lto-output-module-name --save-temps=prelink --lto-obj-path=aaa --lto-emit-asm --thinlto-jobs=all --entry=f obj1.o obj2.o -o bin1
+; RUN: ls -1 *.lto.o *.s | FileCheck %s --check-prefixes=OBJPATHOUT1,PRELINKOUT1
+; Objects with an absolute path.
+; RUN: rm -f *.lto.o *.s
+; RUN: ld.lld --lto-output-module-name --save-temps=prelink --lto-obj-path=aaa --lto-emit-asm --thinlto-jobs=1 --entry=f %t_obj3.o %t_obj4.o -o bin2
+; RUN: ls -1 *.lto.o *.s | FileCheck %s --check-prefixes=OBJPATHOUT2,PRELINKOUT2
+; Objects in an archive
+; RUN: rm -f *.lto.o *.s
+; RUN: llvm-ar rcS ar.a obj1.o obj2.o
+; RUN: ld.lld --lto-output-module-name --save-temps=prelink --lto-obj-path=aaa --lto-emit-asm --thinlto-jobs=1 --entry=f ar.a -o bin1
+; RUN: ls -1 *.lto.o *.s | FileCheck %s --check-prefixes=OBJPATHOUT3,PRELINKOUT3
+; Use with thinlto-cahce
+; RUN: rm -f *.lto.o *.s
+; RUN: ld.lld --lto-output-module-name --save-temps=prelink --thinlto-cache-dir=thinlto-cache --lto-emit-asm --lto-obj-path=aaa --thinlto-jobs=1 --entry=f obj1.o obj2.o -o bin1
+; RUN: ls -1 *.lto.o *.s | FileCheck %s --check-prefixes=OBJPATHOUT1,PRELINKOUT1
+;
+; OBJPATHOUT1-DAG: aaa_1_obj1.lto.o
+; OBJPATHOUT1-DAG: aaa_2_obj2.lto.o
+; PRELINKOUT1-DAG: bin1_1_obj1.lto.o
+; PRELINKOUT1-DAG: bin1_2_obj2.lto.o
+; PRELINKOUT1-DAG: bin1_1_obj1.lto.s
+; PRELINKOUT1-DAG: bin1_2_obj2.lto.s
+; OBJPATHOUT2-DAG: aaa_1_{{.*}}_obj3.lto.o
+; OBJPATHOUT2-DAG: aaa_2_{{.*}}obj4.lto.o
+; PRELINKOUT2-DAG: bin2_1_{{.*}}obj3.lto.o
+; PRELINKOUT2-DAG: bin2_2_{{.*}}obj4.lto.o
+; PRELINKOUT2-DAG: bin2_1_{{.*}}obj3.lto.s
+; PRELINKOUT2-DAG: bin2_2_{{.*}}obj4.lto.s
+; OBJPATHOUT3-DAG: aaa_1_ar.a_obj1.lto.o
+; OBJPATHOUT3-DAG: aaa_2_ar.a_obj2.lto.o
+; PRELINKOUT3-DAG: bin1_1_ar.a_obj1.lto.o
+; PRELINKOUT3-DAG: bin1_2_ar.a_obj2.lto.o
+; PRELINKOUT3-DAG: bin1_1_ar.a_obj1.lto.s
+; PRELINKOUT3-DAG: bin1_2_ar.a_obj2.lto.s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @g(...)
+
+define void @f() {
+entry:
+  call void (...) @g()
+  ret void
+}

llvm/include/llvm/LTO/LTO.h

@@ -299,7 +299,18 @@ class LTO {
 
   /// Static method that returns a list of libcall symbols that can be generated
   /// by LTO but might not be visible from bitcode symbol table.
-  static ArrayRef<const char*> getRuntimeLibcallSymbols();
+  static ArrayRef<const char *> getRuntimeLibcallSymbols();
+
+  /// Return the name of n-th module. This only applies to ThinLTO.
+  StringRef getModuleName(size_t N) const {
+    size_t I = N;
+    if (I >= ThinLTO.ModuleMap.size())
+      return "";
+    auto it = ThinLTO.ModuleMap.begin();
+    while (I--)
+      it++;
+    return (*it).first;
+  }
 
 private:
   Config Conf;

llvm/include/llvm/ProfileData/InstrProfData.inc

@@ -676,6 +676,7 @@ serializeValueProfDataFrom(ValueProfRecordClosure *Closure,
 #define INSTR_PROF_PROFILE_RUNTIME_VAR __llvm_profile_runtime
 #define INSTR_PROF_PROFILE_COUNTER_BIAS_VAR __llvm_profile_counter_bias
 #define INSTR_PROF_PROFILE_SET_TIMESTAMP __llvm_profile_set_timestamp
+#define INSTR_PROF_PROFILE_SAMPLING_VAR __llvm_profile_sampling
 
 /* The variable that holds the name of the profile data
  * specified via command line. */

llvm/include/llvm/Transforms/Instrumentation.h

@@ -116,12 +116,18 @@ struct InstrProfOptions {
   // Use BFI to guide register promotion
   bool UseBFIInPromotion = false;
 
+  // Use sampling to reduce the profile instrumentation runtime overhead.
+  bool Sampling = false;
+
   // Name of the profile file to use as output
   std::string InstrProfileOutput;
 
   InstrProfOptions() = default;
 };
 
+// Create the variable for profile sampling.
+void createProfileSamplingVar(Module &M);
+
 // Options for sanitizer coverage instrumentation.
 struct SanitizerCoverageOptions {
   enum Type {

llvm/include/llvm/Transforms/Instrumentation/InstrProfiling.h

@@ -86,6 +86,9 @@ class InstrProfiling : public PassInfoMixin<InstrProfiling> {
   /// Returns true if profile counter update register promotion is enabled.
   bool isCounterPromotionEnabled() const;
 
+  /// Return true if profile sampling is enabled.
+  bool isSamplingEnabled() const;
+
   /// Count the number of instrumented value sites for the function.
   void computeNumValueSiteCounts(InstrProfValueProfileInst *Ins);
 
@@ -109,6 +112,9 @@ class InstrProfiling : public PassInfoMixin<InstrProfiling> {
   /// acts on.
   Value *getCounterAddress(InstrProfInstBase *I);
 
+  /// Lower the incremental instructions under profile sampling predicates.
+  void doSampling(Instruction *I);
+
   /// Get the region counters for an increment, creating them if necessary.
   ///
   /// If the counter array doesn't yet exist, the profile data variables

llvm/include/llvm/Transforms/Instrumentation/PGOInstrumentation.h

@@ -43,12 +43,14 @@ class FileSystem;
 class PGOInstrumentationGenCreateVar
     : public PassInfoMixin<PGOInstrumentationGenCreateVar> {
 public:
-  PGOInstrumentationGenCreateVar(std::string CSInstrName = "")
-      : CSInstrName(CSInstrName) {}
+  PGOInstrumentationGenCreateVar(std::string CSInstrName = "",
+                                 bool Sampling = false)
+      : CSInstrName(CSInstrName), ProfileSampling(Sampling) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
 
 private:
   std::string CSInstrName;
+  bool ProfileSampling;
 };
 
 /// The instrumentation (profile-instr-gen) pass for IR based PGO.

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -273,6 +273,9 @@ static cl::opt<AttributorRunOption> AttributorRun(
                clEnumValN(AttributorRunOption::NONE, "none",
                           "disable attributor runs")));
 
+static cl::opt<bool> EnableSampledInstr(
+    "enable-sampled-instr", cl::init(false), cl::Hidden,
+    cl::desc("Enable profile instrumentation sampling (default = off)"));
 static cl::opt<bool> UseLoopVersioningLICM(
     "enable-loop-versioning-licm", cl::init(false), cl::Hidden,
     cl::desc("Enable the experimental Loop Versioning LICM pass"));
@@ -805,6 +808,10 @@ void PassBuilder::addPGOInstrPasses(ModulePassManager &MPM,
   // Do counter promotion at Level greater than O0.
   Options.DoCounterPromotion = true;
   Options.UseBFIInPromotion = IsCS;
+  if (EnableSampledInstr) {
+    Options.Sampling = true;
+    Options.DoCounterPromotion = false;
+  }
   Options.Atomic = AtomicCounterUpdate;
   MPM.addPass(InstrProfiling(Options, IsCS));
 }
@@ -1117,7 +1124,8 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
   }
   if (PGOOpt && Phase != ThinOrFullLTOPhase::ThinLTOPostLink &&
       PGOOpt->CSAction == PGOOptions::CSIRInstr)
-    MPM.addPass(PGOInstrumentationGenCreateVar(PGOOpt->CSProfileGenFile));
+    MPM.addPass(PGOInstrumentationGenCreateVar(PGOOpt->CSProfileGenFile,
+                                               EnableSampledInstr));
 
   if (PGOOpt && Phase != ThinOrFullLTOPhase::ThinLTOPostLink &&
       !PGOOpt->MemoryProfile.empty())