[AArch64] Support stackmap/patchpoint in getInstSizeInBytes

[rovka]

Summary

We currently return 4 for stackmaps and patchpoints, which is very optimistic
and can in rare cases cause the branch relaxation pass to fail to relax certain
branches.

This patch causes getInstSizeInBytes to return a pessimistic estimate of the
size as the number of bytes requested in the stackmap/patchpoint. In the future,
we could provide a more accurate estimate by sharing some of the logic in
AArch64::LowerSTACKMAP/PATCHPOINT.

Fixes part of https://llvm.org/bugs/show_bug.cgi?id=28750

[rengolin]

lib/Target/AArch64/AArch64InstrInfo.cpp

64 ↗

(On Diff #69832)

Nit: maybe make the switch set NumBytes and return it in the end?

The unreachable down there is silly, but probably silences a compiler warning. This way, we could get rid of that, too.

[rovka]

Got rid of the unreachable, as Renato suggested.

[rengolin]

Now that looks good to me.

I don't like much the tests, because if the branch relaxation algorithm change, this test can "fail" without really failing, and whoever hits this will have a hard time understanding how to make this test proper. But I have no better idea.

I'll leave this be for a while, see if anyone else has a better idea. If not, having those tests are better than nothing, and we can work on trying some unit tests in the future.

[rovka]

I agree, unit tests would be a lot better, but I don't think we have a directory with AArch64-specific unit tests. We should try to add one at some point.

[kavon]

When the op code is a patchpoint, isn't a call still emitted? I think maybe the total number of bytes returned for patchpoints should include the bytes for the call.

[rovka]

In D24073#539438, @kavon wrote:

When the op code is a patchpoint, isn't a call still emitted? I think maybe the total number of bytes returned for patchpoints should include the bytes for the call.

From my reading of AArch64AsmPrinter::LowerPATCHPOINT, it looks like the number of bytes requested by the patchpoint has to include the bytes for the call - i.e. if the call is generated, its size is subtracted from the number of bytes requested, and only the remaining number of nops are added.

[kavon]

In D24073#539679, @rovka wrote:

In D24073#539438, @kavon wrote:

When the op code is a patchpoint, isn't a call still emitted? I think maybe the total number of bytes returned for patchpoints should include the bytes for the call.

From my reading of AArch64AsmPrinter::LowerPATCHPOINT, it looks like the number of bytes requested by the patchpoint has to include the bytes for the call - i.e. if the call is generated, its size is subtracted from the number of bytes requested, and only the remaining number of nops are added.

Ah, okay. This looks good to me then. I'm a newcomer when it comes to contributing patches, so I'm not sure what else I should do. :)

[rovka]

In D24073#539895, @kavon wrote:

I'm a newcomer when it comes to contributing patches, so I'm not sure what else I should do. :)

You don't have to do anything else, I'll commit it tomorrow :) Thanks for reviewing.

llvm/trunk/lib/Target/AArch64/AArch64InstrInfo.cpp

 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
 #include "MCTargetDesc/AArch64AddressingModes.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/PseudoSourceValue.h"
+#include "llvm/CodeGen/StackMaps.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/TargetRegistry.h"
 #include <algorithm>
 
 using namespace llvm;
 
 #define GET_INSTRINFO_CTOR_DTOR

 unsigned AArch64InstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   const MachineBasicBlock &MBB = *MI.getParent();
   const MachineFunction *MF = MBB.getParent();
   const MCAsmInfo *MAI = MF->getTarget().getMCAsmInfo();
 
   if (MI.getOpcode() == AArch64::INLINEASM)
     return getInlineAsmLength(MI.getOperand(0).getSymbolName(), *MAI);
 
