[MCWin64EH] Check for matches between epilogs and the prolog on ARM64

[mstorsjo]

Summary

This allows sharing opcodes between prolog and epilog even when there
is more than one epilog.

I didn't make any handcrafted special MC level testcases for this (yet
at least), but it does seem to have the expected effect on two existing
CodeGen level testcases.

[mstorsjo]

Updated the patch to remove a fixme for something that is mostly implemented already.

[efriedma]

llvm/lib/MC/MCWin64EH.cpp

622

subset of the *prolog*

[mstorsjo]

Fix the typo that was pointed out

[efriedma]

LGTM

llvm/lib/MC/MCWin64EH.cpp

if (Reverse) {
     for (auto It = Instructions.rbegin(); It != Instructions.rend(); It++)
       VisitInstruction(*It);
   } else {
     for (WinEH::Instruction &Inst : Instructions)
       VisitInstruction(Inst);
   }
 }
 
+// Check if an epilog exists as a subset of the end of a prolog (backwards).
+static int getOffsetInProlog(const std::vector<WinEH::Instruction> &Prolog,
+                             const std::vector<WinEH::Instruction> &Epilog) {
+  // Can't find an epilog as a subset if it is longer than the prolog.
+  if (Epilog.size() > Prolog.size())
+    return -1;
+
+  // Check that the epilog actually is a perfect match for the end (backwrds)
+  // of the prolog.
+  for (int I = Epilog.size() - 1; I >= 0; I--) {
+    if (Prolog[I] != Epilog[Epilog.size() - 1 - I])
+      return -1;
+  }
+
+  // If the epilog was a subset of the epilog, find its offset.
+  if (Epilog.size() == Prolog.size())
+    return 0;
+  return ARM64CountOfUnwindCodes(ArrayRef<WinEH::Instruction>(
+      &Prolog[Epilog.size()], Prolog.size() - Epilog.size()));
+}
+
 static int checkPackedEpilog(MCStreamer &streamer, WinEH::FrameInfo *info,
                              int PrologCodeBytes) {
   // Can only pack if there's one single epilog
   if (info->EpilogMap.size() != 1)
     return -1;
 
   const std::vector<WinEH::Instruction> &Epilog =
       info->EpilogMap.begin()->second;
static int checkPackedEpilog(MCStreamer &streamer, WinEH::FrameInfo *info,
   // Even if we don't end up sharing opcodes with the prolog, we can still
   // write the offset as a packed offset, if the single epilog is located at
   // the end of the function and the offset (pointing after the prolog) fits
   // as a packed offset.
   if (PrologCodeBytes <= 31 &&
       PrologCodeBytes + ARM64CountOfUnwindCodes(Epilog) <= 124)
     RetVal = PrologCodeBytes;
 
-  // Can pack if the epilog is a subset of the prolog but not vice versa
-  if (Epilog.size() > info->Instructions.size())
+  int Offset = getOffsetInProlog(info->Instructions, Epilog);
+  if (Offset < 0)
     return RetVal;
 
-  // Check that the epilog actually is a perfect match for the end (backwrds)
-  // of the prolog.
-  for (int I = Epilog.size() - 1; I >= 0; I--) {
-    if (info->Instructions[I] != Epilog[Epilog.size() - 1 - I])
-      return RetVal;
-  }
-
-  int Offset = Epilog.size() == info->Instructions.size()
-                   ? 0
-                   : ARM64CountOfUnwindCodes(ArrayRef<WinEH::Instruction>(
-                         &info->Instructions[Epilog.size()],
-                         info->Instructions.size() - Epilog.size()));
-
   // Check that the offset and prolog size fits in the first word; it's
   // unclear whether the epilog count in the extension word can be taken
   // as packed epilog offset.
   if (Offset > 31 || PrologCodeBytes > 124)
     return RetVal;
 
   // As we choose to express the epilog as part of the prolog, remove the
   // epilog from the map, so we don't try to emit its opcodes.
static void ARM64EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info,
 
   for (auto &I : info->EpilogMap) {
     MCSymbol *EpilogStart = I.first;
     auto &EpilogInstrs = I.second;
     uint32_t CodeBytes = ARM64CountOfUnwindCodes(EpilogInstrs);
 
     MCSymbol* MatchingEpilog =
       FindMatchingEpilog(EpilogInstrs, AddedEpilogs, info);
+    int PrologOffset;
     if (MatchingEpilog) {
       assert(EpilogInfo.find(MatchingEpilog) != EpilogInfo.end() &&
              "Duplicate epilog not found");
       EpilogInfo[EpilogStart] = EpilogInfo.lookup(MatchingEpilog);
       // Clear the unwind codes in the EpilogMap, so that they don't get output
       // in the logic below.
       EpilogInstrs.clear();
+    } else if ((PrologOffset =
+                    getOffsetInProlog(info->Instructions, EpilogInstrs)) >= 0) {
+      EpilogInfo[EpilogStart] = PrologOffset;
+      // Clear the unwind codes in the EpilogMap, so that they don't get output
+      // in the logic below.
+      EpilogInstrs.clear();
     } else {
       EpilogInfo[EpilogStart] = TotalCodeBytes;
       TotalCodeBytes += CodeBytes;
       AddedEpilogs.push_back(EpilogStart);
     }
   }
 
   // Code Words, Epilog count, E, X, Vers, Function Length
static void ARM64EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info,
   if (PackedEpilogOffset >= 0) // E
     row1 |= 1 << 21;
   row1 |= FuncLength & 0x3FFFF;
   streamer.emitInt32(row1);
 
   // Extended Code Words, Extended Epilog Count
   if (ExtensionWord) {
     // FIXME: We should be able to split unwind info into multiple sections.
-    // FIXME: We should share epilog codes across epilogs, where possible,
-    // which would make this issue show up less frequently.
     if (CodeWords > 0xFF || EpilogCount > 0xFFFF)
       report_fatal_error("SEH unwind data splitting not yet implemented");
     uint32_t row2 = 0x0;
     row2 |= (CodeWords & 0xFF) << 16;
     row2 |= (EpilogCount & 0xFFFF);
     streamer.emitInt32(row2);
   }
 

llvm/test/CodeGen/AArch64/wineh4.mir

 # RUN: llc -o - %s -mtriple=aarch64-windows -start-after=prologepilog \
 # RUN:   -disable-branch-fold  -filetype=obj \
 # RUN: | llvm-readobj --unwind - | FileCheck %s
 # Check that identical multiple epilgoues are correctly shared in .xdata.
 
 # CHECK:        ExceptionData {
 # CHECK-NEXT:      FunctionLength: 164
 # CHECK-NEXT:      Version: 0
 # CHECK-NEXT:      ExceptionData: No
 # CHECK-NEXT:      EpiloguePacked: No
 # CHECK-NEXT:      EpilogueScopes: 2
-# CHECK-NEXT:      ByteCodeLength: 32
+# CHECK-NEXT:      ByteCodeLength: 16
 # CHECK-NEXT:      Prologue [
 # CHECK-NEXT:        0xc80c              ; stp x19, x20, [sp, #96]
 # CHECK-NEXT:        0xc88a              ; stp x21, x22, [sp, #80]
 # CHECK-NEXT:        0xc908              ; stp x23, x24, [sp, #64]
 # CHECK-NEXT:        0xc986              ; stp x25, x26, [sp, #48]
 # CHECK-NEXT:        0xca04              ; stp x27, x28, [sp, #32]
 # CHECK-NEXT:        0xd802              ; stp d8, d9, [sp, #16]
 # CHECK-NEXT:        0xda8d              ; stp d10, d11, [sp, #-112]!
 # CHECK-NEXT:        0xe4                ; end
 # CHECK-NEXT:      ]
 # CHECK-NEXT:      EpilogueScopes [
 # CHECK-NEXT:        EpilogueScope {
 # CHECK-NEXT:          StartOffset: 16
-# CHECK-NEXT:          EpilogueStartIndex: 15
+# CHECK-NEXT:          EpilogueStartIndex: 0
 # CHECK-NEXT:          Opcodes [
 # CHECK-NEXT:            0xc80c              ; ldp x19, x20, [sp, #96]
 # CHECK-NEXT:            0xc88a              ; ldp x21, x22, [sp, #80]
 # CHECK-NEXT:            0xc908              ; ldp x23, x24, [sp, #64]
 # CHECK-NEXT:            0xc986              ; ldp x25, x26, [sp, #48]
 # CHECK-NEXT:            0xca04              ; ldp x27, x28, [sp, #32]
 # CHECK-NEXT:            0xd802              ; ldp d8, d9, [sp, #16]
 # CHECK-NEXT:            0xda8d              ; ldp d10, d11, [sp], #112
 # CHECK-NEXT:            0xe4                ; end
 # CHECK-NEXT:          ]
 # CHECK-NEXT:        }
 # CHECK-NEXT:        EpilogueScope {
 # CHECK-NEXT:          StartOffset: 33
-# CHECK-NEXT:          EpilogueStartIndex: 15
+# CHECK-NEXT:          EpilogueStartIndex: 0
 # CHECK-NEXT:          Opcodes [
 # CHECK-NEXT:            0xc80c              ; ldp x19, x20, [sp, #96]
 # CHECK-NEXT:            0xc88a              ; ldp x21, x22, [sp, #80]
 # CHECK-NEXT:            0xc908              ; ldp x23, x24, [sp, #64]
 # CHECK-NEXT:            0xc986              ; ldp x25, x26, [sp, #48]
 # CHECK-NEXT:            0xca04              ; ldp x27, x28, [sp, #32]
 # CHECK-NEXT:            0xd802              ; ldp d8, d9, [sp, #16]
 # CHECK-NEXT:            0xda8d              ; ldp d10, d11, [sp], #112

llvm/test/CodeGen/AArch64/wineh8.mir

 # RUN: llc -o - %s -mtriple=aarch64-windows -start-after=prologepilog \
 # RUN:   -disable-branch-fold  -filetype=obj \
 # RUN: | llvm-readobj --unwind - | FileCheck %s
 # Check that non-identical multiple epilgoues are correctly shared in .xdata.
 
 # CHECK:        ExceptionData {
 # CHECK-NEXT:      FunctionLength: 160
 # CHECK-NEXT:      Version: 0
 # CHECK-NEXT:      ExceptionData: No
 # CHECK-NEXT:      EpiloguePacked: No
 # CHECK-NEXT:      EpilogueScopes: 2
-# CHECK-NEXT:      ByteCodeLength: 44
+# CHECK-NEXT:      ByteCodeLength: 28
 # CHECK-NEXT:      Prologue [
 # CHECK-NEXT:        0xc80c              ; stp x19, x20, [sp, #96]
 # CHECK-NEXT:        0xc88a              ; stp x21, x22, [sp, #80]
 # CHECK-NEXT:        0xc908              ; stp x23, x24, [sp, #64]
 # CHECK-NEXT:        0xc986              ; stp x25, x26, [sp, #48]
 # CHECK-NEXT:        0xca04              ; stp x27, x28, [sp, #32]
 # CHECK-NEXT:        0xd802              ; stp d8, d9, [sp, #16]
 # CHECK-NEXT:        0xda8d              ; stp d10, d11, [sp, #-112]!

 # CHECK-NEXT:            0xc986              ; ldp x25, x26, [sp, #48]
 # CHECK-NEXT:            0xd802              ; ldp d8, d9, [sp, #16]
 # CHECK-NEXT:            0xda8d              ; ldp d10, d11, [sp], #112
 # CHECK-NEXT:            0xe4                ; end
 # CHECK-NEXT:          ]
 # CHECK-NEXT:        }
 # CHECK-NEXT:        EpilogueScope {
 # CHECK-NEXT:          StartOffset: 32
-# CHECK-NEXT:          EpilogueStartIndex: 28
+# CHECK-NEXT:          EpilogueStartIndex: 0
 # CHECK-NEXT:          Opcodes [
 # CHECK-NEXT:            0xc80c              ; ldp x19, x20, [sp, #96]
 # CHECK-NEXT:            0xc88a              ; ldp x21, x22, [sp, #80]
 # CHECK-NEXT:            0xc908              ; ldp x23, x24, [sp, #64]
 # CHECK-NEXT:            0xc986              ; ldp x25, x26, [sp, #48]
 # CHECK-NEXT:            0xca04              ; ldp x27, x28, [sp, #32]
 # CHECK-NEXT:            0xd802              ; ldp d8, d9, [sp, #16]
 # CHECK-NEXT:            0xda8d              ; ldp d10, d11, [sp], #112