[AArch64] Don't try to compress jump tables if there are any inline asm instructions.

[aemerson]

Summary

Inline asm can contain constructs like .bytes which may have arbitrary size. In some cases, this causes us to miscalculate the size of blocks and therefore offsets, causing us to incorrectly compress a JT.

To be safe, just bail out of the whole thing if we find any inline asm.

Fixes PR48255

[t.p.northover]

I think this is a reasonable change, inline asm is pretty rare.

llvm/lib/Target/AArch64/AArch64CompressJumpTables.cpp

 STATISTIC(NumJT32, "Number of jump-tables with 4-byte entries");
 
 namespace {
 class AArch64CompressJumpTables : public MachineFunctionPass {
   const TargetInstrInfo *TII;
   MachineFunction *MF;
   SmallVector<int, 8> BlockInfo;
 
-  int computeBlockSize(MachineBasicBlock &MBB);
-  void scanFunction();
+  /// Returns the size in instructions of the block \p MBB, or None if we
+  /// couldn't get a safe upper bound.
+  Optional<int> computeBlockSize(MachineBasicBlock &MBB);
+
+  /// Gather information about the function, returns false if we can't perform
+  /// this optimization for some reason.
+  bool scanFunction();
 
   bool compressJumpTable(MachineInstr &MI, int Offset);
 
 public:
   static char ID;
   AArch64CompressJumpTables() : MachineFunctionPass(ID) {
     initializeAArch64CompressJumpTablesPass(*PassRegistry::getPassRegistry());
   }
public:
   }
 };
 char AArch64CompressJumpTables::ID = 0;
 } // namespace
 
 INITIALIZE_PASS(AArch64CompressJumpTables, DEBUG_TYPE,
                 "AArch64 compress jump tables pass", false, false)
 
