Add an option to allow JumpInstrTables to set unnamed_addr and jumptable on all address-taken functions

[tmroeder]

Summary

This small patch adds a new command flag -jump-table-all and an associated TargetOptions field JumpTableAll, and it modifies JumpInstrTables to add unnamed_addr and jumptable to all address-taken functions if the flag is set.

This is important for applications like my implementation of control-flow integrity, which want to operate over all address-taken functions and not just the ones that happened to have been marked jumptable.

include/llvm/CodeGen/CommandFlags.h

                          "Create one table per unique function type."),
               clEnumValEnd));
 
+cl::opt<bool>
+JumpTableAll("jump-table-all",
+             cl::desc("Add unnamed_addr and jumptable to all address-taken"
+                      " functions"),
+             cl::init(false));
+
 // Common utility function tightly tied to the options listed here. Initializes
 // a TargetOptions object with CodeGen flags and returns it.
 static inline TargetOptions InitTargetOptionsFromCodeGenFlags() {

 
   Options.MCOptions = InitMCTargetOptionsFromFlags();
   Options.JTType = JTableType;
+  Options.JumpTableAll = JumpTableAll;
 
   return Options;
 }

include/llvm/CodeGen/JumpInstrTables.h

   static char ID;
 
   JumpInstrTables();
-  JumpInstrTables(JumpTable::JumpTableType JTT);
+  JumpInstrTables(JumpTable::JumpTableType JTT, bool JTAll);
   virtual ~JumpInstrTables();
   bool runOnModule(Module &M) override;
   const char *getPassName() const override { return "Jump-Instruction Tables"; }

 
   /// The type of tables to build.
   JumpTable::JumpTableType JTType;
+
+  /// Whether or not to apply unnamed_addr and jumptable to all address-taken
+  /// functions.
+  bool JumpTableAll;
 };
 
 /// Creates a JumpInstrTables pass for the given type of jump table.
-ModulePass *createJumpInstrTablesPass(JumpTable::JumpTableType JTT);
+ModulePass *createJumpInstrTablesPass(JumpTable::JumpTableType JTT, bool JTAll);
 }
 
 #endif /* LLVM_CODEGEN_JUMPINSTRTABLES_H */

include/llvm/Target/TargetOptions.h

           CompressDebugSections(false), FunctionSections(false),
           DataSections(false), TrapUnreachable(false), TrapFuncName(""),
           FloatABIType(FloatABI::Default),
-          AllowFPOpFusion(FPOpFusion::Standard), JTType(JumpTable::Single) {}
+          AllowFPOpFusion(FPOpFusion::Standard), JTType(JumpTable::Single),
+          JumpTableAll(false) {}
 
     /// PrintMachineCode - This flag is enabled when the -print-machineinstrs
     /// option is specified on the command line, and should enable debugging

     /// create for functions that have the jumptable attribute.
     JumpTable::JumpTableType JTType;
 
+    /// JumpTableAll - This flags controls whether or not jumptable is
+    /// automatically applied to all address-taken functions.
+    bool JumpTableAll;
+
     /// Machine level options.
     MCTargetOptions MCOptions;
   };

lib/CodeGen/JumpInstrTables.cpp

 STATISTIC(NumFuncsInJumpTables, "Number of functions in the jump tables");
 
 ModulePass *llvm::createJumpInstrTablesPass() {
-  // The default implementation uses a single table for all functions.
-  return new JumpInstrTables(JumpTable::Single);
+  // The default implementation uses a single table for all functions, and it
+  // doesn't rewrite all address-taken functions to have unnamed_addr and
+  // jumptable.
+  return new JumpInstrTables(JumpTable::Single, false);
 }
 
-ModulePass *llvm::createJumpInstrTablesPass(JumpTable::JumpTableType JTT) {
-  return new JumpInstrTables(JTT);
+ModulePass *llvm::createJumpInstrTablesPass(JumpTable::JumpTableType JTT,
+                                            bool JTAll) {
+  return new JumpInstrTables(JTT, JTAll);
 }
 
 namespace {

 
 JumpInstrTables::JumpInstrTables()
     : ModulePass(ID), Metadata(), JITI(nullptr), TableCount(0),
-      JTType(JumpTable::Single) {
+      JTType(JumpTable::Single), JumpTableAll(false) {
   initializeJumpInstrTablesPass(*PassRegistry::getPassRegistry());
 }
 
-JumpInstrTables::JumpInstrTables(JumpTable::JumpTableType JTT)
-    : ModulePass(ID), Metadata(), JITI(nullptr), TableCount(0), JTType(JTT) {
+JumpInstrTables::JumpInstrTables(JumpTable::JumpTableType JTT, bool JTAll)
+    : ModulePass(ID), Metadata(), JITI(nullptr), TableCount(0), JTType(JTT),
+    JumpTableAll(JTAll) {
   initializeJumpInstrTablesPass(*PassRegistry::getPassRegistry());
 }
 

       assert(F.hasUnnamedAddr() &&
              "Attribute 'jumptable' requires 'unnamed_addr'");
       Functions[&F] = nullptr;
+    } else if (JumpTableAll && F.hasAddressTaken(nullptr)) {
+      F.setUnnamedAddr(true);
+      F.addFnAttr(Attribute::JumpTable);
+      Functions[&F] = nullptr;
     }
   }
 

lib/CodeGen/LLVMTargetMachine.cpp

   // Passes to handle jumptable function annotations. These can't be handled at
   // JIT time, so we don't add them directly to addPassesToGenerateCode.
   PM.add(createJumpInstrTableInfoPass());
-  PM.add(createJumpInstrTablesPass(Options.JTType));
+  PM.add(createJumpInstrTablesPass(Options.JTType, Options.JumpTableAll));
 
   // Add common CodeGen passes.
   MCContext *Context = addPassesToGenerateCode(this, PM, DisableVerify,

test/CodeGen/X86/jump_table_all.ll

+; RUN: llc <%s -jump-table-all -jump-table-type=single | FileCheck %s
+; RUN: llc <%s -jump-table-type=single | FileCheck --check-prefix=WITHOUT %s
+target triple = "x86_64-unknown-linux-gnu"
+define i32 @f() {
+entry:
+  ret i32 0
+}
+
+define i32 @g() {
+entry:
+  ret i32 0
+}
+
+define i32 @main(i32 %argc, i8** %argv) {
+  %temp = alloca i32 ()*, align 8
+  store i32 ()* @f, i32()** %temp, align 8
+; CHECK: movq    $__llvm_jump_instr_table_0_1
+; WITHOUT-NOT: movq    $__llvm_jump_instr_table_0_1
+  %1 = load i32 ()** %temp, align 8
+; CHECK: movl    $__llvm_jump_instr_table_0_1
+  %2 = call i32 ()* %1()
+  %3 = call i32 ()* @g()
+  ret i32 %3
+}
+
+; This table should be created even without the jumptable annotation
+; because llc is called with -jump-table-all, and @f has its address taken
+; CHECK:         .globl  __llvm_jump_instr_table_0_1
+; CHECK:         .align  8, 0x90
+; CHECK:         .type   __llvm_jump_instr_table_0_1,@function
+; CHECK: __llvm_jump_instr_table_0_1:
+; CHECK:         jmp     f@PLT
+; CHECK-NOT:     jmp     g@PLT
+
+; WITHOUT-NOT:          .globl  __llvm_jump_instr_table_0_1