[Clang][OpenMP] Fix failure with team-wide allocated variable

[doru1004]

Summary

This patch aims to resolve issue: https://github.com/llvm/llvm-project/issues/60345

The following code:

#include <iostream>
#include <omp.h>
#include <stdlib.h>


int main()
{
	int N =1<<30;
	int *a = new int[N];
#pragma omp target data map(tofrom:a[:N])
	{
           #pragma omp target teams distribute parallel for
            for(int i = 0; i < N; i++)
            {
	       int local_a[10];
               #pragma omp allocate(local_a) allocator(omp_pteam_mem_alloc)
	        for(int j = 0; j < 10; j++)
		    local_a[j] = a[(i+j)%N];
	        a[i] = local_a[0];
  }
	}
std::cout << a[0] << "\n";
}

Fails with the following linker errors:

clang-linker-wrapper: error: <unknown>:0: local_a: unsupported initializer for address space

clang-linker-wrapper: error: Errors encountered inside the LTO pipeline.

[nlopes]

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

3355

Please use poison instead of undef wherever possible as we are tying to remove undef. The replacement is usually safe when you just need a placeholder.
Thank you!

[doru1004]

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

3355

I've made the change as requested, this also means that I had to add another check in AMDGPUAsmPrinter.cpp.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

303 ↗

(On Diff #510934)

Isa<UndefValue> covers PoisonValue already

[doru1004]

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

303 ↗

(On Diff #510934)

Perfect! I'll revert this part.

[jhuber6]

clang/test/OpenMP/target_team_variable_codegen.cpp

33

Shouldn't the Nvidia version also be undefined? Not sure why this should vary depending on the target.

[doru1004]

clang/test/OpenMP/target_team_variable_codegen.cpp

33

Perhaps NVIDIA code path can tolerate a zeroinitializer? I don't want to change it if it's not needed. I am basing this check on the code path for AMD GPUs and the initial bug that was reported.

[jdoerfert]

clang/test/OpenMP/target_team_variable_codegen.cpp

33

for AS 3 we should make it always poison.

[jhuber6]

clang/test/OpenMP/target_team_variable_codegen.cpp

33

We should probably change this in HeadToShared in OpenMPOpt as well.

[doru1004]

clang/test/OpenMP/target_team_variable_codegen.cpp

33

Happy to remove the guard and have it always use poison for both NVIDIA and AMD.

[doru1004]

ping

[jhuber6]

LGTM unless anyone else has any concerns.

clang/test/OpenMP/target_team_variable_codegen.cpp

33

These should be a single check line now.

[doru1004]

Commit: 01910787d386584ea5a3d5dc317a908423ba39ed

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

case OMPAllocateDeclAttr::OMPPTeamMemAlloc:
       break;
     case OMPAllocateDeclAttr::OMPLargeCapMemAlloc:
     case OMPAllocateDeclAttr::OMPCGroupMemAlloc:
       break;
     }
     llvm::Type *VarTy = CGF.ConvertTypeForMem(VD->getType());
     auto *GV = new llvm::GlobalVariable(
         CGM.getModule(), VarTy, /*isConstant=*/false,
-        llvm::GlobalValue::InternalLinkage, llvm::Constant::getNullValue(VarTy),
+        llvm::GlobalValue::InternalLinkage,
+        CGM.getTriple().isAMDGCN() ? llvm::PoisonValue::get(VarTy)
+                                   : llvm::Constant::getNullValue(VarTy),
         VD->getName(),
         /*InsertBefore=*/nullptr, llvm::GlobalValue::NotThreadLocal,
         CGM.getContext().getTargetAddressSpace(AS));
     CharUnits Align = CGM.getContext().getDeclAlign(VD);
     GV->setAlignment(Align.getAsAlign());
     return Address(
         CGF.Builder.CreatePointerBitCastOrAddrSpaceCast(
             GV, VarTy->getPointerTo(CGM.getContext().getTargetAddressSpace(

clang/test/OpenMP/target_team_variable_codegen.cpp

+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --check-globals --prefix-filecheck-ir-name _ --global-value-regex "llvm.compiler.used" "_[0-9a-zA-Z]+A[0-9a-zA-Z]+pi[0-9a-zA-Z]+" "_[0-9a-zA-Z]+anotherPi" --version 2
+// REQUIRES: amdgpu-registered-target
+
+
+// Test target codegen - host bc file has to be created first.
+// RUN: %clang_cc1 -verify -fopenmp -x c++ -triple powerpc64le-unknown-unknown -fopenmp-targets=amdgcn-amd-amdhsa -emit-llvm-bc %s -o %t-ppc-host-amd.bc
+// RUN: %clang_cc1 -verify -fopenmp -x c++ -triple amdgcn-amd-amdhsa -fopenmp-targets=amdgcn-amd-amdhsa -emit-llvm %s -fopenmp-target-debug -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host-amd.bc -o - | FileCheck %s --check-prefix=CHECK-AMD
+
+// RUN: %clang_cc1 -verify -fopenmp -x c++ -triple powerpc64le-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm-bc %s -o %t-ppc-host-nvidia.bc
+// RUN: %clang_cc1 -verify -fopenmp -x c++ -triple nvptx64-unknown-unknown -fopenmp-targets=nvptx64-unknown-unknown -emit-llvm %s -fopenmp-target-debug -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host-nvidia.bc -o - | FileCheck %s --check-prefix=CHECK-NVIDIA
+
+// expected-no-diagnostics
+
+#ifndef HEADER
+#define HEADER
+
+typedef enum omp_allocator_handle_t {
+  omp_null_allocator = 0,
+  omp_default_mem_alloc = 1,
+  omp_large_cap_mem_alloc = 2,
+  omp_const_mem_alloc = 3,
+  omp_high_bw_mem_alloc = 4,
+  omp_low_lat_mem_alloc = 5,
+  omp_cgroup_mem_alloc = 6,
+  omp_pteam_mem_alloc = 7,
+  omp_thread_mem_alloc = 8,
+  KMP_ALLOCATOR_MAX_HANDLE = __UINTPTR_MAX__
+} omp_allocator_handle_t;
+
+//.
+// CHECK-AMD: @local_a = internal addrspace(3) global [10 x i32] poison, align 4
+//.
+// CHECK-NVIDIA: @local_a = internal addrspace(3) global [10 x i32] zeroinitializer, align 4
+//.
+int main()
+{
+  int N = 10000;
+  int *a = new int[N];
+#pragma omp target data map(tofrom:a[:N])
+  {
+#pragma omp target teams distribute parallel for
+    for(int i = 0; i < N; i++)
+    {
+      int local_a[10];
+#pragma omp allocate(local_a) allocator(omp_pteam_mem_alloc)
+      for(int j = 0; j < 10; j++)
+        local_a[j] = a[(i + j) % N];
+      a[i] = local_a[0];
+    }
+  }
+  return a[17];
+}
+
+#endif
+//// NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+// CHECK-AMD: {{.*}}
+// CHECK-NVIDIA: {{.*}}