+  // FIXME: We currently only handle pseudoinstructions that don't get expanded
+  //        before the assembly printer.
+  unsigned NumBytes = 0;
   const MCInstrDesc &Desc = MI.getDesc();
   switch (Desc.getOpcode()) {
   default:
     // Anything not explicitly designated otherwise is a normal 4-byte insn.
-    return 4;
+    NumBytes = 4;
+    break;
   case TargetOpcode::DBG_VALUE:
   case TargetOpcode::EH_LABEL:
   case TargetOpcode::IMPLICIT_DEF:
   case TargetOpcode::KILL:
-    return 0;
+    NumBytes = 0;
+    break;
+  case TargetOpcode::STACKMAP:
+    // The upper bound for a stackmap intrinsic is the full length of its shadow
+    NumBytes = StackMapOpers(&MI).getNumPatchBytes();
+    assert(NumBytes % 4 == 0 && "Invalid number of NOP bytes requested!");
+    break;
+  case TargetOpcode::PATCHPOINT:
+    // The size of the patchpoint intrinsic is the number of bytes requested
+    NumBytes = PatchPointOpers(&MI).getNumPatchBytes();
+    assert(NumBytes % 4 == 0 && "Invalid number of NOP bytes requested!");
+    break;
   case AArch64::TLSDESC_CALLSEQ:
     // This gets lowered to an instruction sequence which takes 16 bytes
-    return 16;
+    NumBytes = 16;
+    break;
   }
 
