[OpenMP] Use the align clause value from 'omp allocate' for globals

[mikerice]

Summary

Refactor the code that handles the align clause of 'omp allocate' so
it can be used with globals as well as local variables.

[ABataev]

LG

clang/lib/CodeGen/CGDecl.cpp

llvm::Constant *NewPtrForOldDecl =
             Entry, DummyGV->getType());
 
     // Now we have a casted version of the changed global, the dummy can be
     // replaced and deleted.
     DummyGV->replaceAllUsesWith(NewPtrForOldDecl);
     DummyGV->eraseFromParent();
   }
 }
+
+llvm::Optional<CharUnits>
+CodeGenModule::getOMPAllocateAlignment(const VarDecl *VD) {
+  if (const auto *AA = VD->getAttr<OMPAllocateDeclAttr>()) {
+    if (Expr *Alignment = AA->getAlignment()) {
+      unsigned UserAlign =
+          Alignment->EvaluateKnownConstInt(getContext()).getExtValue();
+      CharUnits NaturalAlign =
+          getNaturalTypeAlignment(VD->getType().getNonReferenceType());
+
+      // OpenMP5.1 pg 185 lines 7-10
+      //   Each item in the align modifier list must be aligned to the maximum
+      //   of the specified alignment and the type's natural alignment.
+      return CharUnits::fromQuantity(
+          std::max<unsigned>(UserAlign, NaturalAlign.getQuantity()));
+    }
+  }
+  return llvm::None;
+}

clang/lib/CodeGen/CGOpenMPRuntime.cpp

static llvm::Value *getAllocatorVal(CodeGenFunction &CGF,
   } else {
     // If no allocator specified, it defaults to the null allocator.
     AllocVal = llvm::Constant::getNullValue(
         CGF.CGM.getTypes().ConvertType(CGF.getContext().VoidPtrTy));
   }
   return AllocVal;
 }
 
