[OpenMP][JIT] Introduce support for AMDGPU

[jdoerfert]

Summary

To JIT kernels for AMDGPUs we need to provide the architecture, the

triple, and a post-link callback. The first two are simple, the last one
is a little more complicated since we need to invoke `lld`. There is
some library interface but for that we need the lld library, which is
not generally available, thus we go with the executable for now. In
either way we need to manifest the (amdgcn) object file and read the
output from another file. We should try to avoid that in the future.
The options for `lld` are copied from the way clang invokes it.

[tschuett]

Instead of invoking lld and copying the flags, could you invoke clang instead? The flags may change over time. Rust calls cc for linking.

[arsenm]

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

1654

SmallString

1656

Don’t need llvm::

1673

Also need to get ecc and xnack target features?

1679

Do we still need this flag with real LTO

1688

ExecuteAndWait should switch to std::error_code

1696

Return StringRef

2329

Why not return a complete triple

[tianshilei1992]

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

1696

It is not always good to return a StringRef here. In the NVPTX implementation, the arch string is built on the fly (in the return statement).

2329

Because we don't know the other components of a complete triple, though the assumption here is if the arch is same, the module should be compatible.

[arsenm]

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

1696

That doesn’t mean it should be done there

2329

That defeats the point of the other triple fields

[jhuber6]

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

1679

I don't think we need it since we should be linking only the object files to a shared object. Probably copied from the LinkerWrapper, I think it should probably be removed there as well.

[tianshilei1992]

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

1696

this is override function

[tianshilei1992]

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

2329

yes, and then when we compare, we have to get arch again, which I don’t see the point.

[jhuber6]

In D140720#4018224, @tschuett wrote:

Instead of invoking lld and copying the flags, could you invoke clang instead? The flags may change over time. Rust calls cc for linking.

I think invoking lld directly should be fine. It's the only linker that can link AMDGPU images and we already call it directly in the clang-linker-wrapper tool. To use clang you would need to pass --target=amdgcn-amd-amdhsa which most likely isn't compatible between different user cc's so we can't use that trick. If you do use clang is just gives this invocation, which probably won't be changed too much in the future:

"ld.lld" "device.o" "-shared" "-o" "a.out"

[jdoerfert]

@arsenm Not sure if/how we should deal with the feature flags (xsnack, etc.) I'll leave that up to AMD people to look into. This should give us basic support, tested with OpenMC on a gfx90a.

[tschuett]

Rust does not really care whether cc is clang or gcc. My point is that it is safer to invoke clang and let it manage the link step.

[jdoerfert]

I will update the nits I marked as done but didn't comment on. The commented stuff are kept for "follow ups", assuming people don't object. Anything else?

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

1673

Let's keep that question for a follow up. I honestly don't know. Basic support should work w/o all of this. AMD folks will look into it (and how to propagate the information).

1688

Not my department.

1696

Can't call a StringRef w/o refactoring stuff. Once we JIT per kernel, not per module, this might become interesting, for now, it's less of an issue.

2329

Similar answer as to above. Once we look at more than the arch + sub-arch, e.g., features, we might want to restructure this. I'll assume AMD (@jhuber6) will take a look how to do that.

[jdoerfert]

In D140720#4018590, @tschuett wrote:

Rust does not really care whether cc is clang or gcc. My point is that it is safer to invoke clang and let it manage the link step.

FWIW, what we really want is to use JITLink or lld as a library (neither we can really do right now). If people think using clang is better than lld directly, I can change it for sure.

[tschuett]

Halide is apparently using lld as a library: https://reviews.llvm.org/D140726

[jdoerfert]

In D140720#4018603, @tschuett wrote:

Halide is apparently using lld as a library: https://reviews.llvm.org/D140726

