[XRay] Detect loops in functions being lowered

[dberris]

Summary

This is an implementation of the loop detection logic that XRay needs to
determine whether a function might take time at runtime. Without this
heuristic, XRay will tend to not instrument short functions that have
loops that might have runtime dependent on inputs or external values.

While this implementation doesn't do any further analysis than just
figuring out whether there is a loop in the MachineFunction being
code-gen'ed, we're paving the way for being able to perform more
sophisticated analysis of the function in the future (for example to
determine whether the trip count for the loop might be constant, and
make a decision on that instead). This enables us to cover more
functions with the default heuristics, and potentially identify ones
that have variable runtime latency just by looking for the presence of
loops.

[dberris]

ping

@chandlerc -- specifically looking for advice on whether the preservation logic is correct in the analysis/transformation dependencies.

[rnk]

Looks good with a small suggestion

lib/CodeGen/XRayInstrumentation.cpp

40 ↗

(On Diff #95899)

This pass at least preserves MachineDomTree, so let's add that too. It seems consistent with other MI passes.

[dberris]

Thanks @rnk!

llvm/trunk/lib/CodeGen/XRayInstrumentation.cpp

 // operations.
 //
 //===---------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/Analysis.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
+#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Target/TargetInstrInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 
 using namespace llvm;
 
 namespace {
 struct XRayInstrumentation : public MachineFunctionPass {
   static char ID;
 
   XRayInstrumentation() : MachineFunctionPass(ID) {
     initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
   }
 
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesCFG();
+    AU.addRequired<MachineLoopInfo>();
+    AU.addPreserved<MachineLoopInfo>();
+    AU.addPreserved<MachineDominatorTree>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
   bool runOnMachineFunction(MachineFunction &MF) override;
 
 private:
   // Replace the original RET instruction with the exit sled code ("patchable
   //   ret" pseudo-instruction), so that at runtime XRay can replace the sled
   //   with a code jumping to XRay trampoline, which calls the tracing handler
   //   and, in the end, issues the RET instruction.
   // This is the approach to go on CPUs which have a single RET instruction,
   //   like x86/x86_64.
   void replaceRetWithPatchableRet(MachineFunction &MF,
                                   const TargetInstrInfo *TII);
 
   // Prepend the original return instruction with the exit sled code ("patchable
   //   function exit" pseudo-instruction), preserving the original return
   //   instruction just after the exit sled code.
   // This is the approach to go on CPUs which have multiple options for the
   //   return instruction, like ARM. For such CPUs we can't just jump into the
   //   XRay trampoline and issue a single return instruction there. We rather
   //   have to call the trampoline and return from it to the original return
   //   instruction of the function being instrumented.
   void prependRetWithPatchableExit(MachineFunction &MF,
                                    const TargetInstrInfo *TII);
 };
 } // anonymous namespace
 
-void XRayInstrumentation::replaceRetWithPatchableRet(MachineFunction &MF,
-  const TargetInstrInfo *TII)
-{
+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:
for (auto &T : MBB.terminators()) {
       }
     }
   }
 
   for (auto &I : Terminators)
     I->eraseFromParent();
 }
 
-void XRayInstrumentation::prependRetWithPatchableExit(MachineFunction &MF,
-  const TargetInstrInfo *TII)
-{
+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));
+        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.
+
+    // Check if we have a loop.
+    // FIXME: Maybe make this smarter, and see whether the loops are dependent
+    // on inputs or side-effects?
+    MachineLoopInfo &MLI = getAnalysis<MachineLoopInfo>();
+    if (MLI.empty() && F.size() < XRayThreshold)
+      return false; // Function is too small and has no loops.
   }
 
   // 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.
   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
 
   switch (MF.getTarget().getTargetTriple().getArch()) {
   case Triple::ArchType::arm:
   case Triple::ArchType::thumb:
default:
     replaceRetWithPatchableRet(MF, TII);
     break;
   }
   return true;
 }
 
 char XRayInstrumentation::ID = 0;
 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
-INITIALIZE_PASS(XRayInstrumentation, "xray-instrumentation", "Insert XRay ops",
-                false, false)
+INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
+                      "Insert XRay ops", false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
+INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
+                    "Insert XRay ops", false, false)

llvm/trunk/test/CodeGen/X86/xray-loop-detection.ll

+; RUN: llc -filetype=asm -o - -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+; RUN: llc -filetype=asm -o - -mtriple=x86_64-darwin-unknown    < %s | FileCheck %s
+
+define i32 @foo(i32 %i) nounwind noinline uwtable "xray-instruction-threshold"="1" {
+entry:
+  br label %Test
+Test:
+  %indvar = phi i32 [0, %entry], [%nextindvar, %Inc]
+  %cond = icmp eq i32 %indvar, %i
+  br i1 %cond, label %Exit, label %Inc
+Inc:
+  %nextindvar = add i32 %indvar, 1
+  br label %Test
+Exit:
+  %retval = phi i32 [%indvar, %Test]
+  ret i32 %retval
+}
+
+; CHECK-LABEL: xray_sled_0:
+; CHECK-NEXT:  .ascii "\353\t"
+; CHECK-NEXT:  nopw 512(%rax,%rax)
+; CHECK-LABEL: Ltmp0:
+