[PseudoProbe] Support emitting to COFF object

Support emitting pseudo probe to .pseudo_probe and .pseudo_probe_desc
COFF sections.

Author: HaohaiWen

clang/include/clang/Driver/Options.td

@@ -1890,7 +1890,7 @@ defm pseudo_probe_for_profiling : BoolFOption<"pseudo-probe-for-profiling",
   CodeGenOpts<"PseudoProbeForProfiling">, DefaultFalse,
   PosFlag<SetTrue, [], [ClangOption], "Emit">,
   NegFlag<SetFalse, [], [ClangOption], "Do not emit">,
-  BothFlags<[], [ClangOption, CC1Option],
+  BothFlags<[], [ClangOption, CC1Option, CLOption],
           " pseudo probes for sample profiling">>;
 def forder_file_instrumentation : Flag<["-"], "forder-file-instrumentation">,
     Group<f_Group>, Visibility<[ClangOption, CC1Option, CLOption]>,

clang/test/Driver/cl-options.c

@@ -720,6 +720,7 @@
 // RUN:     -fno-profile-instr-use \
 // RUN:     -fcs-profile-generate \
 // RUN:     -fcs-profile-generate=dir \
+// RUN:     -fpseudo-probe-for-profiling \
 // RUN:     -ftime-trace \
 // RUN:     -fmodules \
 // RUN:     -fno-modules \

llvm/include/llvm/Target/TargetLoweringObjectFile.h

@@ -91,6 +91,9 @@ class TargetLoweringObjectFile : public MCObjectFileInfo {
   /// Emit Call Graph Profile metadata.
   void emitCGProfileMetadata(MCStreamer &Streamer, Module &M) const;
 
+  /// Emit pseudo_probe_desc metadata.
+  void emitPseudoProbeDescMetadata(MCStreamer &Streamer, Module &M) const;
+
   /// Get the module-level metadata that the platform cares about.
   virtual void getModuleMetadata(Module &M) {}
 

llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp

@@ -335,28 +335,7 @@ void TargetLoweringObjectFileELF::emitModuleMetadata(MCStreamer &Streamer,
     }
   }
 
