[llvm] [AMDGPU][ASAN] Remove amdgpu-no-flat-scratch-init attribute from lowered kernels in amdgpu-sw-lower-lds pass (PR #160541)

[via llvm-commits]

https://github.com/skc7 created https://github.com/llvm/llvm-project/pull/160541

LLVM IR of __asan_malloc_impl has generic pointers accessing private memory. But this call to __asan_malloc_impl is introduces late in the pass pipeline by "amdgpu-sw-lower-lds" pass. So, to make sure Flat Addressing for Scratch Memory works, this PR removes amdgpu-no-flat-scratch-init from the lowered kernels.

>From fba5da61866d5557e24302aa9b5e537d64e92585 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 24 Sep 2025 20:24:18 +0530
Subject: [PATCH] [AMDGPU][ASAN] Remove  mdgpu-no-flat-scratch-init attribute
 from lowered kernels in amdgpu-sw-lower-lds pass

---
 llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp    |  3 ++-
 .../amdgpu-sw-lower-lds-no-scratch-init.ll     | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-no-scratch-init.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 4a9437b37aa39..d610fbf9bb06b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -1213,7 +1213,8 @@ bool AMDGPUSwLowerLDS::run() {
       removeFnAttrFromReachable(
           CG, Func,
           {"amdgpu-no-workitem-id-x", "amdgpu-no-workitem-id-y",
-           "amdgpu-no-workitem-id-z", "amdgpu-no-heap-ptr"});
+           "amdgpu-no-workitem-id-z", "amdgpu-no-heap-ptr",
+           "amdgpu-no-flat-scratch-init"});
       if (!LDSParams.IndirectAccess.StaticLDSGlobals.empty() ||
           !LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
         removeFnAttrFromReachable(CG, Func, {"amdgpu-no-lds-kernel-id"});
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-no-scratch-init.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-no-scratch-init.ll
new file mode 100644
index 0000000000000..660b0f32590ba
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-no-scratch-init.ll
@@ -0,0 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 4
+; RUN: opt < %s -passes=amdgpu-sw-lower-lds -S -amdgpu-asan-instrument-lds=false -mtriple=amdgcn-amd-amdhsa | FileCheck %s
+; RUN: llc < %s -enable-new-pm -stop-after=amdgpu-sw-lower-lds -amdgpu-asan-instrument-lds=false -mtriple=amdgcn-amd-amdhsa | FileCheck %s
+
+; Test to check if amdgpu-no-flat-scratch-init is removed from the kernel.
+ at lds_1 = internal addrspace(3) global [1 x i8] poison, align 4
+
+define amdgpu_kernel void @k0() #0 {
+  store i8 7, ptr addrspace(3) @lds_1, align 4
+  ret void
+}
+
+attributes #0 = { sanitize_address "amdgpu-no-flat-scratch-init"}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"nosanitize_address", i32 1}
+
+; CHECK: attributes #[[ATTR0:.*]] = { sanitize_address "amdgpu-lds-size"="8" }