Repository: pytorch
Issue: torch.solve on Jetson is slower than humans
Labels: module: build,triaged,module: linear algebra

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

The following code took more than 13 minutes to run on a Nvidia AGX Xavier 32GB devkit:

```
import torch

A = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
B = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

A = A.cuda()
B = B.cuda()

torch.solve(B, A)

```

During the 13 minutes, the GPU appears to be idle, and one of the CPU cores runs at 100%.

Curiously, this code only takes long on the first run. After that, some CUDA "compute cache" are generated, and the running time drops to a reasonable level. This problem can be reproduced in a Docker container: whenever the container is restarted, the above takes a long time to run.


## To Reproduce

Steps to reproduce the behavior:

1. Build PyTorch with MAGMA with the following Docker image. The compiling follows the instructions from https://forums.developer.nvidia.com/t/pytorch-for-jetson-version-1-7-0-now-available/72048, with added MAGMA support. 

```
FROM nvcr.io/nvidia/l4t-base:r32.4.4

RUN apt-get update \
    && apt-get install -y python3-pip cmake libopenblas-dev libssl-dev git gfortran

RUN pip3 install cython scikit-build

RUN pip3 install cmake

RUN pip3 install ninja

RUN wget http://icl.utk.edu/projectsfiles/magma/downloads/magma-2.5.4.tar.gz \
    && tar xf magma-2.5.4.tar.gz \
    && cd magma-2.5.4 \
    && mkdir build \
    && cd build \
    && cmake .. \
    && make -j8 \
    && make install

RUN mkdir /pytorch \
    && cd /pytorch \
    && git clone --recurse-submodules -j8 -b v1.7.0 https://github.com/pytorch/pytorch v1.7.0 \
    && cd /pytorch/v1.7.0 \
    && wget https://gist.githubusercontent.com/dusty-nv/ce51796085178e1f38e3c6a1663a93a1/raw/9d7261584a7482e7cc0fcb08a4a232c6d023f812/pytorch-1.7-jetpack-4.4.1.patch \
    && git apply pytorch-1.7-jetpack-4.4.1.patch \
    && pip3 install -r requirements.txt

ENV USE_NCCL=0
ENV USE_DISTRIBUTED=0
ENV USE_QNNPACK=0
ENV USE_PYTORCH_QNNPACK=0
ENV TORCH_CUDA_ARCH_LIST="5.3;6.2;7.2"
ENV PYTORCH_BUILD_VERSION=1.7.0
ENV PYTORCH_BUILD_NUMBER=0

RUN cd /pytorch/v1.7.0 \
    && python3 setup.py bdist_wheel \
    && cd dist \
    && pip3 install torch-1.7.0-cp36-cp36m-linux_aarch64.whl

WORKDIR /pytorch/v1.7.0

CMD ["bash"]

```

1. Run the above Docker image, and execute the code mentioned above with python3.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

 - PyTorch Version (e.g., 1.0): 1.7.0 (and also 1.7.1)
 - OS (e.g., Linux): Jetpack 4.4.1
 - How you installed PyTorch (`conda`, `pip`, source): source
 - Build command you used (if compiling from source): see above
 - Python version:
 - CUDA/cuDNN version: 10.2
 - GPU models and configuration: Volta
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->


cc @malfet @seemethere @walterddr @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @VitalyFedyunin @ngimel