Repository: pytorch
Issue: caffe2: cudaHostRegister() and mbind()-related test failures on ppc64le
Labels: caffe2

## 🐛 Bug

## Environment

```
PyTorch version: 1.0.0a0+7998997 (with some local changes)
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Red Hat Enterprise Linux Server 7.5 (Maipo)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-36)
CMake version: Could not collect

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
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: 410.72
cuDNN version: Probably one of the following:
/usr/local/cuda-10.0/targets/ppc64le-linux/lib/libcudnn.so.7.3.1
/usr/local/cuda-10.0/targets/ppc64le-linux/lib/libcudnn_static.a
```

## Additional context

Building on ppc64le systems with CUDA / NVIDIA GPUs, we're seeing
failures in various caffe2 tests, for example:

- `TestGradientsAccumulationWithPassThroughGradients::testAccumulationRuns`:

```
E       RuntimeError: [enforce fail at context_gpu.h:375] error == cudaSuccess. 61 vs 0. Error at: /opt/anaconda2/conda-bld/pytorch_1547774768751/work/caffe2/core/context_gpu.h:375: part or all of the requested memory range is already mapped
E       frame #0: c10::ThrowEnforceNotMet(char const*, int, char const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, void const*) + 0xb0 (0x7fff76a67510 in /opt/anaconda2/envs/6518b/lib/python2.7/site-packages//torch/lib/libc10.so)
...
```

- `brew_test.py::BrewTest::test_arg_scope_single`

```
...
>       workspace.FeedBlob("x", X)
...
>       return blob._feed(arg, device_option)
E       RuntimeError: [enforce fail at numa.cc:85] mbind( (void*)page_start_ptr, size + offset, MPOL_BIND, &mask, sizeof(mask) * 8, MPOL_MF_MOVE | MPOL_MF_STRICT) == 0. Could not move memory to a NUMA node
E       frame #0: c10::ThrowEnforceNotMet(char const*, int, char const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, void const*) + 0xb0 (0x3fffa2167510 in /opt/anaconda/envs/pytorch-env/lib/python2.7/site-packages//torch/lib/libc10.so)
E       frame #1: caffe2::NUMAMove(void*, unsigned long, int) + 0x210 (0x3fffae2485b0 in /opt/anaconda/envs/pytorch-env/lib/python2.7/site-packages//torch/lib/libcaffe2.so)
...
```

Other tests fail simlarly:

```
DataParallelModelTest.test_equiv
TestCloneNet.testPartialClone
TestCloneNet.testPartialClone
TestCrfDecode.test_crf_viterbi
TestOperatorTraceback.test_async_exception_handling
TestOperatorTraceback.test_operator_runtime_traceback
```

We believe the cause is in `PinnedCPUAllocator` in `caffe2/core/context_gpu.h`.

If NUMA is enabled, then `PinnedCPUAllocator`'s `allocate()` function
- calls to `DefaultCPUAllocator` to allocate a `DataPtr`,
- calls `cudaHostRegister()` on the actual memory pointer
- returns `DefaultCPUAllocator`'s `DataPtr` to the caller

`PinnedCPUAllocator` also provides a `Delete()` function that _intends_
to call `cudaHostUnregister()` when the region is freed. But in the
NUMA / `cudaHostRegister()` case, `PinnedCPUAllocator`'s `Delete()` is
never called because the `DataPtr` was constructed by
`DefaultCPUAllocator` and so refers to _that_ `Delete()` instead.

The end result is that `cudaHostRegister()` is called on memory
regions, but there are never any corresponding `cudaHostUnregister()` calls.

In addition to registering the host memory with CUDA, `cudaHostRegister()`
also pins / locks the memory.

So we think this explains both failure modes shown above. If memory
regions are allocated / freed / re-allocated, then either:

- `DefaultCPUAllocator`'s `mbind()` call will fail (prevented by the
pin/lock if the reallocation is on a different NUMA node from the
original allocation), or

-  the `cudaHostRegister()` call with fail (if `mbind()` didn't and the
region was previously `cudaHostRegister()`'d).


Not sure of the best way to fix this. Maybe any of:

- extend `at::DataPtr` with something like a `PinnedDataPtr` that just
includes an extra field to hang the underlying `DataPtr` from the base
allocator. Have `PinnedCPUAllocator`'s `allocate()` always construct
with a reference to own `Delete()`, and return that to callers

- add an `allocate()` variant to `DefaultCPUAllocator` that allows a
Delete function reference to be passed in

- add an `allocate()` variant to `DefaultCPUAllocator` that just returns
the naked memory pointer, leaving `PinnedCPUAllocator` to construct
the `DataPtr` as in the non-NUMA case

- thread some `set_deleter()` functionality through the DataPtr class
hierarchy (but that seems not well-supported by the base
`std::unique_ptr` class, and you'd really want deleters to nest)