Repository: pytorch
Issue: Torch function runtime seemingly dependent on scipy call
Labels: needs reproduction,module: performance,triaged,module: linear algebra

### 🐛 Describe the bug

From https://discuss.pytorch.org/t/torch-function-runtime-dependent-on-scipy-call/139483:

I've noticed that certain PyTorch functions run slower when I make calls to `scipy.linalg.solve_triangular`. 

I have managed to reproduce this issue with the following script:
```
import time

import numpy as np
import torch
import scipy.linalg

def torch_fn(input_size):
    with torch.no_grad():
        a = torch.randn(input_size, 4, 4, dtype=torch.float64)
        b = torch.randn(input_size, 4, 5, dtype=torch.float64)
        begin = time.perf_counter()
        c = torch.bmm(a, b)
        return time.perf_counter() - begin


def scipy_fn(input_size):
    a = np.asarray([[1, 0], [1, 1]], dtype=np.float64)
    b = np.random.randn(input_size, 2).astype(np.float64)
    begin = time.perf_counter()
    c = scipy.linalg.solve_triangular(a, b.T, lower=True).T
    return time.perf_counter() - begin


def avg_time(torch_input_size, scipy_input_size):
    tot = 0
    for _ in range(100):
        tot += torch_fn(torch_input_size)
        if scipy_input_size > 0:
            scipy_fn(scipy_input_size)
    return tot / 100


if __name__ == '__main__':
    print(avg_time(5, 0))
    print(avg_time(5, 3))
    print(avg_time(5, 75))
    print(avg_time(30, 0))
    print(avg_time(30, 3))
    print(avg_time(30, 75))
```

The output I get is
```
0.0002978778630495071
8.901087567210197e-06
0.006054648053832352
0.0008917823247611523
8.631939999759198e-06
0.004081738111563027
```

So it seems like making a `scipy.linalg.solve_triangular` call with sufficiently large inputs results in `torch.bmm` running much slower. The runtime difference may not seem too significant, but in our codebase, we call other torch functions and we observe much larger differences (e.g. 30ms vs. < 0.5ms if we call/do not call scipy).

I suspect the issue may have to do with the underlying linear algebra library that scipy calls (in the above case, the `trtrs` function in lapack/blas). Beyond this, I can't explain why this runtime discrepancy occurs.

It was also determined that a different ordering also results in slightly different timings:
```
if __name__ == '__main__':
    print(avg_time(30, 75))
    print(avg_time(30, 3))
    print(avg_time(30, 0))
    print(avg_time(5, 75))
    print(avg_time(5, 3))
    print(avg_time(5, 0))
```
```
0.0058318315260112285
0.0006946996226906777
7.811854593455792e-06
0.005560100376605988
0.0008387839421629906
7.02753197401762e-06
```


### Versions

```
PyTorch version: 1.9.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.7 (default, Nov  8 2021, 18:36:03)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.11.0-38-generic-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce RTX 2080 Ti
Nvidia driver version: 460.91.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.812
[pip3] mypy-extensions==0.4.3
[pip3] mypy-protobuf==2.8
[pip3] numpy==1.20.3
[pip3] pytorch3d==0.6.0+cu111.1
[pip3] torch==1.9.0+cu111
[pip3] torch-with-deps==1.0.3
[pip3] torchvision==0.10.0+cu111
[conda] Could not collect
```
Possibly relevant:
```
>>> scipy.__version__
'1.6.3'
```


cc @VitalyFedyunin @ngimel @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @Lezcano