[AArch64CompressJumpTables] Prevent over-compression caused by invalid alignment.

[paulwalker-arm]

Summary

AArch64CompressJumpTables assumes it can calculate exact block
offsets. This assumption is bogus because getInstSizeInBytes()
only returns an upper bound rather than an exact size. The
assumption is also invalid when a block alignment is bigger than
the function's alignment.

To mitigate both scenarios this patch changes the algorithm to
compute the maximum upper bound for all block offsets. This is
pessimistic but safe because all offsets are treated as unsigned.

[paulwalker-arm]

Not sure of the performance ramifications of this patch but this is a bug fix for something I've hit in real world code. Looking at AArch64Subtarget.cpp has the affected targets as:

CortexA76
CortexA77
CortexA78
CortexA78C
CortexR82
CortexX1
CortexX1C
CortexA710
CortexA715
CortexX2
CortexX3
NeoverseN1
NeoverseN2
NeoverseV2
NeoverseV1

I'm not sure where the value is for having function alignments smaller than loop alignments, so matching them up is a potential easy fix to restore jump table compression. Even if that's done I believe this fix is required to catch other places where alignments can be set.

[efriedma]

Don't you just need to add "Alignment - FunctionAlignment" bytes to AlignedOffset? The assembler resolves the final offsets, so we don't need precise offsets at the point we're doing the compression; we just need a conservative estimate so we don't over-compress.

[paulwalker-arm]

In D150009#4323571, @efriedma wrote:

Don't you just need to add "Alignment - FunctionAlignment" bytes to AlignedOffset? The assembler resolves the final offsets, so we don't need precise offsets at the point we're doing the compression; we just need a conservative estimate so we don't over-compress.

Is it that simple? Looking at the logic I convinced myself that overestimating the offsets will be just as bad because it can make it look like MinBlock is closer to MaxBlock than it actually is. My assumption is that to be safe we need to calculate min and max offset estimates, with min == max representing the case where the offset is known. Do you agree? or am I over complicating things?

If it is more involved do you object to me getting the fix in first and the following up with the rework?

[efriedma]

If we have the maximum size of each block, and the alignment of each block, we can compute a conservative estimate of the layout of the function:

-- block1 offset
block1 instruction size estimate
block2 padding estimate (alignment - 4)
--block2 offset
block2 instruction size estimate
block3 padding estimate (alignment - 4)
--block3 offset
[...]

The actual distance between two of the block offsets will always be less than an estimate computed using simple subtraction: there isn't any actual alignment computation involved, so blocks can't appear closer than they actually are in either direction.

Looking more closely, though, it looks like it doesn't actually work that way at the moment, though: instead of using an alignment estimate, it uses alignTo(). That's probably wrong; I think getInstSizeInBytes is a conservative estimate in some cases.

So the right fix is instead of:

if (Alignment == Align(1))
  AlignedOffset = Offset;
else
  AlignedOffset = alignTo(Offset, Alignment);

We should do:

if (Alignment <= Align(4))
  AlignedOffset = Offset;
else
  AlignedOffset = Offset + Alignment.value() - 4;

That should solve the issue without specifically referencing the function's alignment.

[aemerson]

Adding @t.p.northover

[paulwalker-arm]

Sorry of the delay in following up. I've notice that branch relaxation solves the same problem so I've copied its solution and updated the test accordingly.

[efriedma]

I'm pretty sure what branch relaxation is doing is wrong, in a subtle way.

Suppose, at an extreme, you have an instruction that's somewhere between 0 and 124 bytes. So you have a sequence like this:

bb align 128:
ZERO_TO_124_BYTES_INSTR
B another_bb

bb2 align 128:
[...]

Say the start of the block is aligned. We add up the size of the block: ZERO_TO_124_BYTES_INSTR is 124 bytes, B is 4 bytes, so the block is 128 bytes long. Calling alignTo() does nothing, so we conclude there's no padding after the branch. But actually, there can be up to 124 bytes of padding depending on the actual size of ZERO_TO_124_BYTES_INSTR. And that padding increases the distance from the branch to the destination.

To get this right, we must not call alignTo(): trying to compute padding after a block using alignTo() assumes we know the exact offset of the end of the block, which is not a correct assumption. We need to instead estimate the maximum amount of padding the .p2align directive will insert.

I'm not sure if there are any pseudo-instructions that are actually variable-length on aarch64 at the moment, but that's how the interface works in general; such instructions exist on other targets, and they've existed in the past on AArch64. We also need to consider the interaction with MachineOutliner, which can effectively shrink the size of arbitrary sequences of instructions.

Branch relaxation is pretty forgiving, generally, so I'm not surprised nobody has tripped over this.

[paulwalker-arm]

Updated to always calculate pessimistic offsets based on Eli's education as to what getInstSizeInBytes really means.

[efriedma]

LGTM

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;
 
-  /// Returns the size in instructions of the block \p MBB, or std::nullopt if
+  /// Returns the size of instructions in the block \p MBB, or std::nullopt if
   /// we couldn't get a safe upper bound.
   std::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);
AArch64CompressJumpTables::computeBlockSize(MachineBasicBlock &MBB) {
   }
   return Size;
 }
 
 bool AArch64CompressJumpTables::scanFunction() {
   BlockInfo.clear();
   BlockInfo.resize(MF->getNumBlockIDs());
 
+  // NOTE: BlockSize, Offset, OffsetAfterAlignment are all upper bounds.
+
   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;
+    unsigned OffsetAfterAlignment = Offset;
+    // We don't know the exact size of MBB so assume worse case padding.
+    if (Alignment > Align(4))
+      OffsetAfterAlignment += Alignment.value() - 4;
+    BlockInfo[MBB.getNumber()] = OffsetAfterAlignment;
     auto BlockSize = computeBlockSize(MBB);
     if (!BlockSize)
       return false;
-    Offset = AlignedOffset + *BlockSize;
+    Offset = OffsetAfterAlignment + *BlockSize;
   }
   return true;
 }
 
 bool AArch64CompressJumpTables::compressJumpTable(MachineInstr &MI,
                                                   int Offset) {
   if (MI.getOpcode() != AArch64::JumpTableDest32)
     return false;

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 }
+  define void @test_bb_alignment_not_byte_compressable() { ret void }
 
 ...
 ---
 name:            test_jumptable
 alignment:       4
 exposesReturnsTwice: false
 legalized:       false
 regBankSelected: false
body:             |
   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
 
 ...
+---
+name:            test_bb_alignment_not_byte_compressable
+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
+    ; Ensure there's no jump table compression when block alignments are bigger
+    ; than the function alignment because we don't known the padding length at
+    ; the point where compression is done.
+    ; CHECK-LABEL: test_bb_alignment_not_byte_compressable
+    ; CHECK-LABEL: bb.1
+    ; CHECK: JumpTableDest16
+    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
+    $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 is aligned to make it more that 256 instructions away from bb.1, which
+  ; means we can no longer assume the jump table will be byte indexable.
+  bb.5 (align 1024):
+    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
+
+...