[OpenMP][OMPIRBuilder] Add device shared memory allocation support

[skatrak]

This patch adds the __kmpc_alloc_shared and __kmpc_free_shared DeviceRTL functions to the list of those the OMPIRBuilder is able to create.

[llvmbot]

@llvm/pr-subscribers-flang-openmp

Author: Sergio Afonso (skatrak)

Changes

This patch adds the __kmpc_alloc_shared and __kmpc_free_shared DeviceRTL functions to the list of those the OMPIRBuilder is able to create.

---

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

2 Files Affected:

• (modified) llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h (+23)
• (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+27)

diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 206ad4a4ef85f..110b0fde863c5 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -2777,6 +2777,29 @@ class OpenMPIRBuilder {
   LLVM_ABI CallInst *createOMPFree(const LocationDescription &Loc, Value *Addr,
                                    Value *Allocator, std::string Name = "");
 
+  /// Create a runtime call for kmpc_alloc_shared.
+  ///
+  /// \param Loc The insert and source location description.
+  /// \param VarType Type of variable to be allocated.
+  /// \param Name Name of call Instruction.
+  ///
+  /// \returns CallInst to the kmpc_alloc_shared call.
+  LLVM_ABI CallInst *createOMPAllocShared(const LocationDescription &Loc,
+                                          Type *VarType,
+                                          const Twine &Name = Twine(""));
+
+  /// Create a runtime call for kmpc_free_shared.
+  ///
+  /// \param Loc The insert and source location description.
+  /// \param Addr Value obtained from the corresponding kmpc_alloc_shared call.
+  /// \param VarType Type of variable to be freed.
+  /// \param Name Name of call Instruction.
+  ///
+  /// \returns CallInst to the kmpc_free_shared call.
+  LLVM_ABI CallInst *createOMPFreeShared(const LocationDescription &Loc,
+                                         Value *Addr, Type *VarType,
+                                         const Twine &Name = Twine(""));
+
   /// Create a runtime call for kmpc_threadprivate_cached
   ///
   /// \param Loc The insert and source location description.
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 79287729fbfd1..2e8fb5efb7743 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6213,6 +6213,33 @@ CallInst *OpenMPIRBuilder::createOMPFree(const LocationDescription &Loc,
   return Builder.CreateCall(Fn, Args, Name);
 }
 
+CallInst *OpenMPIRBuilder::createOMPAllocShared(const LocationDescription &Loc,
+                                                Type *VarType,
+                                                const Twine &Name) {
+  IRBuilder<>::InsertPointGuard IPG(Builder);
+  updateToLocation(Loc);
+
+  const DataLayout &DL = M.getDataLayout();
+  Value *Args[] = {Builder.getInt64(DL.getTypeStoreSize(VarType))};
+  Function *Fn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_alloc_shared);
+  CallInst *Call = Builder.CreateCall(Fn, Args, Name);
+  Call->addRetAttr(
+      Attribute::getWithAlignment(M.getContext(), DL.getPrefTypeAlign(Int64)));
+  return Call;
+}
+
+CallInst *OpenMPIRBuilder::createOMPFreeShared(const LocationDescription &Loc,
+                                               Value *Addr, Type *VarType,
+                                               const Twine &Name) {
+  IRBuilder<>::InsertPointGuard IPG(Builder);
+  updateToLocation(Loc);
+
+  Value *Args[] = {
+      Addr, Builder.getInt64(M.getDataLayout().getTypeStoreSize(VarType))};
+  Function *Fn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_free_shared);
+  return Builder.CreateCall(Fn, Args, Name);
+}
+
 CallInst *OpenMPIRBuilder::createOMPInteropInit(
     const LocationDescription &Loc, Value *InteropVar,
     omp::OMPInteropType InteropType, Value *Device, Value *NumDependences,

[skatrak]

PR stack:

• [MLIR][OpenMP] Remove Generic-SPMD early detection #150922
• [OpenMP][OMPIRBuilder] Add device shared memory allocation support #150923 ◀️
• [MLIR][OpenMP] Support allocations of device shared memory #150924
• [OpenMP][OMPIRBuilder] Use device shared memory for arg structures #150925
• [OpenMP][OMPIRBuilder] Support parallel in Generic kernels #150926
• [OpenMPOpt] Make parallel regions reachable from new DeviceRTL loop functions #150927

[Meinersbur]

LGTM

I would normally ask for unittests, but these seem sufficiently simple/straightforward.

[bhandarkar-pranav]

A thought that is not terribly urgent simply because my understanding is that reverse offload is not terribly important right now. However, when we reverse offload, we'll have to ensure that parallel threads on the CPU (host) now are able to access incoming data which used to be on the stack. Do the two __kmpc_shared_alloc/dealloc calls degenerate to mallocs on the host?