[SPIRV][SPIRVPrepareGlobals] Map AMD's dynamic LDS 0-element globals to arrays with UINT32_MAX elements

[jmmartinez]

In HIP, dynamic LDS variables are represented using 0-element global arrays in the __shared__ language address-space.

  extern __shared__ int LDS[];

These are not representable in SPIRV directly.

To represent them, for AMD, we use an array with UINT32_MAX-elements. These are reverse translated to 0-element arrays later in AMD's SPIRV runtime pipeline (in SPIRVReader.cpp).

Stacked over #166950

[llvmbot]

@llvm/pr-subscribers-backend-spir-v

Author: Juan Manuel Martinez Caamaño (jmmartinez)

Changes

In HIP, dynamic LDS variables are represented using 0-element global arrays in the __shared__ language address-space.

  extern __shared__ int LDS[];

These are not representable in SPIRV directly.

To represent them, for AMD, we use an array with UINT32_MAX-elements. These are reverse translated to 0-element arrays later in AMD's SPIRV runtime pipeline (in SPIRVReader.cpp).

Stacked over #166950

---

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

2 Files Affected:

• (modified) llvm/lib/Target/SPIRV/SPIRVPrepareGlobals.cpp (+27)
• (added) llvm/test/CodeGen/SPIRV/hip_dyn_lds.ll (+20)

diff --git a/llvm/lib/Target/SPIRV/SPIRVPrepareGlobals.cpp b/llvm/lib/Target/SPIRV/SPIRVPrepareGlobals.cpp
index c44c53129f1e0..42a9577bb2054 100644
--- a/llvm/lib/Target/SPIRV/SPIRVPrepareGlobals.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVPrepareGlobals.cpp
@@ -13,6 +13,7 @@
 
 #include "SPIRV.h"
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Module.h"
 
 using namespace llvm;
@@ -43,6 +44,29 @@ bool tryExtendLLVMBitcodeMarker(GlobalVariable &Bitcode) {
   return true;
 }
 
+bool tryExtendDynamicLDSGlobal(GlobalVariable &GV) {
+  constexpr unsigned WorkgroupAS = 3;
+  const bool IsWorkgroupExternal =
+      GV.hasExternalLinkage() && GV.getAddressSpace() == WorkgroupAS;
+  if (!IsWorkgroupExternal)
+    return false;
+
+  const ArrayType *AT = dyn_cast<ArrayType>(GV.getValueType());
+  if (!AT || AT->getNumElements() != 0)
+    return false;
+
+  constexpr auto Magic = std::numeric_limits<uint32_t>::max();
+  ArrayType *NewAT = ArrayType::get(AT->getElementType(), Magic);
+  GlobalVariable *NewGV = new GlobalVariable(
+      *GV.getParent(), NewAT, GV.isConstant(), GV.getLinkage(), nullptr, "",
+      &GV, GV.getThreadLocalMode(), WorkgroupAS, GV.isExternallyInitialized());
+  NewGV->takeName(&GV);
+  GV.replaceAllUsesWith(NewGV);
+  GV.eraseFromParent();
+
+  return true;
+}
+
 bool SPIRVPrepareGlobals::runOnModule(Module &M) {
   const bool IsAMD = M.getTargetTriple().getVendor() == Triple::AMD;
   if (!IsAMD)
@@ -52,6 +76,9 @@ bool SPIRVPrepareGlobals::runOnModule(Module &M) {
   if (GlobalVariable *Bitcode = M.getNamedGlobal("llvm.embedded.module"))
     Changed |= tryExtendLLVMBitcodeMarker(*Bitcode);
 
+  for (GlobalVariable &GV : make_early_inc_range(M.globals()))
+    Changed |= tryExtendDynamicLDSGlobal(GV);
+
   return Changed;
 }
 char SPIRVPrepareGlobals::ID = 0;
diff --git a/llvm/test/CodeGen/SPIRV/hip_dyn_lds.ll b/llvm/test/CodeGen/SPIRV/hip_dyn_lds.ll
new file mode 100644
index 0000000000000..f0acfdfdede9d
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/hip_dyn_lds.ll
@@ -0,0 +1,20 @@
+; RUN: llc -verify-machineinstrs -mtriple=spirv64-amd-amdhsa %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -mtriple=spirv64-amd-amdhsa %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpName %[[#LDS:]] "lds"
+; CHECK: OpDecorate %[[#LDS]] LinkageAttributes "lds" Import
+; CHECK: %[[#UINT:]] = OpTypeInt 32 0
+; CHECK: %[[#UINT_MAX:]] = OpConstant %[[#UINT]] 4294967295
+; CHECK: %[[#LDS_ARR_TY:]] = OpTypeArray %[[#UINT]] %[[#UINT_MAX]]
+; CHECK: %[[#LDS_ARR_PTR_WG:]] = OpTypePointer Workgroup %[[#LDS_ARR_TY]]
+; CHECK: %[[#LDS]] = OpVariable %[[#LDS_ARR_PTR_WG]] Workgroup
+
+@lds = external addrspace(3) global [0 x i32]
+
+define spir_kernel void @foo(ptr addrspace(4) %in, ptr addrspace(4) %out) {
+entry:
+  %val = load i32, ptr addrspace(4) %in
+  %add = add i32 %val, 1
+  store i32 %add, ptr addrspace(4) %out
+  ret void
+}

[maarquitos14]

LGTM, just one nit.

[s-perron]

I'm wondering if you need to change all 0-sized array or not. If so, we might want to centralize the conversion of 0-sized arrays. We could try to move the code that changes them to 1 element arrays here as well.

I can look into changes HLSL use of zero-sized arrays if needed.

[s-perron]

What do you want to do with 0-sized arrays that are not the type of the global value? Is even possible to do that? Comments explaining why you limit this to just the type of the GV would be useful.

[jmmartinez]

Sorry I'm not sure I understood the question.

This condition is matching globals that have an array type with 0 elements. If the global does not have an array type or if it is an array type with a size different from 0 this function returns false.

Maybe there is a mix from using getValueType vs getType. The first gives the type of the initializer of the global (the value stored in the global), the second returns the type of the global when used as a value in the llvm-ir (a pointer type).

[s-perron]

Sorry, I did not write that properly.

Could you have, say, a global whose type is a struct containing a 0-sized array? What do you want to do in case?

@lds = external addrspace(3) global {i32, [0 x i32]}

[jmmartinez]

I don't think it's possible. These cases get rejected by the frontend normally: https://godbolt.org/z/8PdqqrYMT

[jmmartinez]

I'm wondering if you need to change all 0-sized array or not. If so, we might want to centralize the conversion of 0-sized arrays. We could try to move the code that changes them to 1 element arrays here as well.

If you point me to where this is happening I can try to do that too as a follow-up patch.

[s-perron]

I'm wondering if you need to change all 0-sized array or not. If so, we might want to centralize the conversion of 0-sized arrays. We could try to move the code that changes them to 1 element arrays here as well.

If you point me to where this is happening I can try to do that too as a follow-up patch.

I was wrong. I was thinking for #149522, but that does not changes the global variables, it just change the type on the GEP. They seem different enough that they may not be reasonable to merge into a single place.