Fix NVRTC versioning for CUDA 11.X (X>=3), CUDA 12 and later

[zasdfgbnm]

NVRTC versioning has changed starting 11.3, and will change again for CUDA 12.X. See comment in code for detail. As a result, jit on CUDA 11.3 is broken.

Also, the error message is misleading: When both libname and alt_libname are non-empty, the error message is only reporting alt_libname, it should report both.

To reproduce the error, you can use:

import torch

torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_profiling_executor(False)
torch._C._jit_override_can_fuse_on_cpu(True)
torch._C._jit_override_can_fuse_on_gpu(True)

@torch.jit.script
def jit_relu_dropout(x, prob) :
    # type: (Tensor, float) -> Tensor
    x = torch.nn.functional.relu(x)
    x = torch.nn.functional.dropout(x, p=prob, training=True)
    return x

x = torch.randn((64, 40, 12, 1024), device="cuda:0", dtype=torch.float16, requires_grad=True)
y = jit_relu_dropout(x, 0.5)

with CUDA 11.3, and you will see

Traceback (most recent call last):
  File "/home/gaoxiang/misc/nvrtc-failure.py", line 16, in <module>
    y = jit_relu_dropout(x, 0.5)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: Error in dlopen or dlsym: libnvrtc-8aa72235.so.11.3: cannot open shared object file: No such file or directory

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 18b3b07 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[zasdfgbnm]

cc: @ptrblck @mcarilli

[malfet]

Thank you for the fix, but can you please rework it (using preprocessor or metaprogramming) to keep libname static for performance reasons? (otherwise libname would be computed every time this function is called)
Or initialize explicitly using std::once construct

[zasdfgbnm]

@malfet Done

[codecov]

Codecov Report

Merging #57204 (18b3b07) into master (4cb534f) will decrease coverage by 0.01%.
The diff coverage is 0.00%.

@@            Coverage Diff             @@
##           master   #57204      +/-   ##
==========================================
- Coverage   77.58%   77.56%   -0.02%     
==========================================
  Files        1952     1952              
  Lines      194466   194468       +2     
==========================================
- Hits       150868   150831      -37     
- Misses      43598    43637      +39     

[facebook-github-bot]

@malfet has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@malfet merged this pull request in 42b3fc2.