Adding pin_memory kwarg to zeros, ones, empty, ... tensor constructors (#18952)
authorVitaly Fedyunin <vitalyf@fb.com>
Tue, 16 Apr 2019 17:50:48 +0000 (10:50 -0700)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Tue, 16 Apr 2019 18:06:15 +0000 (11:06 -0700)
commit1c5073fb4bb7fdb73b4137e4fcf9bc7345ab4b9a
tree41fb934e01774bdc1c4c59a991f950c82d96e59e
parent31686805f2c31073168a7710da04d7cbe542fe13
Adding pin_memory kwarg to zeros, ones, empty, ... tensor constructors (#18952)

Summary:
Make it possible to construct a pinned memory tensor without creating a storage first and without calling pin_memory() function. It is also faster, as copy operation is unnecessary.

Supported functions:
```python
torch.rand_like(t, pin_memory=True)
torch.randn_like(t, pin_memory=True)
torch.empty_like(t, pin_memory=True)
torch.full_like(t, 4, pin_memory=True)
torch.zeros_like(t, pin_memory=True)
torch.ones_like(t, pin_memory=True)
torch.tensor([10,11], pin_memory=True)
torch.randn(3, 5, pin_memory=True)
torch.rand(3, pin_memory=True)
torch.zeros(3, pin_memory=True)
torch.randperm(3, pin_memory=True)
torch.empty(6, pin_memory=True)
torch.ones(6, pin_memory=True)
torch.eye(6, pin_memory=True)
torch.arange(3, 5, pin_memory=True)
```

Part of the bigger: `Remove Storage` plan.

Now compatible with both torch scripts:
 `  _1 = torch.zeros([10], dtype=6, layout=0, device=torch.device("cpu"), pin_memory=False)`
and
`  _1 = torch.zeros([10], dtype=6, layout=0, device=torch.device("cpu"))`

Same checked for all similar functions `rand_like`, `empty_like` and others

It is fixed version of #18455
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18952

Differential Revision: D14801792

Pulled By: VitalyFedyunin

fbshipit-source-id: 8dbc61078ff7a637d0ecdb95d4e98f704d5450ba
18 files changed:
aten/src/ATen/native/TensorFactories.cpp
aten/src/ATen/native/cuda/TensorFactories.cu
aten/src/ATen/native/native_functions.yaml
aten/src/ATen/native/sparse/SparseTensor.cpp
aten/src/ATen/native_parse.py
test/test_torch.py
tools/autograd/gen_python_functions.py
tools/autograd/templates/python_torch_functions.cpp
tools/jit/gen_jit_dispatch.py
tools/pyi/gen_pyi.py
torch/_tensor_docs.py
torch/_torch_docs.py
torch/csrc/jit/ir.cpp
torch/csrc/jit/operator.cpp
torch/csrc/jit/passes/shape_analysis.cpp
torch/csrc/jit/tracer.cpp
torch/csrc/utils/tensor_new.cpp
torch/onnx/symbolic.py