[MC] Emit a jump table size section

clang/lib/Driver/ToolChains/PS4CPU.cpp

@@ -265,6 +265,8 @@ void tools::PS5cpu::Linker::ConstructJob(Compilation &C, const JobAction &JA,
     CmdArgs.push_back(D.getLTOMode() == LTOK_Thin ? "--lto=thin"
                                                   : "--lto=full");
 
+  AddLTOFlag("-emit-jump-table-sizes-section");
+
   if (UseJMC)
     AddLTOFlag("-enable-jmc-instrument");
 
@@ -483,6 +485,12 @@ void toolchains::PS4PS5Base::addClangTargetOptions(
     else
       CC1Args.push_back("-fvisibility-externs-nodllstorageclass=keep");
   }
+
+  // Enable jump table sizes section for PS5.
+  if (getTriple().isPS5()) {
+    CC1Args.push_back("-mllvm");
+    CC1Args.push_back("-emit-jump-table-sizes-section");
+  }
 }
 
 // PS4 toolchain.

clang/test/Driver/ps4-ps5-toolchain.c

@@ -11,3 +11,8 @@
 // RUN: %clang %s -### -target x86_64-sie-ps5 -flto 2>&1 | FileCheck %s --check-prefix=LTO
 // LTO-NOT: error:
 // LTO-NOT: unable to pass LLVM bit-code
+
+// Verify that the jump table sizes section is enabled.
+// RUN: %clang %s -target x86_64-sie-ps5 -### 2>&1 | FileCheck -check-prefix=JUMPTABLESIZES %s
+// JUMPTABLESIZES: "-mllvm" "-emit-jump-table-sizes-section"
+// JUMPTABLESIZES: "-plugin-opt=-emit-jump-table-sizes-section"

llvm/docs/Extensions.rst

@@ -554,6 +554,12 @@ time. This section is generated when the compiler enables fat LTO. This section
 has the ``SHF_EXCLUDE`` flag so that it is stripped from the final executable
 or shared library.
 
+``SHT_LLVM_JT_SIZES`` Section (Jump table addresses and sizes)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+This section stores pairs of (jump table address, number of entries).
+This information is useful for tools that need to statically reconstruct
+the control flow of executables.
+
 CodeView-Dependent
 ------------------
 

llvm/include/llvm/BinaryFormat/ELF.h

@@ -1121,6 +1121,7 @@ enum : unsigned {
   SHT_LLVM_BB_ADDR_MAP = 0x6fff4c0a,        // LLVM Basic Block Address Map.
   SHT_LLVM_OFFLOADING = 0x6fff4c0b,         // LLVM device offloading data.
   SHT_LLVM_LTO = 0x6fff4c0c,                // .llvm.lto for fat LTO.
+  SHT_LLVM_JT_SIZES = 0x6fff4c0d,           // LLVM jump tables sizes.
   // Android's experimental support for SHT_RELR sections.
   // https://android.googlesource.com/platform/bionic/+/b7feec74547f84559a1467aca02708ff61346d2a/libc/include/elf.h#512
   SHT_ANDROID_RELR = 0x6fffff00,   // Relocation entries; only offsets.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -107,6 +107,7 @@
 #include "llvm/Pass.h"
 #include "llvm/Remarks/RemarkStreamer.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
@@ -155,6 +156,11 @@ static cl::bits<PGOMapFeaturesEnum> PgoAnalysisMapFeatures(
         "Enable extended information within the SHT_LLVM_BB_ADDR_MAP that is "
         "extracted from PGO related analysis."));
 
+static cl::opt<bool> EmitJumpTableSizesSection(
+    "emit-jump-table-sizes-section",
+    cl::desc("Emit a section containing jump table addresses and sizes"),
+    cl::Hidden, cl::init(false));
+
 STATISTIC(EmittedInsts, "Number of machine instrs printed");
 
 char AsmPrinter::ID = 0;
@@ -2764,6 +2770,27 @@ void AsmPrinter::emitJumpTableInfo() {
     for (const MachineBasicBlock *MBB : JTBBs)
       emitJumpTableEntry(MJTI, MBB, JTI);
   }
