8388283: Add opt-in NVRTC compilation for CUDA backend

[xyyNicole]

HAT generates CUDA source at runtime and immediately compiles it before launching the generated kernel. This flow is a good match for NVRTC because it avoids spawning an external compiler process for each generated CUDA source and can significantly reduce runtime compilation overhead.

Add an opt-in NVRTC compilation path for the HAT CUDA backend while keeping NVCC as the default compiler. Users can select the NVRTC path with:

  HAT_CUDA_COMPILER=nvrtc java @.ffi-cuda-example matmul.Main --kernel=1D

NVRTC is loaded lazily with dlopen/dlsym, so the CUDA backend does not gain a link-time dependency on libnvrtc and existing NVCC/PTX flows can still run without loading NVRTC. The generated CUDA/PTX artifacts are kept under ./var/cuda for parity with the NVCC path.

Performance gains measured with matmul.Main 2DTILING on Grace Hopper:

1. CUDA source compilation
   NVRTC: 50.9 ms
   NVCC: 213.1 ms
   Speedup: ~4x

2. First execution including compilation
   NVRTC: 348.68 ms
   NVCC: 544.35 ms

---

• I confirm that I make this contribution in accordance with the OpenJDK Interim AI Policy.

---

Progress

• Change must not contain extraneous whitespace

Issue

• JDK-8388283: Add opt-in NVRTC compilation for CUDA backend (Enhancement - P4)

Reviewing

Using git

Checkout this PR locally:
$ git fetch https://git.openjdk.org/babylon.git pull/1105/head:pull/1105
$ git checkout pull/1105

Update a local copy of the PR:
$ git checkout pull/1105
$ git pull https://git.openjdk.org/babylon.git pull/1105/head

Using Skara CLI tools

Checkout this PR locally:
$ git pr checkout 1105

View PR using the GUI difftool:
$ git pr show -t 1105

Using diff file

Download this PR as a diff file:
https://git.openjdk.org/babylon/pull/1105.diff

Using Webrev

Link to Webrev Comment

[bridgekeeper]

👋 Welcome back xyyNicole! A progress list of the required criteria for merging this PR into code-reflection will be added to the body of your pull request. There are additional pull request commands available for use with this pull request.

[openjdk]

@xyyNicole This change now passes all automated pre-integration checks.

ℹ️ This project also has non-automated pre-integration requirements. Please see the file CONTRIBUTING.md for details.

After integration, the commit message for the final commit will be:

8388283: Add opt-in NVRTC compilation for CUDA backend

You can use pull request commands such as /summary, /contributor and /issue to adjust it as needed.

At the time when this comment was updated there had been no new commits pushed to the code-reflection branch. If another commit should be pushed before you perform the /integrate command, your PR will be automatically rebased. If you prefer to avoid any potential automatic rebasing, please check the documentation for the /integrate command for further details.

As you do not have Committer status in this project an existing Committer must agree to sponsor your change.

➡️ To flag this PR as ready for integration with the above commit message, type /integrate in a new comment. (Afterwards, your sponsor types /sponsor in a new comment to perform the integration).

[mlbridge]

Webrevs

• 00: Full (4919b247)

[grfrost]

This is interesting. I was not aware of this option. Rather than a separate HAT_CUDA_COMPILER I would suggest we allocate a 'bit' for this in hat.Config. We reserve a 32 int (bits) for comms between java side and FFI native side. I suspect we can free up a bit for this.

[grfrost]

I should have added an example, so we could hopefully use

HAT=CUDA_COMPILER:nvrtc java @.ffi-cuda-example matmul.Main --kernel=1D

See https://github.com/openjdk/babylon/blob/code-reflection/hat/core/src/main/java/hat/Config.java

[xyyNicole]

I should have added an example, so we could hopefully use

HAT=CUDA_COMPILER:nvrtc java @.ffi-cuda-example matmul.Main --kernel=1D

See https://github.com/openjdk/babylon/blob/code-reflection/hat/core/src/main/java/hat/Config.java

