[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

61–62 ↗

(On Diff #81919)

Unrelated formatting change, please land in a separate commit.

93–94 ↗

(On Diff #81919)

Ditto.

114–115 ↗

(On Diff #81919)

Add a test case for this as well?

136–140 ↗

(On Diff #81919)

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

114–115 ↗

(On Diff #81919)

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

136–140 ↗

(On Diff #81919)

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 ↗

(On Diff #81922)

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 ↗

(On Diff #81922)

Awesome! I keep forgetting that. :)

[dberris]

• Use llvm::find_if instead

llvm/trunk/lib/CodeGen/XRayInstrumentation.cpp

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 =
+      find_if(MF, [&](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

llvm/trunk/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
+...

llvm/trunk/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
+...
+