[OpenMP][libomptarget] Improve kernel initialization in plugins

[kevinsala]

Summary

This patch modifies the plugins so that the initialization of <Dev>KernelTy objects
is done in the init method. Part of the initialization was done in the
constructKernelEntry method. Now this method is called constructKernel
and only allocates and constructs a <Dev>KernelTy object.

This patch prepares the kernel class for the new implementation of device
reductions.

[tianshilei1992]

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1911–1912

Potentially this can fail right?

[kevinsala]

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1911–1912

Right, this can potentially fail. However, our allocate function directly calls the BumpPtrAllocator::Allocate(), which doesn't seem to control errors (e.g., when no more memory is available) other than asserts.

But anyway, I'm OK adding the following check for safety:

if (!AMDKernel)
  return Plugin::error("Failed to allocate memory for kernel %s", KernelEntry.name);

[kevinsala]

Fixing issues

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
 
   /// See GenericDeviceTy::getComputeUnitKind().
   std::string getComputeUnitKind() const override { return ComputeUnitKind; }
 
   /// Returns the clock frequency for the given AMDGPU device.
   uint64_t getClockFrequency() const override { return ClockFrequency; }
 
   /// Allocate and construct an AMDGPU kernel.
-  Expected<GenericKernelTy *>
-  constructKernelEntry(const __tgt_offload_entry &KernelEntry,
-                       DeviceImageTy &Image) override {
-
-    Expected<OMPTgtExecModeFlags> ExecModeOrErr =
-        getExecutionModeForKernel(KernelEntry.name, Image);
-    if (!ExecModeOrErr)
-      return ExecModeOrErr.takeError();
-
-    // Allocate and initialize the AMDGPU kernel.
-    AMDGPUKernelTy *AMDKernel = Plugin::get().allocate<AMDGPUKernelTy>();
-    new (AMDKernel) AMDGPUKernelTy(KernelEntry.name, ExecModeOrErr.get());
+  Expected<GenericKernelTy &>
+  constructKernel(const __tgt_offload_entry &KernelEntry,
+                  OMPTgtExecModeFlags ExecMode) override {
+    // Allocate and construct the AMDGPU kernel.
+    AMDGPUKernelTy *AMDGPUKernel = Plugin::get().allocate<AMDGPUKernelTy>();
+    if (!AMDGPUKernel)
+      return Plugin::error("Failed to allocate memory for AMDGPU kernel");
 
-    return AMDKernel;
+    new (AMDGPUKernel) AMDGPUKernelTy(KernelEntry.name, ExecMode);
+
+    return *AMDGPUKernel;
   }
 
   /// Set the current context to this device's context. Do nothing since the
   /// AMDGPU devices do not have the concept of contexts.
   Error setContext() override { return Plugin::success(); }
 
   /// We want to set up the RPC server for host services to the GPU if it is
   /// availible.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

Error registerGlobalOffloadEntry(DeviceImageTy &DeviceImage,
                                    __tgt_offload_entry &DeviceEntry);
 
   /// Register offload entry for kernel function.
   Error registerKernelOffloadEntry(DeviceImageTy &DeviceImage,
                                    const __tgt_offload_entry &KernelEntry,
                                    __tgt_offload_entry &DeviceEntry);
 
   /// Allocate and construct a kernel object.
-  virtual Expected<GenericKernelTy *>
-  constructKernelEntry(const __tgt_offload_entry &KernelEntry,
-                       DeviceImageTy &Image) = 0;
+  virtual Expected<GenericKernelTy &>
+  constructKernel(const __tgt_offload_entry &KernelEntry,
+                  OMPTgtExecModeFlags ExecMode) = 0;
 
   /// Get and set the stack size and heap size for the device. If not used, the
   /// plugin can implement the setters as no-op and setting the output
   /// value to zero for the getters.
   virtual Error getDeviceStackSize(uint64_t &V) = 0;
   virtual Error setDeviceStackSize(uint64_t V) = 0;
   virtual Error getDeviceHeapSize(uint64_t &V) = 0;
   virtual Error setDeviceHeapSize(uint64_t V) = 0;
private:
   /// Environment flag to set the minimum number of threads we use for a
   /// low-trip count combined loop. Instead of using more threads we increase
   /// the outer (block/team) parallelism.
   UInt32Envar OMPX_MinThreadsForLowTripCount =
       UInt32Envar("LIBOMPTARGET_MIN_THREADS_FOR_LOW_TRIP_COUNT", 32);
 
 protected:
   /// Return the execution mode used for kernel \p Name.
-  Expected<OMPTgtExecModeFlags> getExecutionModeForKernel(StringRef Name,
-                                                          DeviceImageTy &Image);
+  virtual Expected<OMPTgtExecModeFlags>
+  getExecutionModeForKernel(StringRef Name, DeviceImageTy &Image);
 
   /// Environment variables defined by the LLVM OpenMP implementation
   /// regarding the initial number of streams and events.
   UInt32Envar OMPX_InitialNumStreams;
   UInt32Envar OMPX_InitialNumEvents;
 
   /// Array of images loaded into the device. Images are automatically
   /// deallocated by the allocator.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

Error GenericDeviceTy::registerGlobalOffloadEntry(
   return Plugin::success();
 }
 
 Error GenericDeviceTy::registerKernelOffloadEntry(
     DeviceImageTy &Image, const __tgt_offload_entry &KernelEntry,
     __tgt_offload_entry &DeviceEntry) {
   DeviceEntry = KernelEntry;
 
+  // Retrieve the execution mode.
+  auto ExecModeOrErr = getExecutionModeForKernel(KernelEntry.name, Image);
+  if (!ExecModeOrErr)
+    return ExecModeOrErr.takeError();
+
   // Create a kernel object.
-  auto KernelOrErr = constructKernelEntry(KernelEntry, Image);
+  auto KernelOrErr = constructKernel(KernelEntry, *ExecModeOrErr);
   if (!KernelOrErr)
     return KernelOrErr.takeError();
 
-  GenericKernelTy *Kernel = *KernelOrErr;
-  assert(Kernel != nullptr && "Invalid kernel");
+  GenericKernelTy &Kernel = *KernelOrErr;
 
   // Initialize the kernel.
-  if (auto Err = Kernel->init(*this, Image))
+  if (auto Err = Kernel.init(*this, Image))
     return Err;
 
   // Set the device entry address to the kernel address and store the entry on
   // the entry table.
-  DeviceEntry.addr = (void *)Kernel;
+  DeviceEntry.addr = (void *)&Kernel;
   Image.getOffloadEntryTable().addEntry(DeviceEntry);
 
   return Plugin::success();
 }
 
 Expected<KernelEnvironmentTy>
 GenericDeviceTy::getKernelEnvironmentForKernel(StringRef Name,
                                                DeviceImageTy &Image) {

openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp

 namespace target {
 namespace plugin {
 
 /// Forward declarations for all specialized data structures.
 struct CUDAKernelTy;
 struct CUDADeviceTy;
 struct CUDAPluginTy;
 
+/// Class implementing the CUDA device images properties.
+struct CUDADeviceImageTy : public DeviceImageTy {
+  /// Create the CUDA image with the id and the target image pointer.
+  CUDADeviceImageTy(int32_t ImageId, const __tgt_device_image *TgtImage)
+      : DeviceImageTy(ImageId, TgtImage), Module(nullptr) {}
+
+  /// Load the image as a CUDA module.
+  Error loadModule() {
+    assert(!Module && "Module already loaded");
+
+    CUresult Res = cuModuleLoadDataEx(&Module, getStart(), 0, nullptr, nullptr);
+    if (auto Err = Plugin::check(Res, "Error in cuModuleLoadDataEx: %s"))
+      return Err;
+
+    return Plugin::success();
+  }
+
+  /// Unload the CUDA module corresponding to the image.
+  Error unloadModule() {
+    assert(Module && "Module not loaded");
+
+    CUresult Res = cuModuleUnload(Module);
+    if (auto Err = Plugin::check(Res, "Error in cuModuleUnload: %s"))
+      return Err;
+
+    Module = nullptr;
+
+    return Plugin::success();
+  }
+
+  /// Getter of the CUDA module.
+  CUmodule getModule() const { return Module; }
+
+private:
+  /// The CUDA module that loaded the image.
+  CUmodule Module;
+};
+
 /// Class implementing the CUDA kernel functionalities which derives from the
 /// generic kernel class.
 struct CUDAKernelTy : public GenericKernelTy {
-  /// Create a CUDA kernel with a name, an execution mode, and the kernel
-  /// function.
-  CUDAKernelTy(const char *Name, OMPTgtExecModeFlags ExecutionMode,
-               CUfunction Func)
-      : GenericKernelTy(Name, ExecutionMode), Func(Func) {}
+  /// Create a CUDA kernel with a name and an execution mode.
+  CUDAKernelTy(const char *Name, OMPTgtExecModeFlags ExecMode)
+      : GenericKernelTy(Name, ExecMode), Func(nullptr) {}
 
-  /// Initialize the CUDA kernel
+  /// Initialize the CUDA kernel.
   Error initImpl(GenericDeviceTy &GenericDevice,
                  DeviceImageTy &Image) override {
+    CUresult Res;
+    CUDADeviceImageTy &CUDAImage = static_cast<CUDADeviceImageTy &>(Image);
+
+    // Retrieve the function pointer of the kernel.
+    Res = cuModuleGetFunction(&Func, CUDAImage.getModule(), getName());
+    if (auto Err = Plugin::check(Res, "Error in cuModuleGetFunction('%s'): %s",
+                                 getName()))
+      return Err;
+
+    // Check that the function pointer is valid.
+    if (!Func)
+      return Plugin::error("Invalid function for kernel %s", getName());
+
     int MaxThreads;
-    CUresult Res = cuFuncGetAttribute(
-        &MaxThreads, CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, Func);
+    Res = cuFuncGetAttribute(&MaxThreads,
+                             CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, Func);
     if (auto Err = Plugin::check(Res, "Error in cuFuncGetAttribute: %s"))
       return Err;
 
-    /// Set the maximum number of threads for the CUDA kernel.
+    // The maximum number of threads cannot exceed the maximum of the kernel.
     MaxNumThreads = std::min(MaxNumThreads, (uint32_t)MaxThreads);
 
     return Plugin::success();
   }
 
-  /// Launch the CUDA kernel function
+  /// Launch the CUDA kernel function.
   Error launchImpl(GenericDeviceTy &GenericDevice, uint32_t NumThreads,
                    uint64_t NumBlocks, KernelArgsTy &KernelArgs, void *Args,
                    AsyncInfoWrapperTy &AsyncInfoWrapper) const override;
 
 private:
   /// The CUDA kernel function to execute.
   CUfunction Func;
 };
struct CUDAEventRef final : public GenericDeviceResourceRef {
   /// Get the underlying CUevent.
   operator HandleTy() const { return Event; }
 
 private:
   /// The reference to the CUDA event.
   HandleTy Event;
 };
 
-/// Class implementing the CUDA device images properties.
-struct CUDADeviceImageTy : public DeviceImageTy {
-  /// Create the CUDA image with the id and the target image pointer.
-  CUDADeviceImageTy(int32_t ImageId, const __tgt_device_image *TgtImage)
-      : DeviceImageTy(ImageId, TgtImage), Module(nullptr) {}
-
-  /// Load the image as a CUDA module.
-  Error loadModule() {
-    assert(!Module && "Module already loaded");
-
-    CUresult Res = cuModuleLoadDataEx(&Module, getStart(), 0, nullptr, nullptr);
-    if (auto Err = Plugin::check(Res, "Error in cuModuleLoadDataEx: %s"))
-      return Err;
-
-    return Plugin::success();
-  }
-
-  /// Unload the CUDA module corresponding to the image.
-  Error unloadModule() {
-    assert(Module && "Module not loaded");
-
-    CUresult Res = cuModuleUnload(Module);
-    if (auto Err = Plugin::check(Res, "Error in cuModuleUnload: %s"))
-      return Err;
-
-    Module = nullptr;
-
-    return Plugin::success();
-  }
-
-  /// Getter of the CUDA module.
-  CUmodule getModule() const { return Module; }
-
-private:
-  /// The CUDA module that loaded the image.
-  CUmodule Module;
-};
-
 /// Class implementing the CUDA device functionalities which derives from the
 /// generic device class.
 struct CUDADeviceTy : public GenericDeviceTy {
   // Create a CUDA device with a device id and the default CUDA grid values.
   CUDADeviceTy(int32_t DeviceId, int32_t NumDevices)
       : GenericDeviceTy(DeviceId, NumDevices, NVPTXGridValues),
         CUDAStreamManager(*this), CUDAEventManager(*this) {}
 
Error deinitImpl() override {
     // Invalidate context and device references.
     Context = nullptr;
     Device = CU_DEVICE_INVALID;
 
     return Plugin::success();
   }
 
   /// Allocate and construct a CUDA kernel.
-  Expected<GenericKernelTy *>
-  constructKernelEntry(const __tgt_offload_entry &KernelEntry,
-                       DeviceImageTy &Image) override {
-    CUDADeviceImageTy &CUDAImage = static_cast<CUDADeviceImageTy &>(Image);
-
-    // Retrieve the function pointer of the kernel.
-    CUfunction Func;
-    CUresult Res =
-        cuModuleGetFunction(&Func, CUDAImage.getModule(), KernelEntry.name);
-    if (auto Err = Plugin::check(Res, "Error in cuModuleGetFunction('%s'): %s",
-                                 KernelEntry.name))
-      return std::move(Err);
-
-    DP("Entry point " DPxMOD " maps to %s (" DPxMOD ")\n", DPxPTR(&KernelEntry),
-       KernelEntry.name, DPxPTR(Func));
-
-    Expected<OMPTgtExecModeFlags> ExecModeOrErr =
-        getExecutionModeForKernel(KernelEntry.name, Image);
-    if (!ExecModeOrErr)
-      return ExecModeOrErr.takeError();
-
-    // Allocate and initialize the CUDA kernel.
+  Expected<GenericKernelTy &>
+  constructKernel(const __tgt_offload_entry &KernelEntry,
+                  OMPTgtExecModeFlags ExecMode) override {
+    // Allocate and construct the CUDA kernel.
     CUDAKernelTy *CUDAKernel = Plugin::get().allocate<CUDAKernelTy>();
-    new (CUDAKernel) CUDAKernelTy(KernelEntry.name, ExecModeOrErr.get(), Func);
+    if (!CUDAKernel)
+      return Plugin::error("Failed to allocate memory for CUDA kernel");
+
+    new (CUDAKernel) CUDAKernelTy(KernelEntry.name, ExecMode);
 
-    return CUDAKernel;
+    return *CUDAKernel;
   }
 
   /// Set the current context to this device's context.
   Error setContext() override {
     CUresult Res = cuCtxSetCurrent(Context);
     return Plugin::check(Res, "Error in cuCtxSetCurrent: %s");
   }
 

openmp/libomptarget/plugins-nextgen/generic-elf-64bit/src/rtl.cpp

 struct GenELF64KernelTy;
 struct GenELF64DeviceTy;
 struct GenELF64PluginTy;
 
 using llvm::sys::DynamicLibrary;
 
 /// Class implementing kernel functionalities for GenELF64.
 struct GenELF64KernelTy : public GenericKernelTy {
-  /// Construct the kernel with a name, execution mode and a function.
-  GenELF64KernelTy(const char *Name, OMPTgtExecModeFlags ExecutionMode,
-                   void (*Func)(void))
-      : GenericKernelTy(Name, ExecutionMode), Func(Func) {}
+  /// Construct the kernel with a name and an execution mode.
+  GenELF64KernelTy(const char *Name, OMPTgtExecModeFlags ExecMode)
+      : GenericKernelTy(Name, ExecMode), Func(nullptr) {}
 
   /// Initialize the kernel.
-  Error initImpl(GenericDeviceTy &GenericDevice,
-                 DeviceImageTy &Image) override {
+  Error initImpl(GenericDeviceTy &Device, DeviceImageTy &Image) override {
+    // Functions have zero size.
+    GlobalTy Global(getName(), 0);
+
+    // Get the metadata (address) of the kernel function.
+    GenericGlobalHandlerTy &GHandler = Plugin::get().getGlobalHandler();
+    if (auto Err = GHandler.getGlobalMetadataFromDevice(Device, Image, Global))
+      return Err;
+
+    // Check that the function pointer is valid.
+    if (!Global.getPtr())
+      return Plugin::error("Invalid function for kernel %s", getName());
+
+    // Save the function pointer.
+    Func = (void (*)())Global.getPtr();
+
     // Set the maximum number of threads to a single.
     MaxNumThreads = 1;
     return Plugin::success();
   }
 
   /// Launch the kernel using the libffi.
   Error launchImpl(GenericDeviceTy &GenericDevice, uint32_t NumThreads,
                    uint64_t NumBlocks, KernelArgsTy &KernelArgs, void *Args,
struct GenELF64DeviceTy : public GenericDeviceTy {
 
   /// Initialize the device, which is a no-op
   Error initImpl(GenericPluginTy &Plugin) override { return Plugin::success(); }
 
   /// Deinitialize the device, which is a no-op
   Error deinitImpl() override { return Plugin::success(); }
 
   /// Construct the kernel for a specific image on the device.
-  Expected<GenericKernelTy *>
-  constructKernelEntry(const __tgt_offload_entry &KernelEntry,
-                       DeviceImageTy &Image) override {
-    GlobalTy Func(KernelEntry);
-
-    // Get the metadata (address) of the kernel function.
-    GenericGlobalHandlerTy &GHandler = Plugin::get().getGlobalHandler();
-    if (auto Err = GHandler.getGlobalMetadataFromDevice(*this, Image, Func))
-      return std::move(Err);
-
-    // Allocate and create the kernel.
+  Expected<GenericKernelTy &>
+  constructKernel(const __tgt_offload_entry &KernelEntry,
+                  OMPTgtExecModeFlags ExecMode) override {
+    // Allocate and construct the kernel.
     GenELF64KernelTy *GenELF64Kernel =
         Plugin::get().allocate<GenELF64KernelTy>();
-    new (GenELF64Kernel) GenELF64KernelTy(
-        KernelEntry.name, OMP_TGT_EXEC_MODE_GENERIC, (void (*)())Func.getPtr());
+    if (!GenELF64Kernel)
+      return Plugin::error("Failed to allocate memory for GenELF64 kernel");
 
-    return GenELF64Kernel;
+    new (GenELF64Kernel) GenELF64KernelTy(KernelEntry.name, ExecMode);
+
+    return *GenELF64Kernel;
   }
 
   /// Set the current context to this device, which is a no-op.
   Error setContext() override { return Plugin::success(); }
 
   /// Load the binary image into the device and allocate an image object.
   Expected<DeviceImageTy *> loadBinaryImpl(const __tgt_device_image *TgtImage,
                                            int32_t ImageId) override {
Error setDeviceStackSize(uint64_t Value) override {
     return Plugin::success();
   }
   Error getDeviceHeapSize(uint64_t &Value) override {
     Value = 0;
     return Plugin::success();
   }
   Error setDeviceHeapSize(uint64_t Value) override { return Plugin::success(); }
 
+protected:
+  /// Retrieve the execution mode for kernels. All kernels use the generic mode.
+  Expected<OMPTgtExecModeFlags>
+  getExecutionModeForKernel(StringRef Name, DeviceImageTy &Image) override {
+    return OMP_TGT_EXEC_MODE_GENERIC;
+  }
+
 private:
   /// Grid values for Generic ELF64 plugins.
   static constexpr GV GenELF64GridValues = {
       1, // GV_Slot_Size
       1, // GV_Warp_Size
       1, // GV_Max_Teams
       1, // GV_Default_Num_Teams
       1, // GV_SimpleBufferSize