I even started using that interface (`lld::elf::link) but:

1. It doesn't solve our problem of materializing the two files for the link step, and
2. It requires the lld libraries, which we don't build by default and there is no guarantee the system lld (e.g., distributed with rocm) provides those.

[tianshilei1992]

In D140720#4018576, @jdoerfert wrote:

@arsenm Not sure if/how we should deal with the feature flags (xsnack, etc.) I'll leave that up to AMD people to look into. This should give us basic support, tested with OpenMC on a gfx90a.

Feature flags are not needed if you don't need portability. I do have a piece of code copied from somewhere that I don't remember: https://github.com/shiltian/llvm-project/commit/02bc7effccc6ff2f5ab3fe5218336094c0485766#diff-321c2038035972ad4994ff9d85b29950ba72c08a79891db5048b8f5d46915314R432.

[jdoerfert]

Addressed comments, I think.

[jdoerfert]

Pick the right patch file.

[tianshilei1992]

JIT part looks good to me. Not sure if the invocation of external tool will have more concerns.

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

 #include <atomic>
 #include <cassert>
 #include <cstddef>
 #include <deque>
 #include <hsa.h>
 #include <hsa_ext_amd.h>
 #include <mutex>
 #include <string>
+#include <system_error>
 #include <unistd.h>
 #include <unordered_map>
 
 #include "Debug.h"
 #include "DeviceEnvironment.h"
 #include "GlobalHandler.h"
 #include "PluginInterface.h"
 #include "Utilities.h"
 #include "UtilitiesRTL.h"
 
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Program.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
 namespace omp {
 namespace target {
 namespace plugin {
 
 /// Forward declarations for all specialized data structures.
 struct AMDGPUKernelTy;
struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   ~AMDGPUDeviceTy() {}
 
   /// Initialize the device, its resources and get its properties.
   Error initImpl(GenericPluginTy &Plugin) override {
     // First setup all the memory pools.
     if (auto Err = initMemoryPools())
       return Err;
 
+    char GPUName[64];
+    if (auto Err = getDeviceAttr(HSA_AGENT_INFO_NAME, GPUName))
+      return Err;
+    Arch = GPUName;
+
     // Get the wavefront size.
     uint32_t WavefrontSize = 0;
     if (auto Err = getDeviceAttr(HSA_AGENT_INFO_WAVEFRONT_SIZE, WavefrontSize))
       return Err;
     GridValues.GV_Warp_Size = WavefrontSize;
 
     // Load the grid values dependending on the wavefront.
     if (WavefrontSize == 32)
Error deinitImpl() override {
     }
 
     // Invalidate agent reference.
     Agent = {0};
 
     return Plugin::success();
   }
 
+  Expected<std::unique_ptr<MemoryBuffer>>
+  doJITPostProcessing(std::unique_ptr<MemoryBuffer> MB) const override {
+
+    // TODO: We should try to avoid materialization but there seems to be no
+    // good linker interface w/o file i/o.
+    SmallString<128> LinkerOutputFilePath;
+    std::error_code EC = sys::fs::createTemporaryFile(
+        "amdgpu-pre-link-jit", ".out", LinkerOutputFilePath);
+    if (EC)
+      return createStringError(EC,
+                               "Failed to create temporary file for linker");
+
+    SmallString<128> LinkerInputFilePath = LinkerOutputFilePath;
+    LinkerInputFilePath.pop_back_n(2);
+
+    auto FD = raw_fd_ostream(LinkerInputFilePath.data(), EC);
+    if (EC)
+      return createStringError(EC, "Failed to open temporary file for linker");
+    FD.write(MB->getBufferStart(), MB->getBufferSize());
+    FD.close();
+
+    const auto &ErrorOrPath = sys::findProgramByName("lld");
+    if (!ErrorOrPath)
+      return createStringError(inconvertibleErrorCode(),
+                               "Failed to find `lld` on the PATH.");
+
+    std::string LLDPath = ErrorOrPath.get();
+    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, getDeviceId(),
+         "Using `%s` to link JITed amdgcn ouput.", LLDPath.c_str());
+
+    std::string MCPU = "-plugin-opt=mcpu=" + getArch();
+
+    StringRef Args[] = {LLDPath,
+                        "-flavor",
+                        "gnu",
+                        "--no-undefined",
+                        "-shared",
+                        MCPU,
+                        "-o",
+                        LinkerOutputFilePath.data(),
+                        LinkerInputFilePath.data()};
+
+    std::string Error;
+    int RC = sys::ExecuteAndWait(LLDPath, Args, std::nullopt, {}, 0, 0, &Error);
+    if (RC)
+      return createStringError(inconvertibleErrorCode(),
+                               "Linking optimized bitcode failed: %s",
+                               Error.c_str());
+
+    return std::move(
+        MemoryBuffer::getFileOrSTDIN(LinkerOutputFilePath.data()).get());
+  }
+
+  std::string getArch() const override { return Arch; }
+
   /// 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)
private:
   AMDGPUEventManagerTy AMDGPUEventManager;
 
   /// Signal manager for AMDGPU signals.
   AMDGPUSignalManagerTy AMDGPUSignalManager;
 
   /// The agent handler corresponding to the device.
   hsa_agent_t Agent;
 
+  /// The GPU architecture.
+  std::string Arch;
+
   /// Reference to the host device.
   AMDHostDeviceTy &HostDevice;
 
   /// List of device packet queues.
   std::vector<AMDGPUQueueTy> Queues;
 };
 
 Error AMDGPUDeviceImageTy::loadExecutable(const AMDGPUDeviceTy &Device) {
Error deinitImpl() override {
     if (auto Err = HostDevice->deinit())
       return Err;
 
     // Finalize the HSA runtime.
     hsa_status_t Status = hsa_shut_down();
     return Plugin::check(Status, "Error in hsa_shut_down: %s");
   }
 
+  Triple::ArchType getTripleArch() const override { return Triple::amdgcn; }
+
   /// Get the ELF code for recognizing the compatible image binary.
   uint16_t getMagicElfBits() const override { return ELF::EM_AMDGPU; }
 
   /// Check whether the image is compatible with an AMDGPU device.
   Expected<bool> isImageCompatible(__tgt_image_info *Info) const override {
     for (hsa_agent_t Agent : KernelAgents) {
       std::string Target;
       auto Err = utils::iterateAgentISAs(Agent, [&](hsa_isa_t ISA) {