Repository: pytorch
Issue: Value of torch.backends.cudnn.benchmark Baked into JIT-Traced Modules ( 150x slowdown on ConvTranspose2d() ) [jit] [libtorch] [cudnn] 
Labels: oncall: jit,triaged

## 🐛 Bug

If you trace a module with `torch.jit.trace(...)` and load that script module in C++ via LibTorch, the resulting behavior in C++ depends on whether or not the `torch.backends.cudnn.benchmark` flag was set. Calls to `at::globalContext().setBenchmarkCuDNN(true/false)` from the C++ API at runtime appear to have no effect.

## To Reproduce

**NOTE**: I was not able to verify this issue still exists on the latest nightly (20190402) because it appears the latest nightly (at least on Windows) cannot run JIT-traced models. Even the simplest model gives the following error:

```
INVALID_ARGUMENT:: Cannot find field. (deserialize at ..\torch\csrc\jit\import.cpp:108)
(no backtrace available)
```

1) Run the python script below: `python test.py 0` or `python test.py 1`
2) Compile + run the C++ code below.
3) Observe:
  a) Average time per call. I see ~0.8ms in the python script and either ~0.8 or ~120ms in C++ depending on the flag used in python. In either case, C++ sets benchmarking ON. (GTX 1080)
  b) Kernel run by CuDNN. w/either setting of the flag, the python code runs `cudnn::detail::dgrad_engine<...>`. With the flag ON, it runs `cudnn::detail::dgrad2d_alg1_1<...>` once (taking ~120ms) and then chooses the faster `dgrad_engine`. If the flag was ON in python, C++ also chooses `dgrad_engine` but if the flag was OFF in python, it always chooses `dgrad2d_alg1_1` regardless of the flag setting in C++.

I observed the choice of kernel using `nvprof python test.py 0/1`.

Python Script (`test.py`):
```python
import sys
import time

import torch as th
th.backends.cudnn.benchmark = bool(int(sys.argv[1]))

mod = th.nn.ConvTranspose2d(8, 3, 4, 2, 1).cuda()
inp = th.zeros(1, 8, 512, 512).cuda()

mod(inp); mod(inp); mod(inp)

smod = th.jit.trace(mod, (inp,), check_trace=False)
smod.save("smod.ptj")

N = 1000
th.cuda.synchronize()
start = time.time()
for _ in range(N):
    mod(inp)
    th.cuda.synchronize()
end = time.time()
print("Time (ms):", 1000*(end-start)/N)
```

C++ Code:
```c++
#include <chrono>
#include <iostream>

#include <c10/cuda/CUDAGuard.h>
#include <torch/script.h>
#include <torch/torch.h>

#include <cuda_runtime_api.h>

int main() {
  at::globalContext().setBenchmarkCuDNN(true);
  auto nograd = torch::NoGradGuard();

  try {
    auto mod = torch::jit::load("smod.ptj");
    mod->to(torch::kCUDA);
    torch::jit::Stack input_stack = {torch::zeros({1, 8, 512, 512}, torch::kCUDA)};

    mod->forward(input_stack);
    mod->forward(input_stack);
    mod->forward(input_stack);

    const int N = 100;
    cudaDeviceSynchronize();
    const auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < N; ++i) {
      mod->forward(input_stack);
      cudaDeviceSynchronize();
    }
    const auto end = std::chrono::high_resolution_clock::now();
    const float elapsed = std::chrono::duration<float, std::milli>(end - start).count() / N;
    std::cout << "Time (ms): " << elapsed << std::endl;
  } catch (c10::Error e) {
    std::cerr << e.what() << std::endl;
    return 1;
  }
  return 0;
}
```

## Expected Behavior

I would expect that either: 1) the C++ setting of `at::globalContext().setBenchmarkCuDNN(true)` should be respected (choosing the correct algorithm) or 2) at least print a warning that it is being overridden by the value of the flag at trace time.

## Additional Info
I printed the JIT graphs generated with benchmarking ON/OFF and got the following with the flag OFF:
```
graph(%input : Float(1, 8, 512, 512),
      %1 : Float(8, 3, 4, 4),
      %2 : Float(3)):
  %3 : int = prim::Constant[value=2](), scope: ConvTranspose2d
  %4 : int = prim::Constant[value=2](), scope: ConvTranspose2d
  %5 : int[] = prim::ListConstruct(%3, %4), scope: ConvTranspose2d
  %6 : int = prim::Constant[value=1](), scope: ConvTranspose2d
  %7 : int = prim::Constant[value=1](), scope: ConvTranspose2d
  %8 : int[] = prim::ListConstruct(%6, %7), scope: ConvTranspose2d
  %9 : int = prim::Constant[value=1](), scope: ConvTranspose2d
  %10 : int = prim::Constant[value=1](), scope: ConvTranspose2d
  %11 : int[] = prim::ListConstruct(%9, %10), scope: ConvTranspose2d
  %12 : bool = prim::Constant[value=1](), scope: ConvTranspose2d
  %13 : int = prim::Constant[value=0](), scope: ConvTranspose2d
  %14 : int = prim::Constant[value=0](), scope: ConvTranspose2d
  %15 : int[] = prim::ListConstruct(%13, %14), scope: ConvTranspose2d
  %16 : int = prim::Constant[value=1](), scope: ConvTranspose2d
  %17 : bool = prim::Constant[value=0](), scope: ConvTranspose2d
  %18 : bool = prim::Constant[value=0](), scope: ConvTranspose2d
  %19 : bool = prim::Constant[value=1](), scope: ConvTranspose2d
  %20 : Float(1, 3, 1024, 1024) = aten::_convolution(%input, %1, %2, %5, %8, %11, %12, %15, %16, %17, %18, %19), scope: ConvTranspose2d
  return (%20)
```

The only change when the flag is ON is that register %17 is 1 instead of 0. I suppose this is where the "hardcoding" of the flag might be happening?

## Environment
Python code was run on Linux, C++ code was run on Windows

 - PyTorch Version (e.g., 1.0): 1.0.0.dev20190311 on linux, 2336f0ba0 on Windows
 - OS (e.g., Linux): Fedora 29, Windows 10 1809
 - How you installed PyTorch (`conda`, `pip`, source): conda (pytorch-nightly)
 - Python version: 3.7
 - CUDA/cuDNN version: CUDA 10, cuDNN 7.4.2
 - GPU models and configuration: Titan RTX (linux), GTX 1080 (windows)

cc @suo