-int AArch64CompressJumpTables::computeBlockSize(MachineBasicBlock &MBB) {
+Optional<int>
+AArch64CompressJumpTables::computeBlockSize(MachineBasicBlock &MBB) {
   int Size = 0;
-  for (const MachineInstr &MI : MBB)
+  for (const MachineInstr &MI : MBB) {
+    // Inline asm may contain some directives like .bytes which we don't
+    // currently have the ability to parse accurately. To be safe, just avoid
+    // computing a size and bail out.
+    if (MI.getOpcode() == AArch64::INLINEASM ||
+        MI.getOpcode() == AArch64::INLINEASM_BR)
+      return None;
     Size += TII->getInstSizeInBytes(MI);
+  }
   return Size;
 }
 
-void AArch64CompressJumpTables::scanFunction() {
+bool AArch64CompressJumpTables::scanFunction() {
   BlockInfo.clear();
   BlockInfo.resize(MF->getNumBlockIDs());
 
   unsigned Offset = 0;
   for (MachineBasicBlock &MBB : *MF) {
     const Align Alignment = MBB.getAlignment();
     unsigned AlignedOffset;
     if (Alignment == Align(1))
       AlignedOffset = Offset;
     else
       AlignedOffset = alignTo(Offset, Alignment);
     BlockInfo[MBB.getNumber()] = AlignedOffset;
-    Offset = AlignedOffset + computeBlockSize(MBB);
+    auto BlockSize = computeBlockSize(MBB);
+    if (!BlockSize)
+      return false;
+    Offset = AlignedOffset + *BlockSize;
   }
+  return true;
 }
 
 bool AArch64CompressJumpTables::compressJumpTable(MachineInstr &MI,
                                                   int Offset) {
   if (MI.getOpcode() != AArch64::JumpTableDest32)
     return false;
 
   int JTIdx = MI.getOperand(4).getIndex();
bool AArch64CompressJumpTables::runOnMachineFunction(MachineFunction &MFIn) {
   MF = &MFIn;
 
   const auto &ST = MF->getSubtarget<AArch64Subtarget>();
   TII = ST.getInstrInfo();
 
   if (ST.force32BitJumpTables() && !MF->getFunction().hasMinSize())
     return false;
 
-  scanFunction();
+  if (!scanFunction())
+    return false;
 
   for (MachineBasicBlock &MBB : *MF) {
     int Offset = BlockInfo[MBB.getNumber()];
     for (MachineInstr &MI : MBB) {
       Changed |= compressJumpTable(MI, Offset);
       Offset += TII->getInstSizeInBytes(MI);
     }
   }
 
   return Changed;
 }
 
 FunctionPass *llvm::createAArch64CompressJumpTablesPass() {
   return new AArch64CompressJumpTables();
 }

llvm/test/CodeGen/AArch64/jump-table-compress.mir

 # RUN: llc -mtriple=aarch64-linux-gnu %s -run-pass=aarch64-jump-tables -o - | FileCheck %s
 --- |
   define i32 @test_jumptable(i32 %in) {
     unreachable
   }
 
+  define void @test_inline_asm_no_compress() { ret void }
+
 ...
 ---
 name:            test_jumptable
 alignment:       4
 exposesReturnsTwice: false
 legalized:       false
 regBankSelected: false
 selected:        false
bb.11:
     dead $xzr = SPACE 1048580, undef $xzr
 
   bb.12:
     ; CHECK-LABEL: bb.12
     ; CHECK: JumpTableDest32
     early-clobber $x10, dead early-clobber $x11 = JumpTableDest32 undef killed $x9, undef killed $x8, %jump-table.5
     BR killed $x10
 ...
+---
+name:            test_inline_asm_no_compress
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$w0' }
+  - { reg: '$w1' }
+  - { reg: '$w2' }
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+jumpTable:
+  kind:            label-difference32
+  entries:
+    - id:              0
+      blocks:          [ '%bb.2', '%bb.4', '%bb.5', '%bb.6', '%bb.7', '%bb.8' ]
+body:             |
+  bb.0:
+    successors: %bb.3(0x12492492), %bb.1(0x6db6db6e)
+    liveins: $w0, $w1, $w2
+  
+    dead $wzr = SUBSWri renamable $w0, 5, 0, implicit-def $nzcv
+    Bcc 8, %bb.3, implicit $nzcv
+  
+  bb.1:
+    successors: %bb.2, %bb.4, %bb.5, %bb.6, %bb.7, %bb.8
+    liveins: $w0, $w1, $w2
+    ; We check that if there's an inline asm instruction in the jump table,
+    ; that we skip compression. This is due to not being able to rely on
+    ; finding the instruction size in some cases.
+    ; CHECK-LABEL: test_inline_asm_no_compress
+    ; CHECK-LABEL: bb.1
+    ; CHECK: JumpTableDest32
+    renamable $w8 = ORRWrs $wzr, killed renamable $w0, 0, implicit-def $x8
+    $x9 = ADRP target-flags(aarch64-page) %jump-table.0
+    renamable $x9 = ADDXri $x9, target-flags(aarch64-pageoff, aarch64-nc) %jump-table.0, 0
+    early-clobber renamable $x10, dead early-clobber renamable $x11 = JumpTableDest32 killed renamable $x9, killed renamable $x8, %jump-table.0
+    BR killed renamable $x10
+  
+  bb.2:
+    liveins: $w1, $w2
+  
+    INLINEASM &".byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09", 1 /* sideeffect attdialect */
+    INLINEASM &".byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09.byte 0x1f,0x20,0x03,0xd5,0x1f,0x20,0x03,0xd5\0A\09", 1 /* sideeffect attdialect */
+    $w0 = ADDWrs killed renamable $w2, killed renamable $w1, 0
+    RET undef $lr, implicit $w0
+  
+  bb.3:
+    $w0 = MOVZWi 0, 0
+    RET undef $lr, implicit $w0
+  
+  bb.4:
+    liveins: $w1, $w2
+  
+    renamable $w0 = nsw MADDWrrr killed renamable $w2, killed renamable $w1, $wzr
+    RET undef $lr, implicit $w0
+  
+  bb.5:
+    liveins: $w1, $w2
+  
+    $w0 = SUBWrs killed renamable $w1, killed renamable $w2, 0
+    RET undef $lr, implicit $w0
+  
+  bb.6:
+    liveins: $w1, $w2
+  
+    $w0 = SUBWrs killed renamable $w2, killed renamable $w1, 0
+    RET undef $lr, implicit $w0
+  
+  bb.7:
+    liveins: $w1, $w2
+  
+    renamable $w0 = MADDWrrr killed renamable $w1, renamable $w1, killed renamable $w2
+    RET undef $lr, implicit $w0
+  
+  bb.8:
+    liveins: $w1, $w2
+  
+    renamable $w8 = nsw MADDWrrr renamable $w2, renamable $w2, $wzr
+    renamable $w0 = MADDWrrr killed renamable $w8, killed renamable $w2, killed renamable $w1
+    RET undef $lr, implicit $w0
+
+...

llvm/test/CodeGen/AArch64/jump-table.ll

 
 def:
   ret i32 0
 
 lbl1:
   ret i32 1
 
 lbl2:
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
-  call void asm sideeffect "1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16", ""()
+  call i64 @llvm.aarch64.space(i32 1024, i64 undef)
   ret i32 2
 
 lbl3:
   ret i32 4
 
 lbl4:
   ret i32 8