[llvm] [AMDGPU] Propagate amdgpu-agpr-alloc to ASan runtime calls in SwLowerLDS (PR #195915)

[Pablo Reble via llvm-commits]

https://github.com/reble updated https://github.com/llvm/llvm-project/pull/195915

>From 016dd6ad050238f5c2b08a4bb9a9e166c57cefe4 Mon Sep 17 00:00:00 2001
From: Pablo Reble <pablo.reble at amd.com>
Date: Tue, 5 May 2026 12:33:33 -0500
Subject: [PATCH 1/3] [AMDGPU] Propagate amdgpu-agpr-alloc to ASan runtime
 calls in SwLowerLDS

---
 llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index ccf6e328b4cbf..04d894fce010c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -210,6 +210,7 @@ class AMDGPUSwLowerLDS {
                                 Value *HiddenDynLDSSize,
                                 SetVector<GlobalVariable *> &DynamicLDSGlobals);
   void initAsanInfo();
+  void propagateAsanAttrs(Function *KernelFunc, FunctionCallee& Callee);
 
 private:
   Module &M;
@@ -744,6 +745,19 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
   }
 }
 
+void AMDGPUSwLowerLDS::propagateAsanAttrs(Function *KernelFunc,
+                                          FunctionCallee& Callee) {
+  auto *F = dyn_cast<Function>(Callee.getCallee());
+  if (!F)
+    return;
+  const GCNSubtarget &ST = AMDGPUTM.getSubtarget<GCNSubtarget>(*KernelFunc);
+  if (!ST.hasGFX90AInsts())
+    return;
+  Attribute AGPRAlloc = KernelFunc->getFnAttribute("amdgpu-agpr-alloc");
+  if (AGPRAlloc.isValid())
+    F->addFnAttr(AGPRAlloc);
+}
+
 void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
   auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
   Type *Int64Ty = IRB.getInt64Ty();
@@ -751,6 +765,7 @@ void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
   FunctionCallee AsanPoisonRegion = M.getOrInsertFunction(
       "__asan_poison_region",
       FunctionType::get(VoidTy, {Int64Ty, Int64Ty}, false));
+  propagateAsanAttrs(Func, AsanPoisonRegion);
 
   auto RedzonesVec = LDSParams.RedzoneOffsetAndSizeVector;
   size_t VecSize = RedzonesVec.size();
@@ -887,6 +902,7 @@ void AMDGPUSwLowerLDS::lowerKernelLDSAccesses(Function *Func,
   FunctionCallee MallocFunc = M.getOrInsertFunction(
       StringRef("__asan_malloc_impl"),
       FunctionType::get(Int64Ty, {Int64Ty, Int64Ty}, false));
+  propagateAsanAttrs(Func, MallocFunc);
   Value *RAPtrToInt = IRB.CreatePtrToInt(ReturnAddress, Int64Ty);
   Value *MallocCall = IRB.CreateCall(MallocFunc, {CurrMallocSize, RAPtrToInt});
 
@@ -948,6 +964,7 @@ void AMDGPUSwLowerLDS::lowerKernelLDSAccesses(Function *Func,
   FunctionCallee AsanFreeFunc = M.getOrInsertFunction(
       StringRef("__asan_free_impl"),
       FunctionType::get(IRB.getVoidTy(), {Int64Ty, Int64Ty}, false));
+  propagateAsanAttrs(Func, AsanFreeFunc);
   Value *ReturnAddr = IRB.CreateIntrinsic(
       Intrinsic::returnaddress, IRB.getPtrTy(DL.getProgramAddressSpace()),
       IRB.getInt32(0));

>From 8619de2566c1455d3e4446bfc6208d1bc52d23f3 Mon Sep 17 00:00:00 2001
From: Pablo Reble <pablo.reble at amd.com>
Date: Tue, 5 May 2026 14:30:14 -0500
Subject: [PATCH 2/3] fix formatting

---
 llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 04d894fce010c..810deee5f874f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -210,7 +210,7 @@ class AMDGPUSwLowerLDS {
                                 Value *HiddenDynLDSSize,
                                 SetVector<GlobalVariable *> &DynamicLDSGlobals);
   void initAsanInfo();
-  void propagateAsanAttrs(Function *KernelFunc, FunctionCallee& Callee);
+  void propagateAsanAttrs(Function *KernelFunc, FunctionCallee &Callee);
 
 private:
   Module &M;
@@ -746,7 +746,7 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
 }
 
 void AMDGPUSwLowerLDS::propagateAsanAttrs(Function *KernelFunc,
-                                          FunctionCallee& Callee) {
+                                          FunctionCallee &Callee) {
   auto *F = dyn_cast<Function>(Callee.getCallee());
   if (!F)
     return;

>From bd62b1fcf623bb633f017960c4de8ed7915273b6 Mon Sep 17 00:00:00 2001
From: Pablo Reble <pablo.reble at amd.com>
Date: Tue, 5 May 2026 22:08:22 -0500
Subject: [PATCH 3/3] Adding test

---
 .../amdgpu-sw-lower-lds-agpr-alloc-asan.ll    | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-agpr-alloc-asan.ll

diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-agpr-alloc-asan.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-agpr-alloc-asan.ll
new file mode 100644
index 0000000000000..d222d1baa0a3d
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-agpr-alloc-asan.ll
@@ -0,0 +1,30 @@
+; RUN: opt < %s -passes=amdgpu-sw-lower-lds -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a \
+; RUN:   | FileCheck %s --check-prefix=GFX90A
+; RUN: opt < %s -passes=amdgpu-sw-lower-lds -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 \
+; RUN:   | FileCheck %s --check-prefix=NOGFX90A
+
+; Test that amdgpu-agpr-alloc is propagated to __asan_malloc_impl,
+; __asan_free_impl, and __asan_poison_region only on gfx90a targets.
+
+ at lds_1 = internal addrspace(3) global i32 poison, align 4
+
+define amdgpu_kernel void @k0() sanitize_address #0 {
+  store i32 1, ptr addrspace(3) @lds_1, align 4
+  ret void
+}
+
+; GFX90A: declare i64 @__asan_malloc_impl(i64, i64) #[[AGPR_ATTR:[0-9]+]]
+; GFX90A: declare void @__asan_poison_region(i64, i64) #[[AGPR_ATTR]]
+; GFX90A: declare void @__asan_free_impl(i64, i64) #[[AGPR_ATTR]]
+
+; GFX90A: attributes #[[AGPR_ATTR]] = { "amdgpu-agpr-alloc"="0" }
+
+; NOGFX90A-NOT: declare i64 @__asan_malloc_impl{{.*}}"amdgpu-agpr-alloc"
+; NOGFX90A-NOT: declare void @__asan_poison_region{{.*}}"amdgpu-agpr-alloc"
+; NOGFX90A-NOT: declare void @__asan_free_impl{{.*}}"amdgpu-agpr-alloc"
+
+attributes #0 = { sanitize_address "amdgpu-agpr-alloc"="0" }
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"nosanitize_address", i32 1}
+