-/// Given the allocate directive list item type and align clause value,
-/// return appropriate alignment.
-static llvm::Value *getAlignmentValue(CodeGenFunction &CGF, QualType ListItemTy,
-                                      const Expr *Alignment) {
-  if (!Alignment)
+/// Return the alignment from an allocate directive if present.
+static llvm::Value *getAlignmentValue(CodeGenModule &CGM, const VarDecl *VD) {
+  llvm::Optional<CharUnits> AllocateAlignment = CGM.getOMPAllocateAlignment(VD);
+
+  if (!AllocateAlignment)
     return nullptr;
 
-  unsigned UserAlign =
-      Alignment->EvaluateKnownConstInt(CGF.getContext()).getExtValue();
-  CharUnits NaturalAlign = CGF.CGM.getNaturalTypeAlignment(ListItemTy);
-
-  // OpenMP5.1 pg 185 lines 7-10
-  //   Each item in the align modifier list must be aligned to the maximum
-  //   of the specified alignment and the type's natural alignment.
-  //
-  // If no alignment specified then use the natural alignment.
-  return llvm::ConstantInt::get(
-      CGF.CGM.SizeTy,
-      std::max<unsigned>(UserAlign, NaturalAlign.getQuantity()));
+  return llvm::ConstantInt::get(CGM.SizeTy,
+                                AllocateAlignment.getValue().getQuantity());
 }
 
 Address CGOpenMPRuntime::getAddressOfLocalVariable(CodeGenFunction &CGF,
                                                    const VarDecl *VD) {
   if (!VD)
     return Address::invalid();
   Address UntiedAddr = Address::invalid();
   Address UntiedRealAddr = Address::invalid();
if (CVD->hasAttr<OMPAllocateDeclAttr>()) {
     } else {
       CharUnits Sz = CGM.getContext().getTypeSizeInChars(CVD->getType());
       Size = CGM.getSize(Sz.alignTo(Align));
     }
     llvm::Value *ThreadID = getThreadID(CGF, CVD->getBeginLoc());
     const auto *AA = CVD->getAttr<OMPAllocateDeclAttr>();
     const Expr *Allocator = AA->getAllocator();
     llvm::Value *AllocVal = getAllocatorVal(CGF, Allocator);
-    llvm::Value *Alignment = getAlignmentValue(
-        CGF, VD->getType().getNonReferenceType(), AA->getAlignment());
+    llvm::Value *Alignment = getAlignmentValue(CGM, CVD);
     SmallVector<llvm::Value *, 4> Args;
     Args.push_back(ThreadID);
     if (Alignment)
       Args.push_back(Alignment);
     Args.push_back(Size);
     Args.push_back(AllocVal);
     llvm::omp::RuntimeFunction FnID =
         Alignment ? OMPRTL___kmpc_aligned_alloc : OMPRTL___kmpc_alloc;

clang/lib/CodeGen/CodeGenModule.h

public:
   /// Emit a code for requires directive.
   /// \param D Requires declaration
   void EmitOMPRequiresDecl(const OMPRequiresDecl *D);
 
   /// Emit a code for the allocate directive.
   /// \param D The allocate declaration
   void EmitOMPAllocateDecl(const OMPAllocateDecl *D);
 
+  /// Return the alignment specified in an allocate directive, if present.
+  llvm::Optional<CharUnits> getOMPAllocateAlignment(const VarDecl *VD);
+
   /// Returns whether the given record has hidden LTO visibility and therefore
   /// may participate in (single-module) CFI and whole-program vtable
   /// optimization.
   bool HasHiddenLTOVisibility(const CXXRecordDecl *RD);
 
   /// Returns whether the given record has public std LTO visibility
   /// and therefore may not participate in (single-module) CFI and whole-program
   /// vtable optimization.

clang/lib/CodeGen/CodeGenModule.cpp

#endif
 
   // If it is in a read-only section, mark it 'constant'.
   if (const SectionAttr *SA = D->getAttr<SectionAttr>()) {
     const ASTContext::SectionInfo &SI = Context.SectionInfos[SA->getName()];
     if ((SI.SectionFlags & ASTContext::PSF_Write) == 0)
       GV->setConstant(true);
   }
 
-  GV->setAlignment(getContext().getDeclAlign(D).getAsAlign());
+  CharUnits AlignVal = getContext().getDeclAlign(D);
+  // Check for alignment specifed in an 'omp allocate' directive.
+  if (llvm::Optional<CharUnits> AlignValFromAllocate =
+          getOMPAllocateAlignment(D))
+    AlignVal = AlignValFromAllocate.getValue();
+  GV->setAlignment(AlignVal.getAsAlign());
 
   // On Darwin, unlike other Itanium C++ ABI platforms, the thread-wrapper
   // function is only defined alongside the variable, not also alongside
   // callers. Normally, all accesses to a thread_local go through the
   // thread-wrapper in order to ensure initialization has occurred, underlying
   // variable will never be used other than the thread-wrapper, so it can be
   // converted to internal linkage.
   //

clang/test/OpenMP/align_clause_global_codegen.cpp

+// RUN: %clang_cc1 -emit-llvm -o - -fopenmp \
+// RUN:  -triple i386-unknown-unknown -fopenmp-version=51 %s \
+// RUN:  | FileCheck %s --check-prefixes=CHECK,CHECK-32
+
+// RUN: %clang_cc1 -emit-llvm -o - -fopenmp \
+// RUN:  -triple x86_64-unknown-linux-gnu -fopenmp-version=51 %s \
+// RUN:  | FileCheck %s --check-prefixes=CHECK,CHECK-64
+
+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;
+
+//
+// Should allow larger alignment.
+//
+
+// CHECK: @foo_global1 = global float 0.000000e+00, align 16
+float foo_global1;
+#pragma omp allocate(foo_global1) align(16)
+
+// CHECK: @foo_global2 = global float 0.000000e+00, align 16
+float foo_global2;
+#pragma omp allocate(foo_global2) allocator(omp_default_mem_alloc) align(16)
+
+// CHECK: @foo_global3 = global float 0.000000e+00, align 16
+float foo_global3;
+#pragma omp allocate(foo_global3) allocator(omp_large_cap_mem_alloc) align(16)
+
+// CHECK: @foop_global1 = global ptr null, align 16
+int *foop_global1;
+#pragma omp allocate(foop_global1) align(16)
+
+//
+// Should use natural alignment when alignment specified is too small.
+//
+
+// CHECK: @foo_global4 = global float 0.000000e+00, align 4
+float foo_global4;
+#pragma omp allocate(foo_global4) align(2)
+
+// CHECK: @foo_global5 = global float 0.000000e+00, align 4
+float foo_global5;
+#pragma omp allocate(foo_global5) allocator(omp_default_mem_alloc) align(2)
+
+// CHECK: @foo_global6 = global float 0.000000e+00, align 4
+float foo_global6;
+#pragma omp allocate(foo_global6) allocator(omp_large_cap_mem_alloc) align(2)
+
+// CHECK-32: @foop_global2 = global ptr null, align 4
+// CHECK-64: @foop_global2 = global ptr null, align 8
+int *foop_global2;
+#pragma omp allocate(foop_global2) align(2)