-  llvm_unreachable("getInstSizeInBytes()- Unable to determine insn size");
+  return NumBytes;
 }
 
 static void parseCondBranch(MachineInstr *LastInst, MachineBasicBlock *&Target,
                             SmallVectorImpl<MachineOperand> &Cond) {
   // Block ends with fall-through condbranch.
   switch (LastInst->getOpcode()) {
   default:
     llvm_unreachable("Unknown branch instruction?");

llvm/trunk/test/CodeGen/MIR/AArch64/inst-size-patchpoint.mir

+# RUN: llc -mtriple=aarch64-unknown -run-pass aarch64-branch-relax -aarch64-tbz-offset-bits=4 %s -o - | FileCheck %s
+--- |
+  ; ModuleID = '/tmp/test.ll'
+  source_filename = "test.ll"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-unknown"
+
+  define void @test_patchpoint_length(i32 %in) {
+    %val = and i32 %in, 1
+    %tst = icmp eq i32 %val, 0
+    br i1 %tst, label %true, label %false
+
+  true:                                             ; preds = %0
+    tail call void (i64, i32, i8*, i32, ...) @llvm.experimental.patchpoint.void(i64 0, i32 32, i8* null, i32 0)
+    ret void
+
+  false:                                            ; preds = %0
+    ret void
+  }
+
+  declare void @llvm.experimental.patchpoint.void(i64, i32, i8*, i32, ...)
+...
+---
+# CHECK-LABEL: name:{{.*}}test_patchpoint_length
+# If the size of the patchpoint is computed correctly, that will push
+# the bb.2.false block too far away from the TBNZW, so the branch will
+# have to be relaxed (note that we're using -aarch64-tbz-offset-bits to
+# constrain the range that can be reached with the TBNZW to something smaller
+# than what the patchpoint is lowered to).
+# CHECK: TBZW killed %w0, 0, %bb.1.true
+# CHECK: B %bb.2.false
+name:            test_patchpoint_length
+alignment:       2
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+tracksRegLiveness: false
+liveins:
+  - { reg: '%w0' }
+calleeSavedRegisters: [ '%fp', '%lr', '%b8', '%b9', '%b10', '%b11', '%b12',
+                        '%b13', '%b14', '%b15', '%d8', '%d9', '%d10', '%d11',
+                        '%d12', '%d13', '%d14', '%d15', '%h8', '%h9', '%h10',
+                        '%h11', '%h12', '%h13', '%h14', '%h15', '%s8',
+                        '%s9', '%s10', '%s11', '%s12', '%s13', '%s14',
+                        '%s15', '%w19', '%w20', '%w21', '%w22', '%w23',
+                        '%w24', '%w25', '%w26', '%w27', '%w28', '%w29',
+                        '%w30', '%x19', '%x20', '%x21', '%x22', '%x23',
+                        '%x24', '%x25', '%x26', '%x27', '%x28', '%d8_d9',
+                        '%d9_d10', '%d10_d11', '%d11_d12', '%d12_d13',
+                        '%d13_d14', '%d14_d15', '%d8_d9_d10_d11', '%d9_d10_d11_d12',
+                        '%d10_d11_d12_d13', '%d11_d12_d13_d14', '%d12_d13_d14_d15',
+                        '%d8_d9_d10', '%d9_d10_d11', '%d10_d11_d12', '%d11_d12_d13',
+                        '%d12_d13_d14', '%d13_d14_d15', '%w19_w20', '%w20_w21',
+                        '%w21_w22', '%w22_w23', '%w23_w24', '%w24_w25',
+                        '%w25_w26', '%w26_w27', '%w27_w28', '%w28_w29',
+                        '%w29_w30', '%fp_lr', '%x28_fp', '%x19_x20', '%x20_x21',
+                        '%x21_x22', '%x22_x23', '%x23_x24', '%x24_x25',
+                        '%x25_x26', '%x26_x27', '%x27_x28' ]
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   true
+  stackSize:       16
+  offsetAdjustment: 0
+  maxAlignment:    8
+  adjustsStack:    true
+  hasCalls:        true
+  maxCallFrameSize: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+stack:
+  - { id: 0, type: spill-slot, offset: -8, size: 8, alignment: 8, callee-saved-register: '%lr' }
+  - { id: 1, type: spill-slot, offset: -16, size: 8, alignment: 8, callee-saved-register: '%fp' }
+body:             |
+  bb.0 (%ir-block.0):
+    successors: %bb.1.true, %bb.2.false
+    liveins: %w0, %lr, %fp
+
+    TBNZW killed %w0, 0, %bb.2.false
+
+  bb.1.true:
+    successors: %bb.2.false
+    liveins: %lr, %fp
+
+    PATCHPOINT 0, 32, 0, 0, 0, csr_aarch64_aapcs, implicit-def dead early-clobber %x16, implicit-def dead early-clobber %x17, implicit-def dead early-clobber %lr, implicit-def %sp
+
+  bb.2.false:
+    liveins: %lr, %fp
+
+    RET killed %lr
+
+...

llvm/trunk/test/CodeGen/MIR/AArch64/inst-size-stackmap.mir

+# RUN: llc -mtriple=aarch64-unknown -run-pass aarch64-branch-relax -aarch64-tbz-offset-bits=4 %s -o - | FileCheck %s
+--- |
+  ; ModuleID = 'test.ll'
+  source_filename = "test.ll"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-unknown"
+
+  define void @test_stackmap_length(i32 %in) {
+    %val = and i32 %in, 1
+    %tst = icmp eq i32 %val, 0
+    br i1 %tst, label %true, label %false
+
+  true:                                             ; preds = %0
+    tail call void (i64, i32, ...) @llvm.experimental.stackmap(i64 0, i32 32)
+    ret void
+
+  false:                                            ; preds = %0
+    ret void
+  }
+
+  declare void @llvm.experimental.stackmap(i64, i32, ...)
+...
+---
+# CHECK-LABEL: name:{{.*}}test_stackmap_length
+# If the size of the stackmap is computed correctly, that will push
+# the bb.2.false block too far away from the TBNZW, so the branch will
+# have to be relaxed (note that we're using -aarch64-tbz-offset-bits to
+# constrain the range that can be reached with the TBNZW to something smaller
+# than what the stackmap is lowered to).
+# CHECK: TBZW killed %w0, 0, %bb.1.true
+# CHECK: B %bb.2.false
+name:            test_stackmap_length
+alignment:       2
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+tracksRegLiveness: false
+liveins:
+  - { reg: '%w0' }
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     true
+  hasPatchPoint:   false
+  stackSize:       16
+  offsetAdjustment: 0
+  maxAlignment:    8
+  adjustsStack:    true
+  hasCalls:        true
+  maxCallFrameSize: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+stack:
+  - { id: 0, type: spill-slot, offset: -8, size: 8, alignment: 8, callee-saved-register: '%lr' }
+  - { id: 1, type: spill-slot, offset: -16, size: 8, alignment: 8, callee-saved-register: '%fp' }
+body:             |
+  bb.0 (%ir-block.0):
+    successors: %bb.1.true, %bb.2.false
+    liveins: %w0, %lr, %fp
+
+    TBNZW killed %w0, 0, %bb.2.false
+
+  bb.1.true:
+    successors: %bb.2.false
+    liveins: %lr, %fp
+
+    STACKMAP 0, 32, implicit-def dead early-clobber %x16, implicit-def dead early-clobber %x17, implicit-def dead early-clobber %lr
+
+  bb.2.false:
+    liveins: %lr, %fp
+
+    RET killed %lr
+