Repository: pytorch
Issue: C++ CUDA assign existing memory to forward method.
Labels: oncall: jit

### 🐛 Describe the bug

Apologies for posting this here. I did post a question in the forum however no one replied for a couple of weeks. I would like to load a pytorch model (which Is trained in python) in a C++ application. The application has already performed device allocations and is quite memory intensive. So, I would like the model to re-use the application memory as much as possible. I am partially able to do so for the input tensors. However, I cannot do so for the output tensor. Below you can find a small indicative example:

```C++
int main(int argc, const char* argv[]) {
  if (argc != 3) {
    std::cerr << "usage: example-app <path-to-exported-script-module> <memory to allocate (MB)>\n";
    return -1;
  }
  long numElements = (1024*1024)/sizeof(double) * atoi(argv[2]);
  int id;
  size_t free, total;
  cudaGetDevice(&id);
  double *a, *b, *c;
  cudaMalloc(&a, sizeof(double)*numElements);
  cudaCheckError();
  cudaMalloc(&b, sizeof(double)*numElements);
  cudaCheckError();
  cudaMalloc(&c, sizeof(double)*numElements*4);
  cudaCheckError();
  dim3 threads(1024L);
  dim3 numBlocks((numElements + threads.x -1) / threads.x);
  init<<<numBlocks, threads>>>(a,b,numElements);
  cudaDeviceSynchronize();
  cudaCheckError();

  cudaMemGetInfo( &free, &total );
  std::cout << "(Before Tensor Creation) GPU " << id << " memory: free=" << (free/(1024.0*1024.0*1024.0)) << ", total=" << total/(1024.0*1024.0*1024.0) << std::endl;

  auto options = torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCUDA, id);
  at::Tensor a_t = torch::from_blob((double*) a, {numElements,1}, options);
  at::Tensor b_t = torch::from_blob((double*) b, {numElements,1}, options);
  
  cudaMemGetInfo( &free, &total );
  std::cout << "(After Tensor Creation) GPU " << id << " memory: free=" << (free/(1024.0*1024.0*1024.0)) << ", total=" << total/(1024.0*1024.0*1024.0) << std::endl;
  
  at::Tensor c_t = at::cat({a_t,b_t}, 1);
  at::Tensor d_t = at::reshape(c_t, {numElements,2});

  cudaMemGetInfo( &free, &total );
  std::cout << "(After Tensor Reshape) GPU " << id << " memory: free=" << (free/(1024.0*1024.0*1024.0)) << ", total=" << total/(1024.0*1024.0*1024.0) << std::endl;

  torch::jit::script::Module module;
  try {
    module = torch::jit::load(argv[1]);
    c10::Device device("cuda");
    module.to(device);
  }
  catch (const c10::Error& e) {
    return -1;
  }
  cudaMemGetInfo( &free, &total );
  std::cout << "(After Module Loading) GPU " << id << " memory: free=" << (free/(1024.0*1024.0*1024.0)) << ", total=" << total/(1024.0*1024.0*1024.0) << std::endl;
  at::Tensor out = torch::from_blob((double*) c, {numElements,4}, options);
  std::move(out) =  module.forward({d_t}).toTensor();
  cudaMemGetInfo( &free, &total );
  std::cout << "(After Forward) GPU " << id << " memory: free=" << (free/(1024.0*1024.0*1024.0)) << ", total=" << total/(1024.0*1024.0*1024.0) << std::endl;

  cudaFree(a);
  cudaFree(b);
  cudaFree(c);
  return 0;
}
```
When I compile and execute the application with the following arguments ./a.out model 1024 I get the following outcome:

```bash
(Before Tensor Creation) GPU 0 memory: free=8.95544, total=15.75
(After Tensor Creation) GPU 0 memory: free=8.95544, total=15.75
(After Tensor Reshape) GPU 0 memory: free=6.95544, total=15.75
(After Module Loading) GPU 0 memory: free=6.95544, total=15.75
(After Forward) GPU 0 memory: free=2.8363, total=15.75
```

So, before creating the tensor the free memory is approximately equal to 9GB. This is expected as each of "a,b" require 1GB of memory and "c" requires 4GB of memory. I create the tensors using the from_blob method and that works fine. I am not entirely sure why  the following code:

```C++
  at::Tensor c_t = at::cat({a_t,b_t}, 1);
  at::Tensor d_t = at::reshape(c_t, {numElements,2});
``` 
is increasing my memory usage by 2GB. I was hoping that c_t and d_t will perform the transformation in-place. 
Finally, my main question is regarding the forward method of the module. How can I make "forward" to use already allocated memory on the device side.  Is there such a mechanism?

``` C++
  at::Tensor out = torch::from_blob((double*) c, {numElements,4}, options);
  std::move(out) =  module.forward({d_t}).toTensor();
```

Thank you.


### Versions

PyTorch version: 1.8.1
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux Server release 7.9 (Maipo) (ppc64le)
GCC version: (GCC) 8.3.1 20190311 (Red Hat 8.3.1-3)
Clang version: 12.0.1
CMake version: version 3.22.4
Libc version: glibc-2.17

Python version: 3.8.6 (default, Apr 19 2022, 12:22:54)  [GCC 8.3.1 20190311 (Red Hat 8.3.1-3)] (64-bit runtime)
Python platform: Linux-4.14.0-115.35.1.3chaos.ch6a.ppc64le-ppc64le-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.1.74
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-16GB
GPU 1: Tesla V100-SXM2-16GB
GPU 2: Tesla V100-SXM2-16GB
GPU 3: Tesla V100-SXM2-16GB

Nvidia driver version: 510.47.03
cuDNN version: cudnn-8.0.4
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

Versions of relevant libraries:
[pip3] numpy==1.22.3
[pip3] torch==1.8.1
[conda] Could not collect

This was installed using spack.