+
+  if (EmitJumpTableSizesSection && TM.getTargetTriple().isOSBinFormatELF() &&
+      !JT.empty()) {
+    MCSymbolELF *LinkedToSym = cast<MCSymbolELF>(CurrentFnSym);
+    int Flags = F.hasComdat() ? ELF::SHF_GROUP : 0;
+    StringRef GroupName = F.hasComdat() ? F.getComdat()->getName() : "";
+
+    MCSection *JumpTableSizesSection = OutContext.getELFSection(
+        ".llvm_jump_table_sizes", ELF::SHT_LLVM_JT_SIZES, Flags, 0, GroupName,
+        F.hasComdat(), MCSection::NonUniqueID, LinkedToSym);
+
+    OutStreamer->switchSection(JumpTableSizesSection);
+
+    for (unsigned JTI = 0, E = JT.size(); JTI != E; ++JTI) {
+      const std::vector<MachineBasicBlock *> &JTBBs = JT[JTI].MBBs;
+      OutStreamer->emitSymbolValue(GetJTISymbol(JTI),
+                                   TM.getProgramPointerSize());
+      OutStreamer->emitIntValue(JTBBs.size(), TM.getProgramPointerSize());
+    }
+  }
+
   if (!JTInDiffSection)
     OutStreamer->emitDataRegion(MCDR_DataRegionEnd);
 }

llvm/lib/MC/MCParser/ELFAsmParser.cpp

@@ -677,6 +677,8 @@ bool ELFAsmParser::ParseSectionArguments(bool IsPush, SMLoc loc) {
       Type = ELF::SHT_LLVM_OFFLOADING;
     else if (TypeName == "llvm_lto")
       Type = ELF::SHT_LLVM_LTO;
+    else if (TypeName == "llvm_jt_sizes")
+      Type = ELF::SHT_LLVM_JT_SIZES;
     else if (TypeName.getAsInteger(0, Type))
       return TokError("unknown section type");
   }

llvm/lib/MC/MCSectionELF.cpp

@@ -172,6 +172,8 @@ void MCSectionELF::printSwitchToSection(const MCAsmInfo &MAI, const Triple &T,
     OS << "llvm_offloading";
   else if (Type == ELF::SHT_LLVM_LTO)
     OS << "llvm_lto";
+  else if (Type == ELF::SHT_LLVM_JT_SIZES)
+    OS << "llvm_jt_sizes";
   else
     OS << "0x" << Twine::utohexstr(Type);
 

llvm/lib/Object/ELF.cpp

@@ -319,6 +319,7 @@ StringRef llvm::object::getELFSectionTypeName(uint32_t Machine, unsigned Type) {
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_BB_ADDR_MAP);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_OFFLOADING);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_LTO);
+    STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_JT_SIZES)
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_ATTRIBUTES);
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_HASH);
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_verdef);

llvm/test/CodeGen/X86/jump-table-size-section.ll

@@ -0,0 +1,97 @@
+; RUN: llc %s -o - -emit-jump-table-sizes-section -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=CHECK %s
+; RUN: llc %s -o - -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOFLAG %s
+
+; This test verifies the jump table size section. Currently only enabled by default on the PS5 target.
+
+$foo1 = comdat any
+
+; Ensure proper comdat handling.
+define void @foo1(i32 %x, ptr %to) comdat {
+
+; CHECK-LABEL: foo1
+; CHECK:      .section        .llvm_jump_table_sizes,"G",@llvm_jt_sizes,foo1,comdat
+; CHECK-NEXT: .quad   .LJTI0_0
+; CHECK-NEXT: .quad   6
+
+; NOFLAG-LABEL: foo1
+; NOFLAG-NOT: .section        .llvm_jump_table_sizes
+
+entry:
+  switch i32 %x, label %default [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+    i32 4, label %bb4
+    i32 5, label %bb4
+  ]
+bb0:
+  store i32 0, ptr %to
+  br label %exit
+bb1:
+  store i32 1, ptr %to
+  br label %exit
+bb2:
+  store i32 2, ptr %to
+  br label %exit
+bb3:
+  store i32 3, ptr %to
+  br label %exit
+bb4:
+  store i32 4, ptr %to
+  br label %exit
+exit:
+  ret void
+default:
+  unreachable
+}
+
+define void @foo2(i32 %x, ptr %to) {
+
+; CHECK-LABEL: foo2
+; CHECK:      .section        .llvm_jump_table_sizes
+; CHECK-NEXT: .quad   .LJTI1_0
+; CHECK-NEXT: .quad   5
+
+; NOFLAG-LABEL: foo2
+; NOFLAG-NOT: .section        .llvm_jump_table_sizes
+
+entry:
+  switch i32 %x, label %default [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+    i32 4, label %bb4
+  ]
+bb0:
+  store i32 0, ptr %to
+  br label %exit
+bb1:
+  store i32 1, ptr %to
+  br label %exit
+bb2:
+  store i32 2, ptr %to
+  br label %exit
+bb3:
+  store i32 3, ptr %to
+  br label %exit
+bb4:
+  store i32 4, ptr %to
+  br label %exit
+exit:
+  ret void
+default:
+  unreachable
+}
+
+; Ensure that the section isn't produced if there is no jump table.
+
+define void @foo3(i32 %x, ptr %to) {
+
+; CHECK-LABEL:    foo3
+; CHECK-NOT:      .section        .llvm_jump_table_sizes
+
+exit:
+  ret void
+}