[XRay] Fix assertion failure on empty machine basic blocks (PR 31424)

[dberris]

Summary

The original version of the code in XRayInstrumentation.cpp assumed that
functions may not have empty machine basic blocks (or that the first one
couldn't be). This change addresses that by special-casing that specific
situation.

We provide a .mir test-case to make sure we're handling this
appropriately.

Fixes llvm.org/PR31424.

[chandlerc]

Also good to add an IR test case with 'unreachable' just for more coverage of these kinds of edge cases?

lib/CodeGen/XRayInstrumentation.cpp

63–64

Unrelated formatting change, please land in a separate commit.

96–97

Ditto.

116–117

Add a test case for this as well?

150–154

Why do you care about instrumenting functions which have no instructions in them?

[dberris]

• Do not instrument empty functions
• Add test for an empty function

lib/CodeGen/XRayInstrumentation.cpp

116–117

I'm not sure how to do that from a .mir properly, but I tried anyway. :)

150–154

Good point -- mostly because I'm not sure whether the assumption that the first MBB being empty implies that there are no instructions in the rest of the function. I'm happy to ignore this case (maybe because it's when an inlined function is there just for the debug info case).

[chandlerc]

LGTM

lib/CodeGen/XRayInstrumentation.cpp

135

You can use just find_if(MF, ...) -- we have a range based version in the llvm namespace.

[dberris]

Thanks @chandlerc!

lib/CodeGen/XRayInstrumentation.cpp

135

Awesome! I keep forgetting that. :)

[dberris]

• Use llvm::find_if instead

lib/CodeGen/XRayInstrumentation.cpp

private:
   //   instruction of the function being instrumented.
   void prependRetWithPatchableExit(MachineFunction &MF,
     const TargetInstrInfo *TII);
 };
 } // anonymous namespace
 
 void XRayInstrumentation::replaceRetWithPatchableRet(MachineFunction &MF,
   const TargetInstrInfo *TII)
 {
   // We look for *all* terminators and returns, then replace those with
   // PATCHABLE_RET instructions.
   SmallVector<MachineInstr *, 4> Terminators;
   for (auto &MBB : MF) {
     for (auto &T : MBB.terminators()) {
       unsigned Opc = 0;
       if (T.isReturn() && T.getOpcode() == TII->getReturnOpcode()) {
         // Replace return instructions with:
         //   PATCHABLE_RET <Opcode>, <Operand>...
void XRayInstrumentation::replaceRetWithPatchableRet(MachineFunction &MF,
   }
 
   for (auto &I : Terminators)
     I->eraseFromParent();
 }
 
 void XRayInstrumentation::prependRetWithPatchableExit(MachineFunction &MF,
   const TargetInstrInfo *TII)
 {
   for (auto &MBB : MF) {
     for (auto &T : MBB.terminators()) {
       unsigned Opc = 0;
       if (T.isReturn()) {
         Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
       }
       if (TII->isTailCall(T)) {
         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
       }
       if (Opc != 0) {
         // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
         //   PATCHABLE_TAIL_CALL .
         BuildMI(MBB, T, T.getDebugLoc(),TII->get(Opc));
       }
     }
   }
 }
 
 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
   auto &F = *MF.getFunction();
   auto InstrAttr = F.getFnAttribute("function-instrument");
   bool AlwaysInstrument = !InstrAttr.hasAttribute(Attribute::None) &&
                           InstrAttr.isStringAttribute() &&
                           InstrAttr.getValueAsString() == "xray-always";
   Attribute Attr = F.getFnAttribute("xray-instruction-threshold");
   unsigned XRayThreshold = 0;
   if (!AlwaysInstrument) {
     if (Attr.hasAttribute(Attribute::None) || !Attr.isStringAttribute())
       return false; // XRay threshold attribute not found.
     if (Attr.getValueAsString().getAsInteger(10, XRayThreshold))
       return false; // Invalid value for threshold.
     if (F.size() < XRayThreshold)
       return false; // Function is too small.
   }
 
-  auto &FirstMBB = *MF.begin();
+  // We look for the first non-empty MachineBasicBlock, so that we can insert
+  // the function instrumentation in the appropriate place.
+  auto MBI =
+      std::find_if(MF.begin(), MF.end(),
+                   [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
+  if (MBI == MF.end())
+    return false; // The function is empty.
+
+  auto *TII = MF.getSubtarget().getInstrInfo();
+  auto &FirstMBB = *MBI;
   auto &FirstMI = *FirstMBB.begin();
 
   if (!MF.getSubtarget().isXRaySupported()) {
     FirstMI.emitError("An attempt to perform XRay instrumentation for an"
       " unsupported target.");
     return false;
   }
 
   // FIXME: Do the loop triviality analysis here or in an earlier pass.
 
   // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
   // MachineFunction.
-  auto *TII = MF.getSubtarget().getInstrInfo();
   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
 
   switch (MF.getTarget().getTargetTriple().getArch()) {
   case Triple::ArchType::arm:
   case Triple::ArchType::thumb:
   case Triple::ArchType::aarch64:
     // For the architectures which don't have a single return instruction
     prependRetWithPatchableExit(MF, TII);

test/CodeGen/X86/xray-empty-firstmbb.mir

+# RUN: llc -run-pass xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
+#
+# Make sure we can handle empty first basic blocks.
+
+--- |
+
+  define i32 @foo() noinline uwtable "xray-instruction-threshold"="1" {
+  entry:
+    unreachable
+  }
+
+...
+
+---
+name: foo
+tracksRegLiveness: true
+liveins:
+  - { reg: '%edi'}
+body:            |
+  bb.0.entry:
+    liveins: %edi
+    ; CHECK-NOT: PATCHABLE_FUNCTION_ENTER
+...

test/CodeGen/X86/xray-empty-function.mir

+# RUN: llc -run-pass xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
+#
+# Make sure we can handle empty functions.
+---
+name: empty
+tracksRegLiveness: true
+liveins:
+  - { reg: '%edi'}
+body:            |
+  bb.0:
+    ; CHECK-NOT: PATCHABLE_FUNCTION_ENTER
+...
+