[llvm] [EntryExitInstrumenter][AArch64][RISCV][LoongArch] Pass `__builtin_return_address(0)` into `_mcount` (PR #121107)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw updated https://github.com/llvm/llvm-project/pull/121107

>From b3d962278b47048f232a4362de7e143cca9ff1dc Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Wed, 25 Dec 2024 20:00:17 +0800
Subject: [PATCH 1/2] [EntryExitInstrumenter][AArch64][RISCV][LoongArch] Pass
 `__builtin_return_address(0)` into `_mcount`

---
 .../Utils/EntryExitInstrumenter.cpp           | 15 ++++++++++++
 .../mcount-with-frompc.ll                     | 24 +++++++++++++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll

diff --git a/llvm/lib/Transforms/Utils/EntryExitInstrumenter.cpp b/llvm/lib/Transforms/Utils/EntryExitInstrumenter.cpp
index 47bb31905d1ac8..5b33edd51cffa6 100644
--- a/llvm/lib/Transforms/Utils/EntryExitInstrumenter.cpp
+++ b/llvm/lib/Transforms/Utils/EntryExitInstrumenter.cpp
@@ -48,6 +48,21 @@ static void insertCall(Function &CurFn, StringRef Func,
                                                   /*isVarArg=*/false)),
           {GV}, "", InsertionPt);
       Call->setDebugLoc(DL);
+    } else if (TargetTriple.isRISCV() || TargetTriple.isAArch64() ||
+               TargetTriple.isLoongArch()) {
+      // On RISC-V, AArch64, and LoongArch, the `_mcount` function takes
+      // `__builtin_return_address(0)` as an argument since
+      // `__builtin_return_address(1)` is not available on these platforms.
+      Instruction *RetAddr = CallInst::Create(
+          Intrinsic::getOrInsertDeclaration(&M, Intrinsic::returnaddress),
+          ConstantInt::get(Type::getInt32Ty(C), 0), "", InsertionPt);
+      RetAddr->setDebugLoc(DL);
+
+      FunctionCallee Fn = M.getOrInsertFunction(
+          Func, FunctionType::get(Type::getVoidTy(C), PointerType::getUnqual(C),
+                                  false));
+      CallInst *Call = CallInst::Create(Fn, RetAddr, "", InsertionPt);
+      Call->setDebugLoc(DL);
     } else {
       FunctionCallee Fn = M.getOrInsertFunction(Func, Type::getVoidTy(C));
       CallInst *Call = CallInst::Create(Fn, "", InsertionPt);
diff --git a/llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll b/llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll
new file mode 100644
index 00000000000000..5123b2270eafa5
--- /dev/null
+++ b/llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll
@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -mtriple=riscv64 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,RISCV64
+; RUN: opt -mtriple=riscv32 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,RISCV32
+; RUN: opt -mtriple=loongarch64 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,LOONGARCH64
+; RUN: opt -mtriple=loongarch32 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,LOONGARCH32
+; RUN: opt -mtriple=aarch64 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64
+; RUN: opt -mtriple=aarch64_be -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64_BE
+; RUN: opt -mtriple=aarch64_32 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64_32
+
+define void @f1() "instrument-function-entry-inlined"="_mcount" {
+; CHECK-LABEL: define void @f1(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    ret void
+;
+  ret void
+}
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; AARCH64: {{.*}}
+; AARCH64_32: {{.*}}
+; AARCH64_BE: {{.*}}
+; LOONGARCH32: {{.*}}
+; LOONGARCH64: {{.*}}
+; RISCV32: {{.*}}
+; RISCV64: {{.*}}

>From 6d49d12616e286f61235813ac51bdf373a97c818 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Mon, 30 Dec 2024 14:19:12 +0800
Subject: [PATCH 2/2] [EntryExitInstrumenter] Fix test. NFC.

---
 .../mcount-with-frompc.ll                     | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll b/llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll
index 5123b2270eafa5..0f8cf5c7354539 100644
--- a/llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll
+++ b/llvm/test/Transforms/EntryExitInstrumenter/mcount-with-frompc.ll
@@ -1,15 +1,16 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -mtriple=riscv64 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,RISCV64
-; RUN: opt -mtriple=riscv32 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,RISCV32
-; RUN: opt -mtriple=loongarch64 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,LOONGARCH64
-; RUN: opt -mtriple=loongarch32 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,LOONGARCH32
-; RUN: opt -mtriple=aarch64 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64
-; RUN: opt -mtriple=aarch64_be -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64_BE
-; RUN: opt -mtriple=aarch64_32 -passes=ee-instrument -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64_32
+; RUN: opt -mtriple=riscv64 -passes="ee-instrument<post-inline>" -S < %s | FileCheck %s --check-prefixes=CHECK,RISCV64
+; RUN: opt -mtriple=riscv32 -passes="ee-instrument<post-inline>" -S < %s | FileCheck %s --check-prefixes=CHECK,RISCV32
+; RUN: opt -mtriple=loongarch64 -passes="ee-instrument<post-inline>" -S < %s | FileCheck %s --check-prefixes=CHECK,LOONGARCH64
+; RUN: opt -mtriple=loongarch32 -passes="ee-instrument<post-inline>" -S < %s | FileCheck %s --check-prefixes=CHECK,LOONGARCH32
+; RUN: opt -mtriple=aarch64 -passes="ee-instrument<post-inline>" -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64
+; RUN: opt -mtriple=aarch64_be -passes="ee-instrument<post-inline>" -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64_BE
+; RUN: opt -mtriple=aarch64_32 -passes="ee-instrument<post-inline>" -S < %s | FileCheck %s --check-prefixes=CHECK,AARCH64_32
 
 define void @f1() "instrument-function-entry-inlined"="_mcount" {
-; CHECK-LABEL: define void @f1(
-; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-LABEL: define void @f1() {
+; CHECK-NEXT:    [[TMP1:%.*]] = call ptr @llvm.returnaddress(i32 0)
+; CHECK-NEXT:    call void @_mcount(ptr [[TMP1]])
 ; CHECK-NEXT:    ret void
 ;
   ret void