[SystemZ] Add proper mcount handling

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

@@ -17,7 +17,9 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
+#include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
 #include "llvm/Target/TargetMachine.h"
 
 using namespace llvm;
@@ -557,6 +559,36 @@ void SystemZELFFrameLowering::emitPrologue(MachineFunction &MF,
   // Debug location must be unknown since the first debug location is used
   // to determine the end of the prologue.
   DebugLoc DL;
+  // Add mcount instrumentation if necessary.
+  if (MF.getFunction()
+          .getFnAttribute("systemz-instrument-function-entry")
+          .getValueAsString() == "mcount") {
+
+    // Store return address 8 bytes above stack pointer.
+    BuildMI(MBB, MBBI, DL, ZII->get(SystemZ::STG))
+        .addReg(SystemZ::R14D)
+        .addReg(SystemZ::R15D)
+        .addImm(8)
+        .addReg(0);
+
+    // Call mcount (Regmask from CC AnyReg since mcount preserves all normal
+    // argument registers).
+    FunctionCallee FC = MF.getFunction().getParent()->getOrInsertFunction(
+        "mcount", Type::getVoidTy(MF.getFunction().getContext()));
+    const uint32_t *Mask = MF.getSubtarget<SystemZSubtarget>()
+                               .getSpecialRegisters()
+                               ->getCallPreservedMask(MF, CallingConv::AnyReg);
+    BuildMI(MBB, MBBI, DL, ZII->get(SystemZ::CallBRASL))
+        .addGlobalAddress(dyn_cast<Function>(FC.getCallee()))
+        .addRegMask(Mask);
+
+    // Reload return address from 8 bytes above stack pointer.
+    BuildMI(MBB, MBBI, DL, ZII->get(SystemZ::LG))
+        .addReg(SystemZ::R14D)
+        .addReg(SystemZ::R15D)
+        .addImm(8)
+        .addReg(0);
+  }
 
   // The current offset of the stack pointer from the CFA.
   int64_t SPOffsetFromCFA = -SystemZMC::ELFCFAOffsetFromInitialSP;

llvm/lib/Transforms/Utils/EntryExitInstrumenter.cpp

@@ -63,6 +63,11 @@ static void insertCall(Function &CurFn, StringRef Func,
                                   false));
       CallInst *Call = CallInst::Create(Fn, RetAddr, "", InsertionPt);
       Call->setDebugLoc(DL);
+    } else if (TargetTriple.isSystemZ()) {
+      // skip insertion for `mcount` on SystemZ. This will be handled later in
+      // `emitPrologue`. Add custom attribute to denote this.
+      CurFn.addFnAttr(
+          llvm::Attribute::get(C, "systemz-instrument-function-entry", Func));
     } else {
       FunctionCallee Fn = M.getOrInsertFunction(Func, Type::getVoidTy(C));
       CallInst *Call = CallInst::Create(Fn, "", InsertionPt);

llvm/test/CodeGen/SystemZ/mcount.ll

@@ -0,0 +1,35 @@
+; Test proper insertion of mcount instrumentation
+;
+; RUN: llc < %s -mtriple=s390x-linux-gnu -o - | FileCheck %s
+;
+; CHECK: # %bb.0:
+; CHECK-NEXT: stg %r14, 8(%r15)
+; CHECK-NEXT: brasl %r14, mcount@PLT
+; CHECK-NEXT: lg %r14, 8(%r15)
+define dso_local signext i32 @fib(i32 noundef signext %n) #0 {
+entry:
+  %n.addr = alloca i32, align 4
+  store i32 %n, ptr %n.addr, align 4
+  %0 = load i32, ptr %n.addr, align 4
+  %cmp = icmp sle i32 %0, 1
+  br i1 %cmp, label %cond.true, label %cond.false
+
+cond.true:                                        ; preds = %entry
+  br label %cond.end
+
+cond.false:                                       ; preds = %entry
+  %1 = load i32, ptr %n.addr, align 4
+  %sub = sub nsw i32 %1, 1
+  %call = call signext i32 @fib(i32 noundef signext %sub)
+  %2 = load i32, ptr %n.addr, align 4
+  %sub1 = sub nsw i32 %2, 2
+  %call2 = call signext i32 @fib(i32 noundef signext %sub1)
+  %add = add nsw i32 %call, %call2
+  br label %cond.end
+
+cond.end:                                         ; preds = %cond.false, %cond.true
+  %cond = phi i32 [ 1, %cond.true ], [ %add, %cond.false ]
+  ret i32 %cond
+}
+
+attributes #0 = { "instrument-function-entry-inlined"="mcount" }