-  if (NamedMDNode *FuncInfo = M.getNamedMetadata(PseudoProbeDescMetadataName)) {
-    // Emit a descriptor for every function including functions that have an
-    // available external linkage. We may not want this for imported functions
-    // that has code in another thinLTO module but we don't have a good way to
-    // tell them apart from inline functions defined in header files. Therefore
-    // we put each descriptor in a separate comdat section and rely on the
-    // linker to deduplicate.
-    for (const auto *Operand : FuncInfo->operands()) {
-      const auto *MD = cast<MDNode>(Operand);
-      auto *GUID = mdconst::dyn_extract<ConstantInt>(MD->getOperand(0));
-      auto *Hash = mdconst::dyn_extract<ConstantInt>(MD->getOperand(1));
-      auto *Name = cast<MDString>(MD->getOperand(2));
-      auto *S = C.getObjectFileInfo()->getPseudoProbeDescSection(
-          TM->getFunctionSections() ? Name->getString() : StringRef());
-
-      Streamer.switchSection(S);
-      Streamer.emitInt64(GUID->getZExtValue());
-      Streamer.emitInt64(Hash->getZExtValue());
-      Streamer.emitULEB128IntValue(Name->getString().size());
-      Streamer.emitBytes(Name->getString());
-    }
-  }
+  emitPseudoProbeDescMetadata(Streamer, M);
 
   if (NamedMDNode *LLVMStats = M.getNamedMetadata("llvm.stats")) {
     // Emit the metadata for llvm statistics into .llvm_stats section, which is
@@ -1864,6 +1843,7 @@ void TargetLoweringObjectFileCOFF::emitModuleMetadata(MCStreamer &Streamer,
   }
 
   emitCGProfileMetadata(Streamer, M);
+  emitPseudoProbeDescMetadata(Streamer, M);
 }
 
 void TargetLoweringObjectFileCOFF::emitLinkerDirectives(

llvm/lib/MC/MCAsmStreamer.cpp

@@ -2492,7 +2492,8 @@ void MCAsmStreamer::emitPseudoProbe(uint64_t Guid, uint64_t Index,
   for (const auto &Site : InlineStack)
     OS << " @ " << std::get<0>(Site) << ":" << std::get<1>(Site);
 
-  OS << " " << FnSym->getName();
+  OS << " ";
+  FnSym->print(OS, MAI);
 
   EmitEOL();
 }

llvm/lib/MC/MCContext.cpp

@@ -23,6 +23,7 @@
 #include "llvm/MC/MCFragment.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCLabel.h"
+#include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCSectionCOFF.h"
 #include "llvm/MC/MCSectionDXContainer.h"
 #include "llvm/MC/MCSectionELF.h"
@@ -722,6 +723,13 @@ MCSectionCOFF *MCContext::getCOFFSection(StringRef Section,
   Iter->second = Result;
   auto *F = allocInitialFragment(*Result);
   Begin->setFragment(F);
+  // Normally the comdat symbol is function begin label and will be set a
+  // fragment in emitLabel. It is not hold for a pseudo_probe_desc comdat
+  // symbol, so we need to set its fragment here.
+  if (COMDATSymbol && !COMDATSymbol->getFragment() &&
+      Section == MOFI->getPseudoProbeDescSection("")->getName()) {
+    COMDATSymbol->setFragment(F);
+  }
   return Result;
 }
 

llvm/lib/MC/MCObjectFileInfo.cpp

@@ -792,6 +792,16 @@ void MCObjectFileInfo::initCOFFMCObjectFileInfo(const Triple &T) {
   StackMapSection = Ctx->getCOFFSection(".llvm_stackmaps",
                                         COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
                                             COFF::IMAGE_SCN_MEM_READ);
+
+  // Set IMAGE_SCN_MEM_DISCARDABLE so that lld will not truncate section name.
+  PseudoProbeSection = Ctx->getCOFFSection(
+      ".pseudo_probe", COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
+                           COFF::IMAGE_SCN_MEM_DISCARDABLE |
+                           COFF::IMAGE_SCN_MEM_READ);
+  PseudoProbeDescSection = Ctx->getCOFFSection(
+      ".pseudo_probe_desc", COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
+                                COFF::IMAGE_SCN_MEM_DISCARDABLE |
+                                COFF::IMAGE_SCN_MEM_READ);
 }
 
 void MCObjectFileInfo::initSPIRVMCObjectFileInfo(const Triple &T) {
@@ -1141,41 +1151,68 @@ MCObjectFileInfo::getKCFITrapSection(const MCSection &TextSec) const {
 
 MCSection *
 MCObjectFileInfo::getPseudoProbeSection(const MCSection &TextSec) const {
-  if (Ctx->getObjectFileType() != MCContext::IsELF)
-    return PseudoProbeSection;
+  auto ObjFileType = Ctx->getObjectFileType();
+  if (ObjFileType == MCContext::IsELF) {
+    const auto &ElfSec = static_cast<const MCSectionELF &>(TextSec);
+    unsigned Flags = ELF::SHF_LINK_ORDER;
+    StringRef GroupName;
+    if (const MCSymbol *Group = ElfSec.getGroup()) {
+      GroupName = Group->getName();
+      Flags |= ELF::SHF_GROUP;
+    }
 
-  const auto &ElfSec = static_cast<const MCSectionELF &>(TextSec);
-  unsigned Flags = ELF::SHF_LINK_ORDER;
-  StringRef GroupName;
-  if (const MCSymbol *Group = ElfSec.getGroup()) {
-    GroupName = Group->getName();
-    Flags |= ELF::SHF_GROUP;
+    return Ctx->getELFSection(PseudoProbeSection->getName(), ELF::SHT_PROGBITS,
+                              Flags, 0, GroupName, true, ElfSec.getUniqueID(),
+                              cast<MCSymbolELF>(TextSec.getBeginSymbol()));
+  } else if (ObjFileType == MCContext::IsCOFF) {
+    StringRef COMDATSymName = "";
+    int Selection = 0;
+    unsigned Characteristics =
+        COFF::IMAGE_SCN_CNT_INITIALIZED_DATA | COFF::IMAGE_SCN_MEM_DISCARDABLE |
+        COFF::IMAGE_SCN_MEM_READ | COFF::IMAGE_SCN_LNK_COMDAT;
+    auto &COFFSec = cast<MCSectionCOFF>(TextSec);
+    if (const MCSymbol *COMDATSym = COFFSec.getCOMDATSymbol()) {
+      // Associate .pseudo_probe to its function section.
+      COMDATSymName = COMDATSym->getName();
+      Selection = COFF::IMAGE_COMDAT_SELECT_ASSOCIATIVE;
+    }
+
+    return Ctx->getCOFFSection(PseudoProbeSection->getName(), Characteristics,
+                               COMDATSymName, Selection, COFFSec.getUniqueID());
   }
 
-  return Ctx->getELFSection(PseudoProbeSection->getName(), ELF::SHT_PROGBITS,
-                            Flags, 0, GroupName, true, ElfSec.getUniqueID(),
-                            cast<MCSymbolELF>(TextSec.getBeginSymbol()));
+  return PseudoProbeSection;
 }
 
 MCSection *
 MCObjectFileInfo::getPseudoProbeDescSection(StringRef FuncName) const {
-  if (Ctx->getObjectFileType() == MCContext::IsELF) {
-    // Create a separate comdat group for each function's descriptor in order
-    // for the linker to deduplicate. The duplication, must be from different
-    // tranlation unit, can come from:
-    //  1. Inline functions defined in header files;
-    //  2. ThinLTO imported funcions;
-    //  3. Weak-linkage definitions.
-    // Use a concatenation of the section name and the function name as the
-    // group name so that descriptor-only groups won't be folded with groups of
-    // code.
-    if (Ctx->getTargetTriple().supportsCOMDAT() && !FuncName.empty()) {
+  // Create a separate comdat group for each function's descriptor in order
+  // for the linker to deduplicate. The duplication, must be from different
+  // tranlation unit, can come from:
+  //  1. Inline functions defined in header files;
+  //  2. ThinLTO imported funcions;
+  //  3. Weak-linkage definitions.
+  // Use a concatenation of the section name and the function name as the
+  // group name so that descriptor-only groups won't be folded with groups of
+  // code.
+  if (Ctx->getTargetTriple().supportsCOMDAT() && !FuncName.empty()) {
+    auto ObjFileType = Ctx->getObjectFileType();
+    if (ObjFileType == MCContext::IsELF) {
       auto *S = static_cast<MCSectionELF *>(PseudoProbeDescSection);
       auto Flags = S->getFlags() | ELF::SHF_GROUP;
       return Ctx->getELFSection(S->getName(), S->getType(), Flags,
                                 S->getEntrySize(),
                                 S->getName() + "_" + FuncName,
                                 /*IsComdat=*/true);
+    } else if (ObjFileType == MCContext::IsCOFF) {
+      unsigned Characteristics = COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
+                                 COFF::IMAGE_SCN_MEM_DISCARDABLE |
+                                 COFF::IMAGE_SCN_MEM_READ |
+                                 COFF::IMAGE_SCN_LNK_COMDAT;
+      auto *S = cast<MCSectionCOFF>(PseudoProbeDescSection);
+      std::string COMDATSymName = (S->getName() + "_" + FuncName).str();
+      return Ctx->getCOFFSection(S->getName(), Characteristics, COMDATSymName,
+                                 COFF::IMAGE_COMDAT_SELECT_ANY);
     }
   }
   return PseudoProbeDescSection;

llvm/lib/Target/TargetLoweringObjectFile.cpp

@@ -192,6 +192,33 @@ void TargetLoweringObjectFile::emitCGProfileMetadata(MCStreamer &Streamer,
   }
 }
 
+void TargetLoweringObjectFile::emitPseudoProbeDescMetadata(MCStreamer &Streamer,
+                                                           Module &M) const {
+  if (NamedMDNode *FuncInfo = M.getNamedMetadata(PseudoProbeDescMetadataName)) {
+    // Emit a descriptor for every function including functions that have an
+    // available external linkage. We may not want this for imported functions
+    // that has code in another thinLTO module but we don't have a good way to
+    // tell them apart from inline functions defined in header files. Therefore
+    // we put each descriptor in a separate comdat section and rely on the
+    // linker to deduplicate.
+    auto &C = getContext();
+    for (const auto *Operand : FuncInfo->operands()) {
+      const auto *MD = cast<MDNode>(Operand);
+      auto *GUID = mdconst::dyn_extract<ConstantInt>(MD->getOperand(0));
+      auto *Hash = mdconst::dyn_extract<ConstantInt>(MD->getOperand(1));
+      auto *Name = cast<MDString>(MD->getOperand(2));
+      auto *S = C.getObjectFileInfo()->getPseudoProbeDescSection(
+          TM->getFunctionSections() ? Name->getString() : StringRef());
+
+      Streamer.switchSection(S);
+      Streamer.emitInt64(GUID->getZExtValue());
+      Streamer.emitInt64(Hash->getZExtValue());
+      Streamer.emitULEB128IntValue(Name->getString().size());
+      Streamer.emitBytes(Name->getString());
+    }
+  }
+}
+
 /// getKindForGlobal - This is a top-level target-independent classifier for
 /// a global object.  Given a global variable and information from the TM, this
 /// function classifies the global in a target independent manner. This function

llvm/test/Transforms/SampleProfile/pseudo-probe-callee-profile-mismatch.ll

@@ -1,4 +1,4 @@
-; REQUIRES: x86_64-linux
+; REQUIRES: target={{x86_64-.*-(linux|windows).*}}
 ; REQUIRES: asserts
 ; RUN: opt < %s -passes='thinlto<O2>' -pgo-kind=pgo-sample-use-pipeline  -sample-profile-file=%S/Inputs/pseudo-probe-callee-profile-mismatch.prof --salvage-stale-profile -S --debug-only=sample-profile,sample-profile-matcher,sample-profile-impl  -pass-remarks=inline 2>&1 | FileCheck %s
 
@@ -14,7 +14,7 @@
 
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
-target triple = "x86_64-unknown-linux-gnu"
+target triple = "x86_64-unknown--"
 
 define available_externally i32 @main() #0 {
   %1 = call i32 @bar(), !dbg !13

llvm/test/Transforms/SampleProfile/pseudo-probe-dangle.ll

@@ -1,4 +1,4 @@
-; REQUIRES: x86_64-linux
+; REQUIRES: target={{x86_64-.*-(linux|windows).*}}
 ; RUN: opt < %s -passes='pseudo-probe,jump-threading' -S -o %t
 ; RUN: FileCheck %s < %t --check-prefix=JT
 ; RUN: llc  -function-sections <%t -filetype=asm | FileCheck %s --check-prefix=ASM