[Offload] Add olLinkProgram

offload/liboffload/API/Program.td

@@ -25,6 +25,34 @@ def : Function {
     let returns = [];
 }
 
+def : Struct {
+    let name = "ol_program_link_buffer_t";
+    let desc = "An image to link with `olLinkProgram`.";
+    let members = [
+        StructMember<"void *", "Address", "base address of memory image">,
+        StructMember<"size_t", "Size", "size in bytes of memory image">,
+    ];
+}
+
+def : Function {
+    let name = "olLinkProgram";
+    let desc = "Compile and link multiple bitcode images into a single binary.";
+    let details = [
+        "No caching is performed; multiple calls to `olLinkProgram` with the same images will result in multiple linking operations",
+    ];
+    let params = [
+        Param<"ol_device_handle_t", "Device", "handle of the device to link for", PARAM_IN>,
+        Param<"ol_program_link_buffer_t *", "Images", "a pointer to an array of `ImagesSize` entries, one for each image to link", PARAM_IN>,
+        Param<"size_t", "ImagesSize", "the number of elements in `Images`", PARAM_IN>,
+        Param<"ol_program_handle_t*", "Program", "output handle for the created program", PARAM_OUT>
+    ];
+    let returns = [
+        Return<"OL_ERRC_INVALID_SIZE", ["`ImagesSize == 0`"]>,
+        Return<"OL_ERRC_INVALID_BINARY", ["Any image is not in the bitcode format"]>,
+        Return<"OL_ERRC_UNSUPPORTED", ["Linking is not supported for this device and `ImagesSize` > 1"]>,
+    ];
+}
+
 def : Function {
     let name = "olDestroyProgram";
     let desc = "Destroy the program and free all underlying resources.";

offload/liboffload/src/OffloadImpl.cpp

@@ -650,6 +650,39 @@ Error olCreateProgram_impl(ol_device_handle_t Device, const void *ProgData,
   return Error::success();
 }
 
+Error olLinkProgram_impl(ol_device_impl_t *Device,
+                         ol_program_link_buffer_t *Images, size_t ImagesSize,
+                         ol_program_handle_t *Program) {
+  std::vector<__tgt_device_image> DevImages;
+  for (size_t I = 0; I < ImagesSize; I++) {
+    auto &ProgData = Images[I];
+    DevImages.push_back({ProgData.Address,
+                         utils::advancePtr(ProgData.Address, ProgData.Size),
+                         nullptr, nullptr});
+  }
+
+  auto LinkResult =
+      Device->Device->jitLinkBinary(Device->Device->Plugin, DevImages);
+  if (!LinkResult)
+    return LinkResult.takeError();
+
+  ol_program_handle_t Prog =
+      new ol_program_impl_t(nullptr, nullptr, *LinkResult);
+
+  auto Res =
+      Device->Device->loadBinary(Device->Device->Plugin, &Prog->DeviceImage);
+  if (!Res) {
+    delete Prog;
+    return Res.takeError();
+  }
+  assert(*Res != nullptr && "loadBinary returned nullptr");
+
+  Prog->Image = *Res;
+  *Program = Prog;
+
+  return Error::success();
+}
+
 Error olDestroyProgram_impl(ol_program_handle_t Program) {
   auto &Device = Program->Image->getDevice();
   if (auto Err = Device.unloadBinary(Program->Image))

offload/plugins-nextgen/common/include/JIT.h

@@ -55,6 +55,10 @@ struct JITEngine {
   process(const __tgt_device_image &Image,
           target::plugin::GenericDeviceTy &Device);
 
+  /// Link and compile multiple bitcode images into a single binary
+  Expected<__tgt_device_image> link(std::vector<__tgt_device_image> &Images,
+                                    target::plugin::GenericDeviceTy &Device);
+
 private:
   /// Compile the bitcode image \p Image and generate the binary image that can
   /// be loaded to the target device of the triple \p Triple architecture \p

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -749,6 +749,10 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   /// Load the binary image into the device and return the target table.
   Expected<DeviceImageTy *> loadBinary(GenericPluginTy &Plugin,
                                        const __tgt_device_image *TgtImage);
+  /// Link and compile multiple bitcode images into a single image.
+  Expected<__tgt_device_image>
+  jitLinkBinary(GenericPluginTy &Plugin,
+                std::vector<__tgt_device_image> InputImages);
   virtual Expected<DeviceImageTy *>
   loadBinaryImpl(const __tgt_device_image *TgtImage, int32_t ImageId) = 0;
 

offload/plugins-nextgen/common/src/JIT.cpp

@@ -327,3 +327,44 @@ JITEngine::process(const __tgt_device_image &Image,
 
   return &Image;
 }
+
+Expected<__tgt_device_image>
+JITEngine::link(std::vector<__tgt_device_image> &Images,
+                target::plugin::GenericDeviceTy &Device) {
+  const std::string &ComputeUnitKind = Device.getComputeUnitKind();
+  ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind];
+
+  PostProcessingFn PostProcessing =
+      [&Device](llvm::SmallVector<std::unique_ptr<MemoryBuffer>> &&MB)
+      -> Expected<std::unique_ptr<MemoryBuffer>> {
+    return Device.doJITPostProcessing(std::move(MB));
+  };
+
+  std::lock_guard<std::mutex> Lock(ComputeUnitMapMutex);
+
+  llvm::SmallVector<std::unique_ptr<MemoryBuffer>> Buffers;
+  size_t Index = 0;
+  for (auto &I : Images) {
+    if (!isImageBitcode(I))
+      return error::createOffloadError(
+          error::ErrorCode::INVALID_BINARY,
+          "binary %i provided to link operation is not bitcode", Index);
+    Index++;
+
+    auto ObjMBOrErr = getOrCreateObjFile(I, CUI.Context, ComputeUnitKind);
+    if (!ObjMBOrErr)
+      return ObjMBOrErr.takeError();
+    Buffers.push_back(std::move(*ObjMBOrErr));
+  }
+
+  auto ImageMBOrErr = PostProcessing(std::move(Buffers));
+  if (!ImageMBOrErr)
+    return ImageMBOrErr.takeError();
+
+  auto &ImageMB = CUI.JITImages.emplace_back(std::move(*ImageMBOrErr));
+  __tgt_device_image JITedImage{};
+  JITedImage.ImageStart = const_cast<char *>(ImageMB->getBufferStart());
+  JITedImage.ImageEnd = const_cast<char *>(ImageMB->getBufferEnd());
+
+  return JITedImage;
+}

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -903,6 +903,13 @@ Error GenericDeviceTy::deinit(GenericPluginTy &Plugin) {
 
   return deinitImpl();
 }
+
+Expected<__tgt_device_image>
+GenericDeviceTy::jitLinkBinary(GenericPluginTy &Plugin,
+                               std::vector<__tgt_device_image> InputImages) {
+  return Plugin.getJIT().link(InputImages, *this);
+}
+
 Expected<DeviceImageTy *>
 GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
                             const __tgt_device_image *InputTgtImage) {

offload/unittests/OffloadAPI/CMakeLists.txt

@@ -32,7 +32,8 @@ add_offload_unittest("platform"
 
 add_offload_unittest("program"
     program/olCreateProgram.cpp
-    program/olDestroyProgram.cpp)
+    program/olDestroyProgram.cpp
+    program/olLinkProgram.cpp)
 
 add_offload_unittest("queue"
     queue/olCreateQueue.cpp

offload/unittests/OffloadAPI/device_code/CMakeLists.txt

@@ -8,6 +8,8 @@ add_offload_test_device_code(localmem_static.c localmem_static)
 add_offload_test_device_code(global.c global)
 add_offload_test_device_code(global_ctor.c global_ctor)
 add_offload_test_device_code(global_dtor.c global_dtor)
+add_offload_test_device_code(link_a.c link_a)
+add_offload_test_device_code(link_b.c link_b)
 
 add_custom_target(offload_device_binaries DEPENDS
     foo.bin
@@ -19,5 +21,7 @@ add_custom_target(offload_device_binaries DEPENDS
     global.bin
     global_ctor.bin
     global_dtor.bin
+    link_a.bin
+    link_b.bin
 )
 set(OFFLOAD_TEST_DEVICE_CODE_PATH ${CMAKE_CURRENT_BINARY_DIR} PARENT_SCOPE)

offload/unittests/OffloadAPI/device_code/link_a.c

@@ -0,0 +1,11 @@
+#include <gpuintrin.h>
+#include <stdint.h>
+
+uint32_t global;
+
+extern uint32_t funky();
+
+__gpu_kernel void link_a(uint32_t *out) {
+  out[0] = funky();
+  out[1] = global;
+}

offload/unittests/OffloadAPI/device_code/link_b.c

@@ -0,0 +1,10 @@
+#include <gpuintrin.h>
+#include <stdint.h>
+
+extern uint32_t global[64];
+
+[[gnu::visibility("default")]]
+uint32_t funky() {
+  global[0] = 100;
+  return 200;
+}

offload/unittests/OffloadAPI/program/olLinkProgram.cpp

@@ -0,0 +1,99 @@
+//===------- Offload API tests - olCreateProgram --------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "../common/Fixtures.hpp"
+#include <OffloadAPI.h>
+#include <gtest/gtest.h>
+
+using olLinkProgramTest = OffloadQueueTest;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olLinkProgramTest);
+
+TEST_P(olLinkProgramTest, SuccessSingle) {
+
+  std::unique_ptr<llvm::MemoryBuffer> DeviceBin;
+  ASSERT_TRUE(TestEnvironment::loadDeviceBinary("foo", Device, DeviceBin));
+  ASSERT_GE(DeviceBin->getBufferSize(), 0lu);
+
+  ol_program_link_buffer_t Buffers[1] = {
+      {const_cast<char *>(DeviceBin->getBufferStart()),
+       DeviceBin->getBufferSize()},
+  };
+
+  ol_program_handle_t Program;
+  ASSERT_SUCCESS(olLinkProgram(Device, Buffers, 1, &Program));
+  ASSERT_NE(Program, nullptr);
+
+  ASSERT_SUCCESS(olDestroyProgram(Program));
+}
+
+TEST_P(olLinkProgramTest, SuccessBuild) {
+  std::unique_ptr<llvm::MemoryBuffer> ABin;
+  ASSERT_TRUE(TestEnvironment::loadDeviceBinary("link_a", Device, ABin));
+  std::unique_ptr<llvm::MemoryBuffer> BBin;
+  ASSERT_TRUE(TestEnvironment::loadDeviceBinary("link_b", Device, BBin));
+
+  ol_program_link_buffer_t Buffers[2] = {
+      {const_cast<char *>(ABin->getBufferStart()), ABin->getBufferSize()},
+      {const_cast<char *>(BBin->getBufferStart()), BBin->getBufferSize()},
+  };
+
+  ol_program_handle_t Program;
+  auto LinkResult = olLinkProgram(Device, Buffers, 2, &Program);
+  if (LinkResult && LinkResult->Code == OL_ERRC_UNSUPPORTED)
+    GTEST_SKIP() << "Linking unsupported: " << LinkResult->Details;
+  ASSERT_SUCCESS(LinkResult);
+  ASSERT_NE(Program, nullptr);
+
+  ol_symbol_handle_t Kernel;
+  ASSERT_SUCCESS(
+      olGetSymbol(Program, "link_a", OL_SYMBOL_KIND_KERNEL, &Kernel));
+
+  void *Mem;
+  ASSERT_SUCCESS(
+      olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, 2 * sizeof(uint32_t), &Mem));
+  struct {
+    void *Mem;
+  } Args{Mem};
+  ol_kernel_launch_size_args_t LaunchArgs{};
+  LaunchArgs.GroupSize = {1, 1, 1};
+  LaunchArgs.NumGroups = {1, 1, 1};
+  LaunchArgs.Dimensions = 1;
+
+  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &Args, sizeof(Args),
+                                &LaunchArgs, nullptr));
+  ASSERT_SUCCESS(olWaitQueue(Queue));
+
+  uint32_t *Data = (uint32_t *)Mem;
+  ASSERT_EQ(Data[0], 200);
+  ASSERT_EQ(Data[1], 100);
+
+  ASSERT_SUCCESS(olMemFree(Mem));
+  ASSERT_SUCCESS(olDestroyProgram(Program));
+}
+
+TEST_P(olLinkProgramTest, InvalidNotBitcode) {
+  char FakeElf[] =
+      "\177ELF0000000000000000000000000000000000000000000000000000"
+      "00000000000000000000000000000000000000000000000000000000000";
+
+  ol_program_link_buffer_t Buffers[1] = {
+      {FakeElf, sizeof(FakeElf)},
+  };
+
+  ol_program_handle_t Program;
+  ASSERT_ERROR(OL_ERRC_INVALID_BINARY,
+               olLinkProgram(Device, Buffers, 1, &Program));
+}
+
+TEST_P(olLinkProgramTest, InvalidSize) {
+  ol_program_link_buffer_t Buffers[0] = {};
+
+  ol_program_handle_t Program;
+  ASSERT_ERROR(OL_ERRC_INVALID_SIZE,
+               olLinkProgram(Device, Buffers, 0, &Program));
+}