Thanks for the review and the concrete example. From the perspective of HAT's existing configuration model, I agree that HAT=CUDA_COMPILER:nvrtc would be more consistent than introducing a separate HAT_CUDA_COMPILER environment variable.

One concern is that the current 32-bit config word appears to have no free bits left. I am not fully familiar with the intended use of every existing field, so I would appreciate your advice on whether any bit could be freed or reused. For
example, PROFILE and PROFILE_CUDA_KERNEL seem somewhat related. Would it be reasonable to retire or repurpose the PROFILE_CUDA_KERNEL bit?

Another option would be to extend the config word from int to long, which would also provide room for future options. However, that appears to require a broader config layout and Java-to-FFI ABI change, so I am not sure it belongs in this NVRTC patch.

[grfrost]

Ah Juan and I just chatted about this is slack. We will find a way to free up some bits ;) I think we have some dead/unused bits. Also we recognize that some bits are FFI only (so passed to native backend), we should have separated these flags/bits into Java only and FFI only. So Config itself may use two banks of 32 bits. Since this nvrtc bit is likely FFI only we should have plenty once we separate them out. In the short term ;) We will let this PR use it's own env var, then later will repack the bits and then refactor this code to use a freed up slot. To avoid merge issues.

[xyyNicole]

Ah Juan and I just chatted about this is slack. We will find a way to free up some bits ;) I think we have some dead/unused bits. Also we recognize that some bits are FFI only (so passed to native backend), we should have separated these flags/bits into Java only and FFI only. So Config itself may use two banks of 32 bits. Since this nvrtc bit is likely FFI only we should have plenty once we separate them out. In the short term ;) We will let this PR use it's own env var, then later will repack the bits and then refactor this code to use a freed up slot. To avoid merge issues.

Thanks, that sounds good. I will keep the separate environment variable in this patch to avoid conflicts, and the NVRTC compiler selection can be migrated to an FFI config bit after the Config bits are repacked.

[xyyNicole]

One issue worth noting is that native object ownership in the CUDA backend appears to cause memory leaks. For example, nvcc() returns a heap-allocated PtxSource, while compile(const PtxSource *) creates Log objects and returns a CudaModule containing raw ptx->text and infLog->text pointers, without a clear ownership or release path.

Fixing this properly requires a broader refactoring to define the lifetimes of PtxSource, Log, their buffers, and CudaModule consistently. This patch therefore keeps the existing ownership model. A proper cleanup could be addressed in a separate future change.

[jjfumero]

Thanks @xyyNicole for this PR. I could check a benefit of 60-70% faster for compilation.

When running the unit-tests as follows:

HAT_CUDA_COMPILER=nvrtc java @.test-suite ffi-cuda  

Many unit-tests failed to run:

Global passed: 34, failed: 0, unsupported: 0, precision-errors: 1, pass-rate: 97.14
[REPORT] Test failed. Some tests were not launched. Common reasons: seg-faults, driver-issues. Please, check again.
    - Expected to run: 138. But only launched 35

For example:

   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestBFloat16Type
   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestF16Type
   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestFlashAttention
   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestFloat2
   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestHATMathLib
   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestTensors
   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestVectorArrayView
   HAT Engine Testing Framework. Backend: hat.backend.ffi.CudaBackend
failed to run test TestVectorTypes

I tracked down one of them (TestTensors) and I get the following error:

$ export HAT_CUDA_COMPILER=nvrtc 
$ HAT=CHECK_SSA_LOWERING,INFO,SHOW_CODE java @.ffi-cuda-test hat.test.TestTensors#testTensor03
... 
> NVRTC log:
./var/cuda/tmp_1784724386117.cu(71): error: unrecognized token
      unsigned char pad$lNwuQ[12];
                       ^

./var/cuda/tmp_1784724386117.cu(71): error: expected a ";"
      unsigned char pad$lNwuQ[12];
                       ^

2 errors detected in the compilation of "./var/cuda/tmp_1784724386117.cu".

We insert this symbol when we insert padding automatically, or construct new symbols at runtime (during codegen). Is there any way to bypass this in nvrtc?

