[mlir] [clang-tools-extra] [llvm] [clang] [ASan][AMDGPU] Fix Assertion Failure. (PR #78242)

[via cfe-commits]

https://github.com/ampandey-1995 updated https://github.com/llvm/llvm-project/pull/78242

>From fe6d4abebb12e063e10b2266f76ff1a604783c0c Mon Sep 17 00:00:00 2001
From: Amit Pandey <amit.pandey at amd.com>
Date: Mon, 11 Dec 2023 14:23:44 +0530
Subject: [PATCH] [ASan][AMDGPU] Fix Assertion Failure.

Assertion failure `(i >= FTy->getNumParams() || FTy->getParamType(i) ==
Args[i]->getType()) && "Calling a function with a bad signature!"'.  The
'llvm.memcpy' intercepted by ASan instrumentation pass is implemented by it's
own __asan_memcpy implementation. The second argument of llvm.memcpy accepts ptr
to addrspace(4), __asan_memcpy also has to follow ptr to addrspace(4)
convention.
---
 .../Instrumentation/AddressSanitizer.cpp      |  8 ++--
 .../AMDGPU/asan_instrument_mem_intrinsics.ll  | 48 +++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Instrumentation/AddressSanitizer/AMDGPU/asan_instrument_mem_intrinsics.ll

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index fb5838bb7941ad..32e7f1b6fc3374 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -1254,9 +1254,11 @@ Value *AddressSanitizer::memToShadow(Value *Shadow, IRBuilder<> &IRB) {
 void AddressSanitizer::instrumentMemIntrinsic(MemIntrinsic *MI) {
   InstrumentationIRBuilder IRB(MI);
   if (isa<MemTransferInst>(MI)) {
-    IRB.CreateCall(isa<MemMoveInst>(MI) ? AsanMemmove : AsanMemcpy,
-                   {MI->getOperand(0), MI->getOperand(1),
-                    IRB.CreateIntCast(MI->getOperand(2), IntptrTy, false)});
+    IRB.CreateCall(
+        isa<MemMoveInst>(MI) ? AsanMemmove : AsanMemcpy,
+        {MI->getOperand(0),
+         IRB.CreatePointerBitCastOrAddrSpaceCast(MI->getOperand(1), PtrTy),
+         IRB.CreateIntCast(MI->getOperand(2), IntptrTy, false)});
   } else if (isa<MemSetInst>(MI)) {
     IRB.CreateCall(
         AsanMemset,
diff --git a/llvm/test/Instrumentation/AddressSanitizer/AMDGPU/asan_instrument_mem_intrinsics.ll b/llvm/test/Instrumentation/AddressSanitizer/AMDGPU/asan_instrument_mem_intrinsics.ll
new file mode 100644
index 00000000000000..26258ef384795d
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/AMDGPU/asan_instrument_mem_intrinsics.ll
@@ -0,0 +1,48 @@
+;RUN: opt < %s -passes=asan -S | FileCheck %s
+
+target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"
+target triple = "amdgcn-amd-amdhsa"
+
+ at __const.__assert_fail.fmt = private unnamed_addr addrspace(4) constant [47 x i8] c"%s:%u: %s: Device-side assertion `%s' failed.\0A\00", align 16
+
+; Function Attrs: nocallback nofree nounwind willreturn memory(argmem: readwrite)
+declare void @llvm.memcpy.p0.p4.i64(ptr noalias nocapture writeonly, ptr addrspace(4) noalias nocapture readonly, i64, i1 immarg) #1
+
+; Function Attrs: convergent mustprogress noinline nounwind optnone
+define weak hidden void @test_mem_intrinsic() sanitize_address #0 {
+; CHECK: define weak hidden void @test_mem_intrinsic() #1 {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[FMT:%.*]] = alloca [47 x i8], align 16, addrspace(5)
+; CHECK-NEXT: [[FADDRC:%.*]] = addrspacecast ptr addrspace(5) [[FMT]] to ptr
+; CHECK-NEXT: [[ITMP:%.*]] = call ptr @__asan_memcpy(ptr [[FADDRC]], ptr addrspacecast (ptr addrspace(4) @__const.__assert_fail.fmt to ptr), i64 47)
+; CHECK-NEXT: ret
+entry:
+%fmt = alloca [47 x i8], align 16, addrspace(5)
+%fmt.ascast = addrspacecast ptr addrspace(5) %fmt to ptr
+call void @llvm.memcpy.p0.p4.i64(ptr align 16 %fmt.ascast, ptr addrspace(4) align 16 @__const.__assert_fail.fmt, i64 47, i1 false)
+ret void
+}
+
+
+
+attributes #0 = { convergent mustprogress noinline nounwind optnone "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="gfx906" "target-features"="+16-bit-insts,+ci-insts,+dl-insts,+dot1-insts,+dot10-insts,+dot2-insts,+dot7-insts,+dpp,+gfx8-insts,+gfx9-insts,+s-memrealtime,+s-memtime-inst,+wavefrontsize64,+xnack" }
+attributes #1 = { nocallback nofree nounwind willreturn memory(argmem: readwrite) }
+attributes #2 = { convergent nounwind "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="gfx906" "target-features"="+16-bit-insts,+ci-insts,+dl-insts,+dot1-insts,+dot10-insts,+dot2-insts,+dot7-insts,+dpp,+gfx8-insts,+gfx9-insts,+s-memrealtime,+s-memtime-inst,+wavefrontsize64,+xnack" }
+attributes #3 = { cold noreturn nounwind memory(inaccessiblemem: write) }
+attributes #4 = { convergent mustprogress noinline norecurse nounwind optnone "amdgpu-flat-work-group-size"="1,1024" "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="gfx906" "target-features"="+16-bit-insts,+ci-insts,+dl-insts,+dot1-insts,+dot10-insts,+dot2-insts,+dot7-insts,+dpp,+gfx8-insts,+gfx9-insts,+s-memrealtime,+s-memtime-inst,+wavefrontsize64,+xnack" "uniform-work-group-size"="true" }
+attributes #5 = { convergent nounwind }
+
+!llvm.module.flags = !{!0, !1, !2, !3, !4}
+!llvm.ident = !{!5}
+
+!0 = !{i32 1, !"amdgpu_code_object_version", i32 400}
+!1 = !{i32 1, !"amdgpu_printf_kind", !"hostcall"}
+!2 = !{i32 1, !"wchar_size", i32 4}
+!3 = !{i32 8, !"PIC Level", i32 2}
+!4 = !{i32 7, !"frame-pointer", i32 2}
+!5 = !{!"clang version 18.0.0git (git at github.com:llvm/llvm-project.git a919fe73e2c4c26fbb2a6f0459ff6564a5c83943)"}
+!6 = distinct !{!6, !7}
+!7 = !{!"llvm.loop.mustprogress"}
+!8 = distinct !{!8, !7}
+!9 = distinct !{!9, !7}
+!10 = distinct !{!10, !7}