[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/include/llvm/CodeGen/AsmPrinter.h

@@ -894,6 +894,10 @@ class AsmPrinter : public MachineFunctionPass {
 
   void emitJumpTableEntry(const MachineJumpTableInfo *MJTI,
                           const MachineBasicBlock *MBB, unsigned uid) const;
+
+  void emitJumpTableSizesSection(const MachineJumpTableInfo *MJTI,
+                                 const Function &F) const;
+
   void emitLLVMUsedList(const ConstantArray *InitList);
   /// Emit llvm.ident metadata in an '.ident' directive.
   void emitModuleIdents(Module &M);

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -99,6 +99,7 @@
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/MCSymbol.h"
+#include "llvm/MC/MCSymbolCOFF.h"
 #include "llvm/MC/MCSymbolELF.h"
 #include "llvm/MC/MCTargetOptions.h"
 #include "llvm/MC/MCValue.h"
@@ -107,6 +108,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 +157,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,10 +2771,62 @@ void AsmPrinter::emitJumpTableInfo() {
     for (const MachineBasicBlock *MBB : JTBBs)
       emitJumpTableEntry(MJTI, MBB, JTI);
   }
+
+  if (EmitJumpTableSizesSection)
+    emitJumpTableSizesSection(MJTI, F);
+
   if (!JTInDiffSection)
     OutStreamer->emitDataRegion(MCDR_DataRegionEnd);
 }
 
+void AsmPrinter::emitJumpTableSizesSection(const MachineJumpTableInfo *MJTI,
+                                           const Function &F) const {
+  const std::vector<MachineJumpTableEntry> &JT = MJTI->getJumpTables();
+
+  if (JT.empty())
+    return;
+
+  StringRef GroupName = F.hasComdat() ? F.getComdat()->getName() : "";
+  MCSection *JumpTableSizesSection = nullptr;
+  StringRef sectionName = ".llvm_jump_table_sizes";
+
+  bool isElf = TM.getTargetTriple().isOSBinFormatELF();
+  bool isCoff = TM.getTargetTriple().isOSBinFormatCOFF();
+
+  if (!isCoff && !isElf)
+    return;
+
+  if (isElf) {
+    MCSymbolELF *LinkedToSym = dyn_cast<MCSymbolELF>(CurrentFnSym);
+    int Flags = F.hasComdat() ? ELF::SHF_GROUP : 0;
+
+    JumpTableSizesSection = OutContext.getELFSection(
+        sectionName, ELF::SHT_LLVM_JT_SIZES, Flags, 0, GroupName, F.hasComdat(),
+        MCSection::NonUniqueID, LinkedToSym);
+  } else if (isCoff) {
+    if (F.hasComdat()) {
+      JumpTableSizesSection = OutContext.getCOFFSection(
+          sectionName,
+          COFF::IMAGE_SCN_CNT_INITIALIZED_DATA | COFF::IMAGE_SCN_MEM_READ |
+              COFF::IMAGE_SCN_LNK_COMDAT | COFF::IMAGE_SCN_MEM_DISCARDABLE,
+          F.getComdat()->getName(), COFF::IMAGE_COMDAT_SELECT_ASSOCIATIVE);
+    } else {
+      JumpTableSizesSection = OutContext.getCOFFSection(
+          sectionName, COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
+                           COFF::IMAGE_SCN_MEM_READ |
+                           COFF::IMAGE_SCN_MEM_DISCARDABLE);
+    }
+  }
+
+  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());
+  }
+}
+
 /// EmitJumpTableEntry - Emit a jump table entry for the specified MBB to the
 /// current stream.
 void AsmPrinter::emitJumpTableEntry(const MachineJumpTableInfo *MJTI,

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,210 @@
+; RUN: llc %s -o - -mtriple x86_64-sie-ps5 -emit-jump-table-sizes-section -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=PS5-CHECK %s
+; RUN: llc %s -o - -mtriple x86_64-sie-ps5 -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOFLAG %s
+; RUN: llc %s -o - -mtriple x86_64-sie-ps5 -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOTABLE %s
+
+; RUN: llc %s -o - -mtriple x86_64-unknown-linux-gnu -emit-jump-table-sizes-section -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=ELF-CHECK %s
+; RUN: llc %s -o - -mtriple x86_64-unknown-linux-gnu -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOFLAG %s
+; RUN: llc %s -o - -mtriple x86_64-unknown-linux-gnu -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOTABLE %s
+
+; RUN: llc %s -o - -mtriple x86_64-pc-windows-msvc -emit-jump-table-sizes-section -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=COFF-CHECK %s
+; RUN: llc %s -o - -mtriple x86_64-pc-windows-msvc -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOFLAG %s
+; RUN: llc %s -o - -mtriple x86_64-pc-windows-msvc -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOTABLE %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 {
+
+; PS5-CHECK-LABEL: foo1
+; PS5-CHECK:       .section        .llvm_jump_table_sizes,"G",@llvm_jt_sizes,foo1,comdat
+; PS5-CHECK-NEXT: .quad   .LJTI0_0
+; PS5-CHECK-NEXT: .quad   6
+
+; ELF-CHECK-LABEL: foo1
+; ELF-CHECK:       .section        .llvm_jump_table_sizes,"G",@llvm_jt_sizes,foo1,comdat
+; ELF-CHECK-NEXT: .quad   .LJTI0_0
+; ELF-CHECK-NEXT: .quad   6
+
+; COFF-CHECK-LABEL: foo1
+; COFF-CHECK:      .section         .llvm_jump_table_sizes,"drD",associative,foo1
+; COFF-CHECK-NEXT: .quad   .LJTI0_0
+; COFF-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) {
+
+; PS5-CHECK-LABEL:    foo2
+; PS5-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; PS5-CHECK-NEXT:     .quad .LJTI1_0
+; PS5-CHECK-NEXT:     .quad   5
+
+; ELF-CHECK-LABEL:    foo2
+; ELF-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; ELF-CHECK-NEXT:     .quad .LJTI1_0
+; ELF-CHECK-NEXT:     .quad   5
+
+; COFF-CHECK-LABEL:   foo2
+; COFF-CHECK:         .section         .llvm_jump_table_sizes,"drD"
+; COFF-CHECK-NEXT:    .quad .LJTI1_0
+; COFF-CHECK-NEXT:    .quad   5
+
+; 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
+  ]
+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) {
+
+; NOTABLE-LABEL:    foo3
+; NOTABLE-NOT:      .section        .llvm_jump_table_sizes
+
+exit:
+  ret void
+}
+
+; Ensure we can deal with nested jump tables.
+
+define void @nested(i32 %x, i32 %y, ptr %to) {
+
+; PS5-CHECK-LABEL:    nested
+; PS5-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; PS5-CHECK-NEXT:     .quad .LJTI3_0
+; PS5-CHECK-NEXT:     .quad   5
+; PS5-CHECK-NEXT:     .quad .LJTI3_1
+; PS5-CHECK-NEXT:     .quad 6
+
+; ELF-CHECK-LABEL:    nested
+; ELF-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; ELF-CHECK-NEXT:     .quad .LJTI3_0
+; ELF-CHECK-NEXT:     .quad   5
+; ELF-CHECK-NEXT:     .quad .LJTI3_1
+; ELF-CHECK-NEXT:     .quad 6
+
+; COFF-CHECK-LABEL:   nested
+; COFF-CHECK:         .section         .llvm_jump_table_sizes,"drD"
+; COFF-CHECK-NEXT:     .quad .LJTI3_0
+; COFF-CHECK-NEXT:     .quad   5
+; COFF-CHECK-NEXT:     .quad .LJTI3_1
+; COFF-CHECK-NEXT:     .quad 6
+
+; NOFLAG-LABEL:       nested
+; 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:
+  switch i32 %y, label %default [
+    i32 1, label %bb5
+    i32 2, label %bb6
+    i32 3, label %bb7
+    i32 4, label %bb8
+    i32 5, label %bb9
+    i32 6, label %bb10
+  ]
+  br label %exit2
+bb5:
+  store i32 4, ptr %to
+  br label %exit
+bb6:
+  store i32 4, ptr %to
+  br label %exit
+bb7:
+  store i32 4, ptr %to
+  br label %exit
+bb8:
+  store i32 4, ptr %to
+  br label %exit
+bb9:
+  store i32 4, ptr %to
+  br label %exit
+bb10:
+  store i32 4, ptr %to
+  br label %exit
+exit:
+  ret void
+exit2:
+  ret void
+default:
+  unreachable
+}