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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Pablo Reble (reble)

<details>
<summary>Changes</summary>

When the software LDS lowering pass inserts calls to `__asan_malloc_impl`,
`__asan_free_imp`, and `__asan_poison_region`, the `amdgpu-agpr-alloc`
attribute from the kernel function was not being forwarded to those
callees. On gfx90a this attribute controls AGPR allocation and must be
consistent across the call graph to avoid incorrect register allocation.

Adding a helper `propagateAsanAttrs` that copies `amdgpu-agpr-alloc` from the
kernel to each inserted ASan callee, gated on `hasGFX90AInsts()` since the
attribute is only meaningful on that subtarget.

Co-Authored-By: Claude Sonnet 4.6 <noreply@<!-- -->anthropic.com>

---
Full diff: https://github.com/llvm/llvm-project/pull/195915.diff


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp (+17) 
- (added) llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-agpr-alloc-asan.ll (+30) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index ccf6e328b4cbf..810deee5f874f 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));
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}
+

``````````

</details>


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