[llvm] [WIP] Extend data layout to add non zero null value for address space. (PR #83109)

[Rana Pratap Reddy via llvm-commits]

https://github.com/ranapratap55 updated https://github.com/llvm/llvm-project/pull/83109

>From 67dea8f4f665fa2cbd944674c5d83eae6a0199fd Mon Sep 17 00:00:00 2001
From: Rana Pratap Reddy N <RanaPratapReddy.Nimmakayala at amd.com>
Date: Tue, 27 Feb 2024 13:07:00 +0530
Subject: [PATCH] [WIP] Extend data layout to add non zero null value for
 address space

---
 llvm/include/llvm/IR/DataLayout.h             | 14 +++++++++++++
 llvm/lib/IR/DataLayout.cpp                    | 21 +++++++++++++++++++
 .../AMDGPU/datalayout-non-zero-lds-value.ll   | 20 ++++++++++++++++++
 3 files changed, 55 insertions(+)
 create mode 100644 llvm/test/CodeGen/AMDGPU/datalayout-non-zero-lds-value.ll

diff --git a/llvm/include/llvm/IR/DataLayout.h b/llvm/include/llvm/IR/DataLayout.h
index 71f7f51d8ee431..75fd19efd7a26b 100644
--- a/llvm/include/llvm/IR/DataLayout.h
+++ b/llvm/include/llvm/IR/DataLayout.h
@@ -24,6 +24,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Alignment.h"
@@ -164,6 +165,8 @@ class DataLayout {
   /// well-defined bitwise representation.
   SmallVector<unsigned, 8> NonIntegralAddressSpaces;
 
+  DenseMap<uint64_t, uint64_t> AddrSpaceToNonZeroValueMap;
+
   /// Attempts to set the alignment of the given type. Returns an error
   /// description on failure.
   Error setAlignment(AlignTypeEnum AlignType, Align ABIAlign, Align PrefAlign,
@@ -299,6 +302,17 @@ class DataLayout {
     return ManglingMode == MM_WinCOFFX86;
   }
 
+  uint64_t getNonZeroValueForAddrSpace(uint64_t AddrSpace) {
+  auto It = AddrSpaceToNonZeroValueMap.find(AddrSpace);
+  if (It == AddrSpaceToNonZeroValueMap.end())
+    return 0x000000000;
+  return AddrSpaceToNonZeroValueMap[AddrSpace];
+  }
+
+  void setNonZeroValueForAddrSpace(uint64_t AddrSpace, uint64_t Value) {
+    AddrSpaceToNonZeroValueMap[AddrSpace] = Value;
+  }
+
   /// Returns true if symbols with leading question marks should not receive IR
   /// mangling. True for Windows mangling modes.
   bool doNotMangleLeadingQuestionMark() const {
diff --git a/llvm/lib/IR/DataLayout.cpp b/llvm/lib/IR/DataLayout.cpp
index a2f5714c706874..591b0fda665387 100644
--- a/llvm/lib/IR/DataLayout.cpp
+++ b/llvm/lib/IR/DataLayout.cpp
@@ -502,6 +502,27 @@ Error DataLayout::parseSpecifier(StringRef Desc) {
         return Err;
       break;
     }
+    case 'z': {
+      uint64_t AddrSpace = 0;
+      if (!Tok.empty())
+        if (Error Err = getInt(Tok, AddrSpace))
+          return Err;
+      if (!isUInt<24>(AddrSpace))
+        return reportError("Invalid address space, must be a 24-bit integer");
+
+      if (Rest.empty())
+        return reportError(
+            "Missing address space specification for pointer in datalayout string");
+      if (Error Err = ::split(Rest, ':', Split))
+        return Err;
+      uint64_t Value;
+      if (Error Err = getInt(Tok, Value))
+        return Err;
+      
+      setNonZeroValueForAddrSpace(AddrSpace, Value);
+
+      break;
+    }
     case 'G': { // Default address space for global variables.
       if (Error Err = getAddrSpace(Tok, DefaultGlobalsAddrSpace))
         return Err;
diff --git a/llvm/test/CodeGen/AMDGPU/datalayout-non-zero-lds-value.ll b/llvm/test/CodeGen/AMDGPU/datalayout-non-zero-lds-value.ll
new file mode 100644
index 00000000000000..2cea2ae28d0cce
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/datalayout-non-zero-lds-value.ll
@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -S -mtriple=amdgcn-- | FileCheck %s
+
+; CHECK: 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-z3:1-S32-A5-G1-ni:7:8:9"
+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-z3:1-S32-A5-G1-ni:7:8:9"
+ at lds = addrspace(3) global [8 x i32] [i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8]
+
+define amdgpu_kernel void @load_init_lds_global(ptr addrspace(1) %out, i1 %p) {
+; CHECK-LABEL: define amdgpu_kernel void @load_init_lds_global(
+; CHECK-SAME: ptr addrspace(1) [[OUT:%.*]], i1 [[P:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [8 x i32], ptr addrspace(3) @lds, i32 0, i32 10
+; CHECK-NEXT:    [[LD:%.*]] = load i32, ptr addrspace(3) [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], ptr addrspace(1) [[OUT]], align 4
+; CHECK-NEXT:    ret void
+;
+  %gep = getelementptr [8 x i32], ptr addrspace(3) @lds, i32 0, i32 10
+  %ld = load i32, ptr addrspace(3) %gep
+  store i32 %ld, ptr addrspace(1) %out
+  ret void
+}