[llvm] [ASAN] Make asan pass idempotent. (PR #99439)

[via llvm-commits]

https://github.com/skc7 updated https://github.com/llvm/llvm-project/pull/99439

>From 14d64a36a2e5ce636bf3ff3989ec50c3ec5bef33 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 18 Jul 2024 11:49:24 +0530
Subject: [PATCH] [ASAN] Make asan pass idempotent.

---
 .../Instrumentation/AddressSanitizer.cpp      | 11 ++++
 .../AddressSanitizer/asan-pass-second-run.ll  | 56 +++++++++++++++++++
 2 files changed, 67 insertions(+)
 create mode 100644 llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index adf77f20cb1c7..f2cc63b345154 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -1249,8 +1249,19 @@ AddressSanitizerPass::AddressSanitizerPass(
       UseOdrIndicator(UseOdrIndicator), DestructorKind(DestructorKind),
       ConstructorKind(ConstructorKind) {}
 
+static bool hasAsanModuleCtor(Module &M) {
+  if (M.getFunction(kAsanModuleCtorName))
+    return true;
+  return false;
+}
+
 PreservedAnalyses AddressSanitizerPass::run(Module &M,
                                             ModuleAnalysisManager &MAM) {
+  // Return early if asan.module_ctor is already present in the module.
+  // This implies that asan pass has already run before.
+  if (hasAsanModuleCtor(M))
+    return PreservedAnalyses::all();
+
   ModuleAddressSanitizer ModuleSanitizer(
       M, Options.InsertVersionCheck, Options.CompileKernel, Options.Recover,
       UseGlobalGC, UseOdrIndicator, DestructorKind, ConstructorKind);
diff --git a/llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll b/llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll
new file mode 100644
index 0000000000000..65d72a72f605b
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll
@@ -0,0 +1,56 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 5
+; This test checks in the second run, function is not instrumented again.
+; RUN: opt < %s -passes=asan,asan -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Function with sanitize_address is instrumented.
+; Function Attrs: nounwind uwtable
+;.
+; CHECK: @llvm.used = appending global [1 x ptr] [ptr @asan.module_ctor], section "llvm.metadata"
+; CHECK: @___asan_globals_registered = common hidden global i64 0
+; CHECK: @__start_asan_globals = extern_weak hidden global i64
+; CHECK: @__stop_asan_globals = extern_weak hidden global i64
+; CHECK: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 1, ptr @asan.module_ctor, ptr @asan.module_ctor }]
+;.
+define void @instr_sa(ptr %a) sanitize_address {
+; CHECK-LABEL: define void @instr_sa(
+; CHECK-SAME: ptr [[A:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint ptr [[A]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 [[TMP0]], 3
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[TMP1]], 2147450880
+; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to ptr
+; CHECK-NEXT:    [[TMP4:%.*]] = load i8, ptr [[TMP3]], align 1
+; CHECK-NEXT:    [[TMP5:%.*]] = icmp ne i8 [[TMP4]], 0
+; CHECK-NEXT:    br i1 [[TMP5]], label %[[BB6:.*]], label %[[BB12:.*]], !prof [[PROF0:![0-9]+]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[TMP7:%.*]] = and i64 [[TMP0]], 7
+; CHECK-NEXT:    [[TMP8:%.*]] = add i64 [[TMP7]], 3
+; CHECK-NEXT:    [[TMP9:%.*]] = trunc i64 [[TMP8]] to i8
+; CHECK-NEXT:    [[TMP10:%.*]] = icmp sge i8 [[TMP9]], [[TMP4]]
+; CHECK-NEXT:    br i1 [[TMP10]], label %[[BB11:.*]], label %[[BB12]]
+; CHECK:       [[BB11]]:
+; CHECK-NEXT:    call void @__asan_report_load4(i64 [[TMP0]]) #[[ATTR3:[0-9]+]]
+; CHECK-NEXT:    unreachable
+; CHECK:       [[BB12]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, ptr [[A]], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[TMP1]], 1
+; CHECK-NEXT:    store i32 [[TMP2]], ptr [[A]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+;.
+; CHECK: attributes #[[ATTR0]] = { sanitize_address }
+; CHECK: attributes #[[ATTR1:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+; CHECK: attributes #[[ATTR2:[0-9]+]] = { nounwind }
+; CHECK: attributes #[[ATTR3]] = { nomerge }
+;.
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 1, i32 1048575}
+;.