[all-commits] [llvm/llvm-project] 3ede81: [Libomptarget] Fix JIT on the NVPTX target by call...

[Joseph Huber via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 3ede817f5bd947cb0da63187f333a6274bf1f418
      https://github.com/llvm/llvm-project/commit/3ede817f5bd947cb0da63187f333a6274bf1f418
  Author: Joseph Huber <huberjn at outlook.com>
  Date:   2024-01-11 (Thu, 11 Jan 2024)

  Changed paths:
    M openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp

  Log Message:
  -----------
  [Libomptarget] Fix JIT on the NVPTX target by calling ptx manually (#77801)

Summary:
Recently a patch added an assertion in the GlobalHandler to indicate
when an ELF was not used. This began to fire whenever NVPTX JIT was
used, because the JIT pass output a PTX file instead of an ELF. The
CUModuleLoad method consumes `.s` internally and compiles it to a cubin,
however, this is too late as we perform several checks on the ELF
directly for the presence of certain symbols and to read some necessary
constants. This results in inconsistent behaviour.

To address this, this patch simply calls `ptxas` manually, similar to
how `lld` is called for the AMDGPU JIT pass. This is inevitably going to
be slower than simply passing it to the CUDA routine due to the overhead
involved in file IO and a fork call, but it's necessary for correctness.

CUDA provides an API for compiling PTX manually. However, this only
started showing up in CUDA 11.1 and is only provided "officially" in a
static library. The `libnvidia-ptxjitcompiler.so` next to the CUDA
driver has the same symbols and can likely be used as a replacement.
This would be the faster solution. However, given that it's not
documented it may have some issues.