torch.jit.script'd function very slow on first invocation on latest nightly

[cpuhrsch]

🐛 Describe the bug

It takes about a minute to run this function for the first time. It takes only a second if it's running on a version of PyTorch built from source.

To reproduce run the code in this gist. All credit to finding this to @Linux-cpp-lisp. I suspect this is an environment issue, i.e. an old version of that we ship as a nightly vs. a newer version I'm using locally.

Clearly this prevents the optimization from being useful.

Versions

The nightly in question here is 1.12.0.dev20220415-py3.9_cuda11.3_cudnn8.3.2_0

My local environment is

Collecting environment information...
PyTorch version: 1.12.0a0+git075974e
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 10.0.1
CMake version: version 3.22.3
Libc version: glibc-2.27

Python version: 3.9.5 (default, Jun  4 2021, 12:28:51)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1051-aws-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: A100-SXM4-40GB
Nvidia driver version: 450.119.03
cuDNN version: Probably one of the following:
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.812
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.3
[pip3] torch==1.12.0a0+git075974e
[pip3] torch2trt==0.3.0
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-include               2021.2.0           h06a4308_296
[conda] mkl-random                1.2.1                    pypi_0    pypi
[conda] mkl-service               2.3.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.0            py39h42c9631_2
[conda] mkl_random                1.2.2            py39h51133e4_0
[conda] mypy                      0.812              pyhd8ed1ab_0    conda-forge
[conda] mypy_extensions           0.4.3            py39h06a4308_0
[conda] numpy                     1.20.3                   pypi_0    pypi
[conda] numpy-base                1.20.2           py39hfae3a4d_0
[conda] torch                     1.11.0                   pypi_0    pypi

One notable difference is CUDA 11.1 locally vs 11.3 in the nightlies (note that the gist doesn't use CUDA).

[Linux-cpp-lisp]

In case it's useful, some longer and shorter functions of the same class can be found here: https://github.com/Linux-cpp-lisp/e3nn/blob/spharm-no-controlflow/e3nn/o3/_spherical_harmonics_generated.py.

[ngimel]

It still takes about a minute on my nightly build if the input tensors are created on cuda, and it hangs (I killed it after 10 minutes) if nvfuser is used, cc @jjsjann123

[ngimel]

It's possible that nightlies are built with llvm so they try to compile something on cpu, whereas local build doesn't have llvm

[davidberard98]

@ngimel to clarify - this happens with both NNC and NVFuser, is that correct?

[ngimel]

1. The original gist uses cpu only, and the first or second run of the function is slow on nightly and fast on a local build, I suspect because nightly tries to compile something for cpu, and my and Christian's local build doesn't
2. If gist is changed to use GPU, then with nnc it compiles for about a minute the first time, then runs fast
3. If gist is changed to use GPU and nvfuser, it is compiling (?) for a very long time or it hangs, I killed it after 10 minutes.

[Linux-cpp-lisp]

additional context on that is that when we run unit tests on this set of compiled functions in our GitHub Actions CI (CPU, PyTorch 1.10 and 1.11) it hangs ~forever with NNC; at some point it just gets stuck and then killed 6 hours later:
https://github.com/e3nn/e3nn/runs/6014157098?check_suite_focus=true
At no point in this CI set is NVFuser ever enabled or CUDA available.

[davidberard98]

In NNC I narrowed this down to the IRMutators; they try to modify the graph by recursing into each input to a node and copying that; but for a dag, it will end up visiting some nodes more than once

  -> x -      -> x -      -> x -      -> x -
 /      \    /      \    /      \    /      \
x        -> x        -> x        -> x        -> x
 \      /    \      /    \      /    \      /
  -> x -      -> x -      -> x -      -> x -

Simpler repro:

import torch

def fn(x, y):
    a = x + y
    for i in range(10):  # <- exponential in this number
        x = torch.cos(a)
        y = torch.sin(a)
        a = x + y
    return a

inp = [torch.randn(13) for _ in range(2)]
fn_s = torch.jit.trace(fn, inp)

print(fn_s.graph)

x, y = [torch.randn(13) for _ in range(2)]
print("hi")
a = fn_s(x, y)
print("hi2")
a = fn_s(x, y)
print("hi3")
a = fn_s(x, y)
print("hi4")

cc @ZolotukhinM @jjsjann123

[eellison]

cc @ZolotukhinM, thoughts here ? maybe this would be a good issue for intel to pick up?

[davidberard98]

https://github.com/pytorch/pytorch/compare/master...davidberard98:mutator-caching?expand=1 - would something like this work?

[eellison]

This might be solved by off-thread compilation as well... tbd on implementation in nvfuser or jit

[jjsjann123]

cc'ing @zasdfgbnm

I believe you are looking at the compilation time issue for nvfuser. David's small example doesn't repro with nvfuser unfortunately. But I did try with the original script with nvfuser turned on and observed the hang.

[cpuhrsch]

@eellison @jjsjann123 @davidberard98 - Thanks for taking a look! Given that this is more complex than just an environment issue how difficult will it be to resolve this? Any updates here?