Using HAT_CUDA_COMPILER=nvcc works fine.

Global passed: 134, failed: 1, unsupported: 0, precision-errors: 3, pass-rate: 97.10
[REPORT] OK: All tests launched. Total: 138

[jjfumero]

BTW, I am using CUDA 13.3 and NVIDIA 610.43.02 driver.

[grfrost]

Looks like the cuda compiler no longer accepts '$' as a valid identifier char. I was always suprised we got away with this in CUDA and OpenCL.

[jjfumero]

WRT F16 and BFloat16, we get the following errors:

[INFO] CUDA source compiler: NVRTC
NVRTC compilation failed: NVRTC_ERROR_COMPILATION. CUDA source saved to ./var/cuda/tmp_1784724931385.cu
> NVRTC log:
./var/cuda/tmp_1784724931385.cu(68): error: expected an expression
          F16_t result = (F16_t){(ha->value + hb->value)};
                                ^

1 error detected in the compilation of "./var/cuda/tmp_1784724931385.cu".

nvrtc might be more restricted compared to nvcc. Happy to work on a work-around that we can improve HAT.

[xyyNicole]

Thanks, @jjfumero, for investigating and providing these details. Sorry that I did not run the full CUDA test suite before submitting the patch; I will do so for future changes.

[xyyNicole]

The failures appear to come from generated CUDA code that relies on extensions accepted by NVCC but rejected by NVRTC. Although both compile CUDA C++, their frontends do not accept exactly the same set of extensions. In particular, the generated $ identifiers are non-standard, and (F16_t){...} uses C99 compound-literal syntax rather than CUDA C++ initialization syntax.

I do not see an NVRTC option that would reliably bypass either issue, so fixing the generated CUDA syntax seems to be the appropriate solution. Do you already plan to address these code-generation issues? If not, I can work on them in a follow-up.

[xyyNicole]

BTW, I am fine with either merging this opt-in NVRTC support first and fixing the codegen compatibility separately, or waiting until the affected tests pass, depending on your preference ;)

[jjfumero]

Thanks @xyyNicole for the feedback. So, chatting with Gary, we think the following:

1. Given that NVRTC does not consume the same extensions as CUDA C++, we expect the changes to only affect the CUDA codegen.
2. Ideally those changes could be also implemented in the OpenCL C99 codegen, making the HAT codegen easier to manage. We want to reuse the same codegen patterns as much as possible.
   a. So specifically I suspect swapping out the '$' -> '_' would work just fine for OpenCL, nvcc and nvrtc.
   b. Its also likely we can come up with a generated pattern that satisfies initializing struct fields (F16_t) that would work for OpenCL
3. To move forward this PR I suggest adapting the codegen to pass nvrtc. Although this PR will be a bit bigger, this will give us more confidence and insights to provide common patters across different backends.

Question: given these differences between nvcc and nvrtc, does the quality of the generated code changes? Wondering about performance.

Note: The two discovered issues I found were just by checking a few tests. Probably there are also other minor differences to cover.

[xyyNicole]

Thanks, @jjfumero. I agree with this approach and will address the codegen issues in this PR, while preserving common patterns across CUDA and OpenCL where practical.

As you noted, the required changes may go beyond the two issues identified so far. I will run the full CUDA unit-test suite with NVRTC and also review the generated source paths to make the CUDA codegen consistently produce syntax
accepted by NVRTC.

For the same CUDA C++ source, target architecture, and optimization settings, NVCC and NVRTC should generally produce kernels with similar runtime performance because they share many compilation stages, although identical results are not guaranteed. Once the codegen fixes are complete, we can compare both compilation time and kernel runtime performance.

One expectation to set: my progress on the codegen work may be incremental, as my bandwidth is currently limited. Thanks for your understanding.

[jjfumero]

Thank you for the update @xyyNicole.

One expectation to set: my progress on the codegen work may be incremental, as my bandwidth is currently limited. Thanks for your understanding.

Sure, no problem. Please, don't hesitate to contact us if you have questions for the codegen.