PyTorch tensor / Faiss index interoperability (#1484)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/1484

This diff allows for native usage of PyTorch tensors for Faiss indexes on both CPU and GPU. It is currently only implemented in this diff for things that inherit from `faiss.Index`, which covers the non-binary indices, and it patches the same functions on `faiss.Index` that were also covered by `__init__.py` for numpy interoperability.

There must be uniformity among the inputs: if any array input is a Torch tensor, then all array inputs must be Torch tensors. Similarly, if any array input is a numpy ndarray, then all array inputs must be numpy ndarrays.

If `faiss.contrib.torch_utils` is imported, it ensures that `import faiss` has already been performed to patch all of the functions using the base `__init__.py` numpy wrappers, and then patches the following functions again:

```
add
add_with_ids
assign
train
search
remove_ids
reconstruct
reconstruct_n
range_search
update_vectors
search_and_reconstruct
sa_encode
sa_decode
```

to allow usage of PyTorch CPU tensors, and additionally PyTorch GPU tensors if the index being used is on the GPU.

numpy functionality is still available when `faiss.contrib.torch_utils` is imported; we pass through to the original patched numpy function when we detect numpy inputs.

In addition, to allow for better (asynchronous) GPU usage without requiring the CPU to be involved, all of these functions which construct tensors/arrays for output now take optional arguments for storage (numpy or torch.Tensor) to be provided that will contain the output data. `range_search` is the only exception to this, as the size of the output data is indeterminate. The eventual GPU implementation will likely require the user to provide a maximum cap on the output size, and allow that to be passed instead. If the optional pre-allocated output values are presented by the user, they are used; otherwise, new return ndarray / Tensors are constructed as before and used for the return. If this feature were not provided on the GPU, then every execution would be completely serial as we would depend upon the CPU to allocate GPU memory before every operation. Instead, now this can function much like NN graph execution on the GPU, assuming that all of the data requirements are pre-allocated, so the execution will run at the full speed of the GPU and not be stalled sequentially launching kernels.

This diff also exposes the `GpuResources` shared_ptr object owned by a GPU index. This is required for pytorch GPU so that we can perform proper stream ordering in Faiss with respect to the current pytorch stream. So, Faiss indices now perform more or less as any NN operation in Torch does.

Note, however, that a Faiss index has its own setting on current device, and if the pytorch GPU tensor inputs are resident on a different device than what the Faiss index expects, a cross-device copy will be initiated. I may choose to make this an error in the future and require matching device to device.

This diff also found a bug when passing GPU data directly to `train()` for `GpuIndexIVFFlat` and `GpuIndexIVFScalarQuantizer`, as I guess we never tested passing GPU data directly to these functions before. `GpuIndexIVFPQ` was doing the right thing however.

The assign function is now also implemented on the GPU as well, and is now marked `const` to be in line with the `search` function.

Also added better checking of non-contiguous inputs for both Torch tensors and numpy ndarrays.

Updated the `knn_gpu` function with a base implementation always present that allows for usage of numpy arrays, which is overridden when `torch_utils` is imported to allow torch usage. This supports row/column major layout, float32/float16 data and int64/int32 indices for both numpy and torch.

Reviewed By: mdouze

Differential Revision: D24299400

fbshipit-source-id: b4f117b9c120bd1ad83e7702087051ab7b303b29

.circleci/config.yml

@@ -35,7 +35,8 @@ jobs:
           command: |
             pip3 install pytest
             export PYTHONPATH="$(ls -d ./build/faiss/python/build/lib*/)"
-            pytest --junitxml=test-results/pytest/results.xml ./tests
+            pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
+            pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py
       - store_test_results:
           path: test-results
 
@@ -81,7 +82,8 @@ jobs:
           command: |
             pip3 install pytest torch
             export PYTHONPATH="$(ls -d ./build/faiss/python/build/lib*/)"
-            pytest --junitxml=test-results/pytest/results.xml ./tests
+            pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
+            pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py
       - store_test_results:
           path: test-results
 
@@ -111,7 +113,8 @@ jobs:
           command: |
             docker build -t faiss -f .circleci/Dockerfile.faiss_gpu .
             docker run --gpus all faiss make -C build test
-            docker run --gpus all faiss sh -c '(cd build/faiss/python; python3 setup.py install) && python3 -m unittest discover faiss/gpu/test'
+            docker run --gpus all faiss sh -c '(cd build/faiss/python; python3 setup.py install) && python3 -m unittest discover -s faiss/gpu/test -p "test_*"'
+            docker run --gpus all faiss sh -c '(cd build/faiss/python; python3 setup.py install) && python3 -m unittest discover -s faiss/gpu/test -p "torch_*.py"'
 
   deploy_linux:
     parameters:

conda/faiss/meta.yaml

@@ -75,7 +75,9 @@ outputs:
       requires:
         - numpy
         - scipy
+        - pytorch
       commands:
-        - python -m unittest discover tests/
+        - python -m unittest discover -s tests -p "test_*"
+        - python -m unittest discover -s tests -p "torch_*"
       source_files:
         - tests/

contrib/README.md

@@ -31,12 +31,13 @@ See [On-disk storage](https://github.com/facebookresearch/faiss/wiki/Indexes-tha
 Computes the ground-truth search results for a dataset that possibly does not fit in RAM. Uses GPU if available.
 Tested in `tests/test_contrib.TestComputeGT`
 
-### gpu.py
+### torch_utils.py
 
-(requires GPU Faiss)
+Interoperability functions for pytorch and Faiss: Importing this will allow pytorch Tensors (CPU or GPU) to be used as arguments to Faiss indexes and other functions. Torch GPU tensors can only be used with Faiss GPU indexes. If this is imported with a package that supports Faiss GPU, the necessary stream synchronization with the current pytorch stream will be automatically performed.
 
-Interoperability functions for pytorch and Faiss: pass GPU data without copying back to CPU.
-Tested in `gpu/test/test_pytorch_faiss`
+Numpy ndarrays can continue to be used in the Faiss python interface after importing this file. All arguments must be uniformly either numpy ndarrays or Torch tensors; no mixing is allowed.
+
+Tested in `tests/test_contrib_torch.py` (CPU) and `gpu/test/test_contrib_torch_gpu.py` (GPU).
 
 ### datasets.py
 

contrib/exhaustive_search.py

@@ -72,78 +72,3 @@ def knn(xq, xb, k, distance_type=faiss.METRIC_L2):
             d, nq, nb, heaps
         )
     return D, I
-
-
-def knn_gpu(res, xb, xq, k, D=None, I=None, metric=faiss.METRIC_L2):
-    """Brute-force k-nearest neighbor on the GPU using CPU-resident numpy arrays
-    Supports float16 arrays and Fortran-order arrays.
-    """
-    if xb.ndim != 2 or xq.ndim != 2:
-        raise TypeError('xb and xq must be matrices')
-
-    nb, d = xb.shape
-    nq, d2 = xq.shape
-    if d != d2:
-        raise TypeError('xq not the same dimension as xb')
-
-    if xb.flags.c_contiguous:
-        xb_row_major = True
-    elif xb.flags.f_contiguous:
-        xb = xb.T
-        xb_row_major = False
-    else:
-        raise TypeError('xb must be either C or Fortran contiguous')
-
-    if xq.flags.c_contiguous:
-        xq_row_major = True
-    elif xq.flags.f_contiguous:
-        xq = xq.T
-        xq_row_major = False
-    else:
-        raise TypeError('xq must be either C or Fortran contiguous')
-
-    if xb.dtype == np.float32 and xq.dtype == np.float32:
-        xb_xq_type = faiss.DistanceDataType_F32
-    elif xb.dtype == np.float16 and xq.dtype == np.float16:
-        xb_xq_type = faiss.DistanceDataType_F16
-    else:
-        raise TypeError('xb and xq must both be np.float32 or np.float16')
-
-    if D is None:
-        D = np.empty((nq, k), dtype=np.float32)
-    else:
-        assert D.shape == (nq, k)
-        assert D.dtype == np.float32
-
-    if I is None:
-        I = np.empty((nq, k), dtype=np.int64)
-        indices_type = faiss.IndicesDataType_I64
-    else:
-        assert I.shape == (nq, k)
-        if I.dtype == np.int64:
-            indices_type = faiss.IndicesDataType_I64
-        elif I.dtype == np.int32:
-            indices_type = faiss.IndicesDataType_I32
-        else:
-            raise TypeError('I must be either np.int64 or np.int32')
-
-    print('row major', xb_row_major, xq_row_major)
-
-    args = faiss.GpuDistanceParams()
-    args.metric = metric
-    args.k = k
-    args.dims = d
-    args.vectors = faiss.swig_ptr(xb)
-    args.vectorType = xb_xq_type
-    args.vectorsRowMajor = xb_row_major
-    args.numVectors = nb
-    args.queries = faiss.swig_ptr(xq)
-    args.queryType = xb_xq_type
-    args.queriesRowMajor = xq_row_major
-    args.numQueries = nq
-    args.outDistances = faiss.swig_ptr(D)
-    args.outIndices = faiss.swig_ptr(I)
-    args.outIndicesType = indices_type
-    faiss.bfKnn(res, args)
-
-    return D, I

contrib/pytorch_tensors.py

@@ -1,134 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-#
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-import faiss
-import torch
-import contextlib
-
-def swig_ptr_from_FloatTensor(x):
-    """ gets a Faiss SWIG pointer from a pytorch trensor (on CPU or GPU) """
-    assert x.is_contiguous()
-    assert x.dtype == torch.float32
-    return faiss.cast_integer_to_float_ptr(
-        x.storage().data_ptr() + x.storage_offset() * 4)
-
-def swig_ptr_from_LongTensor(x):
-    """ gets a Faiss SWIG pointer from a pytorch trensor (on CPU or GPU) """
-    assert x.is_contiguous()
-    assert x.dtype == torch.int64, 'dtype=%s' % x.dtype
-    return faiss.cast_integer_to_long_ptr(
-        x.storage().data_ptr() + x.storage_offset() * 8)
-
-@contextlib.contextmanager
-def using_stream(res, pytorch_stream=None):
-    r""" Creates a scoping object to make Faiss GPU use the same stream
-         as pytorch, based on torch.cuda.current_stream().
-         Or, a specific pytorch stream can be passed in as a second
-         argument, in which case we will use that stream.
-    """
-
-    if pytorch_stream is None:
-        pytorch_stream = torch.cuda.current_stream()
-
-    # This is the cudaStream_t that we wish to use
-    cuda_stream_s = faiss.cast_integer_to_cudastream_t(pytorch_stream.cuda_stream)
-
-    # So we can revert GpuResources stream state upon exit
-    prior_dev = torch.cuda.current_device()
-    prior_stream = res.getDefaultStream(torch.cuda.current_device())
-
-    res.setDefaultStream(torch.cuda.current_device(), cuda_stream_s)
-
-    # Do the user work
-    try:
-        yield
-    finally:
-        res.setDefaultStream(prior_dev, prior_stream)
-
-def search_index_pytorch(res, index, x, k, D=None, I=None):
-    """call the search function of an index with pytorch tensor I/O (CPU
-    and GPU supported)"""
-    assert x.is_contiguous()
-    n, d = x.size()
-    assert d == index.d
-
-    if D is None:
-        D = torch.empty((n, k), dtype=torch.float32, device=x.device)
-    else:
-        assert D.size() == (n, k)
-
-    if I is None:
-        I = torch.empty((n, k), dtype=torch.int64, device=x.device)
-    else:
-        assert I.size() == (n, k)
-
-    with using_stream(res):
-        xptr = swig_ptr_from_FloatTensor(x)
-        Iptr = swig_ptr_from_LongTensor(I)
-        Dptr = swig_ptr_from_FloatTensor(D)
-        index.search_c(n, xptr, k, Dptr, Iptr)
-
-    return D, I
-
-
-def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None,
-                             metric=faiss.METRIC_L2):
-    """search xq in xb, without building an index"""
-    assert xb.device == xq.device
-
-    nq, d = xq.size()
-    if xq.is_contiguous():
-        xq_row_major = True
-    elif xq.t().is_contiguous():
-        xq = xq.t()    # I initially wrote xq:t(), Lua is still haunting me :-)
-        xq_row_major = False
-    else:
-        raise TypeError('matrix should be row or column-major')
-
-    xq_ptr = swig_ptr_from_FloatTensor(xq)
-
-    nb, d2 = xb.size()
-    assert d2 == d
-    if xb.is_contiguous():
-        xb_row_major = True
-    elif xb.t().is_contiguous():
-        xb = xb.t()
-        xb_row_major = False
-    else:
-        raise TypeError('matrix should be row or column-major')
-    xb_ptr = swig_ptr_from_FloatTensor(xb)
-
-    if D is None:
-        D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
-    else:
-        assert D.shape == (nq, k)
-        assert D.device == xb.device
-
-    if I is None:
-        I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
-    else:
-        assert I.shape == (nq, k)
-        assert I.device == xb.device
-
-    D_ptr = swig_ptr_from_FloatTensor(D)
-    I_ptr = swig_ptr_from_LongTensor(I)
-
-    args = faiss.GpuDistanceParams()
-    args.metric = metric
-    args.k = k
-    args.dims = d
-    args.vectors = xb_ptr
-    args.vectorsRowMajor = xb_row_major
-    args.numVectors = nb
-    args.queries = xq_ptr
-    args.queriesRowMajor = xq_row_major
-    args.numQueries = nq
-    args.outDistances = D_ptr
-    args.outIndices = I_ptr
-
-    with using_stream(res):
-        faiss.bfKnn(res, args)
-
-    return D, I

contrib/torch_utils.py

@@ -0,0 +1,561 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import faiss
+import torch
+import contextlib
+import inspect
+import sys
+import numpy as np
+
+def swig_ptr_from_UInt8Tensor(x):
+    """ gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
+    assert x.is_contiguous()
+    assert x.dtype == torch.uint8
+    return faiss.cast_integer_to_uint8_ptr(
+        x.storage().data_ptr() + x.storage_offset())
+
+def swig_ptr_from_HalfTensor(x):
+    """ gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
+    assert x.is_contiguous()
+    assert x.dtype == torch.float16
+    # no canonical half type in C/C++
+    return faiss.cast_integer_to_void_ptr(
+        x.storage().data_ptr() + x.storage_offset() * 4)
+
+def swig_ptr_from_FloatTensor(x):
+    """ gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
+    assert x.is_contiguous()
+    assert x.dtype == torch.float32
+    return faiss.cast_integer_to_float_ptr(
+        x.storage().data_ptr() + x.storage_offset() * 4)
+
+def swig_ptr_from_IntTensor(x):
+    """ gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
+    assert x.is_contiguous()
+    assert x.dtype == torch.int32, 'dtype=%s' % x.dtype
+    return faiss.cast_integer_to_int_ptr(
+        x.storage().data_ptr() + x.storage_offset() * 8)
+
+def swig_ptr_from_IndicesTensor(x):
+    """ gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
+    assert x.is_contiguous()
+    assert x.dtype == torch.int64, 'dtype=%s' % x.dtype
+    return faiss.cast_integer_to_idx_t_ptr(
+        x.storage().data_ptr() + x.storage_offset() * 8)
+
+@contextlib.contextmanager
+def using_stream(res, pytorch_stream=None):
+    """ Creates a scoping object to make Faiss GPU use the same stream
+        as pytorch, based on torch.cuda.current_stream().
+        Or, a specific pytorch stream can be passed in as a second
+        argument, in which case we will use that stream.
+    """
+
+    if pytorch_stream is None:
+        pytorch_stream = torch.cuda.current_stream()
+
+    # This is the cudaStream_t that we wish to use
+    cuda_stream_s = faiss.cast_integer_to_cudastream_t(pytorch_stream.cuda_stream)
+
+    # So we can revert GpuResources stream state upon exit
+    prior_dev = torch.cuda.current_device()
+    prior_stream = res.getDefaultStream(torch.cuda.current_device())
+
+    res.setDefaultStream(torch.cuda.current_device(), cuda_stream_s)
+
+    # Do the user work
+    try:
+        yield
+    finally:
+        res.setDefaultStream(prior_dev, prior_stream)
+
+def torch_replace_method(the_class, name, replacement,
+                         ignore_missing=False, ignore_no_base=False):
+    try:
+        orig_method = getattr(the_class, name)
+    except AttributeError:
+        if ignore_missing:
+            return
+        raise
+    if orig_method.__name__ == 'torch_replacement_' + name:
+        # replacement was done in parent class
+        return
+
+    # We should already have the numpy replacement methods patched
+    assert ignore_no_base or (orig_method.__name__ == 'replacement_' + name)
+    setattr(the_class, name + '_numpy', orig_method)
+    setattr(the_class, name, replacement)
+
+def handle_torch_Index(the_class):
+    def torch_replacement_add(self, x):
+        if type(x) is np.ndarray:
+            # forward to faiss __init__.py base method
+            return self.add_numpy(x)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.add_c(n, x_ptr)
+        else:
+            # CPU torch
+            self.add_c(n, x_ptr)
+
+    def torch_replacement_add_with_ids(self, x, ids):
+        if type(x) is np.ndarray:
+            # forward to faiss __init__.py base method
+            return self.add_with_ids_numpy(x, ids)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        assert type(ids) is torch.Tensor
+        assert ids.shape == (n, ), 'not same number of vectors as ids'
+        ids_ptr = swig_ptr_from_IndicesTensor(ids)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.add_with_ids_c(n, x_ptr, ids_ptr)
+        else:
+            # CPU torch
+            self.add_with_ids_c(n, x_ptr, ids_ptr)
+
+    def torch_replacement_assign(self, x, k, labels=None):
+        if type(x) is np.ndarray:
+            # forward to faiss __init__.py base method
+            return self.assign_numpy(x, k, labels)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if labels is None:
+            labels = torch.empty(n, k, device=x.device, dtype=torch.int64)
+        else:
+            assert type(labels) is torch.Tensor
+            assert labels.shape == (n, k)
+        L_ptr = swig_ptr_from_IndicesTensor(labels)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.assign_c(n, x_ptr, L_ptr, k)
+        else:
+            # CPU torch
+            self.assign_c(n, x_ptr, L_ptr, k)
+
+        return labels
+
+    def torch_replacement_train(self, x):
+        if type(x) is np.ndarray:
+            # forward to faiss __init__.py base method
+            return self.train_numpy(x)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.train_c(n, x_ptr)
+        else:
+            # CPU torch
+            self.train_c(n, x_ptr)
+
+    def torch_replacement_search(self, x, k, D=None, I=None):
+        if type(x) is np.ndarray:
+            # forward to faiss __init__.py base method
+            return self.search_numpy(x, k, D, I)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if D is None:
+            D = torch.empty(n, k, device=x.device, dtype=torch.float32)
+        else:
+            assert type(D) is torch.Tensor
+            assert D.shape == (n, k)
+        D_ptr = swig_ptr_from_FloatTensor(D)
+
+        if I is None:
+            I = torch.empty(n, k, device=x.device, dtype=torch.int64)
+        else:
+            assert type(I) is torch.Tensor
+            assert I.shape == (n, k)
+        I_ptr = swig_ptr_from_IndicesTensor(I)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.search_c(n, x_ptr, k, D_ptr, I_ptr)
+        else:
+            # CPU torch
+            self.search_c(n, x_ptr, k, D_ptr, I_ptr)
+
+        return D, I
+
+    def torch_replacement_search_and_reconstruct(self, x, k, D=None, I=None, R=None):
+        if type(x) is np.ndarray:
+            # Forward to faiss __init__.py base method
+            return self.search_and_reconstruct_numpy(x, k, D, I, R)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if D is None:
+            D = torch.empty(n, k, device=x.device, dtype=torch.float32)
+        else:
+            assert type(D) is torch.Tensor
+            assert D.shape == (n, k)
+        D_ptr = swig_ptr_from_FloatTensor(D)
+
+        if I is None:
+            I = torch.empty(n, k, device=x.device, dtype=torch.int64)
+        else:
+            assert type(I) is torch.Tensor
+            assert I.shape == (n, k)
+        I_ptr = swig_ptr_from_IndicesTensor(I)
+
+        if R is None:
+            R = torch.empty(n, k, d, device=x.device, dtype=torch.float32)
+        else:
+            assert type(R) is torch.Tensor
+            assert R.shape == (n, k, d)
+        R_ptr = swig_ptr_from_FloatTensor(R)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.search_and_reconstruct_c(n, x_ptr, k, D_ptr, I_ptr, R_ptr)
+        else:
+            # CPU torch
+            self.search_and_reconstruct_c(n, x_ptr, k, D_ptr, I_ptr, R_ptr)
+
+        return D, I, R
+
+    def torch_replacement_remove_ids(self, x):
+        # Not yet implemented
+        assert type(x) is not torch.Tensor, 'remove_ids not yet implemented for torch'
+        return self.remove_ids_numpy(x)
+
+    def torch_replacement_reconstruct(self, key, x=None):
+        # No tensor inputs are required, but with importing this module, we
+        # assume that the default should be torch tensors. If we are passed a
+        # numpy array, however, assume that the user is overriding this default
+        if (x is not None) and (type(x) is np.ndarray):
+            # Forward to faiss __init__.py base method
+            return self.reconstruct_numpy(key, x)
+
+        # If the index is a CPU index, the default device is CPU, otherwise we
+        # produce a GPU tensor
+        device = torch.device('cpu')
+        if hasattr(self, 'getDevice'):
+            # same device as the index
+            device = torch.device('cuda', self.getDevice())
+
+        if x is None:
+            x = torch.empty(self.d, device=device, dtype=torch.float32)
+        else:
+            assert type(x) is torch.Tensor
+            assert x.shape == (self.d, )
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.reconstruct_c(key, x_ptr)
+        else:
+            # CPU torch
+            self.reconstruct_c(key, x_ptr)
+
+        return x
+
+    def torch_replacement_reconstruct_n(self, n0, ni, x=None):
+        # No tensor inputs are required, but with importing this module, we
+        # assume that the default should be torch tensors. If we are passed a
+        # numpy array, however, assume that the user is overriding this default
+        if (x is not None) and (type(x) is np.ndarray):
+            # Forward to faiss __init__.py base method
+            return self.reconstruct_n_numpy(n0, ni, x)
+
+        # If the index is a CPU index, the default device is CPU, otherwise we
+        # produce a GPU tensor
+        device = torch.device('cpu')
+        if hasattr(self, 'getDevice'):
+            # same device as the index
+            device = torch.device('cuda', self.getDevice())
+
+        if x is None:
+            x = torch.empty(ni, self.d, device=device, dtype=torch.float32)
+        else:
+            assert type(x) is torch.Tensor
+            assert x.shape == (ni, self.d)
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.reconstruct_n_c(n0, ni, x_ptr)
+        else:
+            # CPU torch
+            self.reconstruct_n_c(n0, ni, x_ptr)
+
+        return x
+
+    def torch_replacement_update_vectors(self, keys, x):
+        if type(keys) is np.ndarray:
+            # Forward to faiss __init__.py base method
+            return self.update_vectors_numpy(keys, x)
+
+        assert type(keys) is torch.Tensor
+        (n, ) = keys.shape
+        keys_ptr = swig_ptr_from_IndicesTensor(keys)
+
+        assert type(x) is torch.Tensor
+        assert x.shape == (n, self.d)
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.update_vectors_c(n, keys_ptr, x_ptr)
+        else:
+            # CPU torch
+            self.update_vectors_c(n, keys_ptr, x_ptr)
+
+    # Until the GPU version is implemented, we do not support pre-allocated
+    # output buffers
+    def torch_replacement_range_search(self, x, thresh):
+        if type(x) is np.ndarray:
+            # Forward to faiss __init__.py base method
+            return self.range_search_numpy(x, thresh)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        assert not x.is_cuda, 'Range search using GPU tensor not yet implemented'
+        assert not hasattr(self, 'getDevice'), 'Range search on GPU index not yet implemented'
+
+        res = faiss.RangeSearchResult(n)
+        self.range_search_c(n, x_ptr, thresh, res)
+
+        # get pointers and copy them
+        # FIXME: no rev_swig_ptr equivalent for torch.Tensor, just convert
+        # np to torch
+        # NOTE: torch does not support np.uint64, just np.int64
+        lims = torch.from_numpy(faiss.rev_swig_ptr(res.lims, n + 1).copy().astype('int64'))
+        nd = int(lims[-1])
+        D = torch.from_numpy(faiss.rev_swig_ptr(res.distances, nd).copy())
+        I = torch.from_numpy(faiss.rev_swig_ptr(res.labels, nd).copy())
+
+        return lims, D, I
+
+    def torch_replacement_sa_encode(self, x, codes=None):
+        if type(x) is np.ndarray:
+            # Forward to faiss __init__.py base method
+            return self.sa_encode_numpy(x, codes)
+
+        assert type(x) is torch.Tensor
+        n, d = x.shape
+        assert d == self.d
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if codes is None:
+            codes = torch.empty(n, self.sa_code_size(), dtype=torch.uint8)
+        else:
+            assert codes.shape == (n, self.sa_code_size())
+        codes_ptr = swig_ptr_from_UInt8Tensor(codes)
+
+        if x.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.sa_encode_c(n, x_ptr, codes_ptr)
+        else:
+            # CPU torch
+            self.sa_encode_c(n, x_ptr, codes_ptr)
+
+        return codes
+
+    def torch_replacement_sa_decode(self, codes, x=None):
+        if type(codes) is np.ndarray:
+            # Forward to faiss __init__.py base method
+            return self.sa_decode_numpy(codes, x)
+
+        assert type(codes) is torch.Tensor
+        n, cs = codes.shape
+        assert cs == self.sa_code_size()
+        codes_ptr = swig_ptr_from_UInt8Tensor(codes)
+
+        if x is None:
+            x = torch.empty(n, self.d, dtype=torch.float32)
+        else:
+            assert type(x) is torch.Tensor
+            assert x.shape == (n, self.d)
+        x_ptr = swig_ptr_from_FloatTensor(x)
+
+        if codes.is_cuda:
+            assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
+
+            # On the GPU, use proper stream ordering
+            with using_stream(self.getResources()):
+                self.sa_decode_c(n, codes_ptr, x_ptr)
+        else:
+            # CPU torch
+            self.sa_decode_c(n, codes_ptr, x_ptr)
+
+        return x
+
+
+    torch_replace_method(the_class, 'add', torch_replacement_add)
+    torch_replace_method(the_class, 'add_with_ids', torch_replacement_add_with_ids)
+    torch_replace_method(the_class, 'assign', torch_replacement_assign)
+    torch_replace_method(the_class, 'train', torch_replacement_train)
+    torch_replace_method(the_class, 'search', torch_replacement_search)
+    torch_replace_method(the_class, 'remove_ids', torch_replacement_remove_ids)
+    torch_replace_method(the_class, 'reconstruct', torch_replacement_reconstruct)
+    torch_replace_method(the_class, 'reconstruct_n', torch_replacement_reconstruct_n)
+    torch_replace_method(the_class, 'range_search', torch_replacement_range_search)
+    torch_replace_method(the_class, 'update_vectors', torch_replacement_update_vectors,
+                         ignore_missing=True)
+    torch_replace_method(the_class, 'search_and_reconstruct',
+                         torch_replacement_search_and_reconstruct, ignore_missing=True)
+    torch_replace_method(the_class, 'sa_encode', torch_replacement_sa_encode)
+    torch_replace_method(the_class, 'sa_decode', torch_replacement_sa_decode)
+
+faiss_module = sys.modules['faiss']
+
+# Re-patch anything that inherits from faiss.Index to add the torch bindings
+for symbol in dir(faiss_module):
+    obj = getattr(faiss_module, symbol)
+    if inspect.isclass(obj):
+        the_class = obj
+        if issubclass(the_class, faiss.Index):
+            handle_torch_Index(the_class)
+
+# allows torch tensor usage with bfKnn
+def torch_replacement_knn_gpu(res, xb, xq, k, D=None, I=None, metric=faiss.METRIC_L2):
+    if type(xb) is np.ndarray:
+        # Forward to faiss __init__.py base method
+        return faiss.knn_gpu_numpy(res, xb, xq, k, D, I, metric)
+
+    nb, d = xb.size()
+    if xb.is_contiguous():
+        xb_row_major = True
+    elif xb.t().is_contiguous():
+        xb = xb.t()
+        xb_row_major = False
+    else:
+        raise TypeError('matrix should be row or column-major')
+
+    if xb.dtype == torch.float32:
+        xb_type = faiss.DistanceDataType_F32
+        xb_ptr = swig_ptr_from_FloatTensor(xb)
+    elif xb.dtype == torch.float16:
+        xb_type = faiss.DistanceDataType_F16
+        xb_ptr = swig_ptr_from_HalfTensor(xb)
+    else:
+        raise TypeError('xb must be f32 or f16')
+
+    nq, d2 = xq.size()
+    assert d2 == d
+    if xq.is_contiguous():
+        xq_row_major = True
+    elif xq.t().is_contiguous():
+        xq = xq.t()
+        xq_row_major = False
+    else:
+        raise TypeError('matrix should be row or column-major')
+
+    if xq.dtype == torch.float32:
+        xq_type = faiss.DistanceDataType_F32
+        xq_ptr = swig_ptr_from_FloatTensor(xq)
+    elif xq.dtype == torch.float16:
+        xq_type = faiss.DistanceDataType_F16
+        xq_ptr = swig_ptr_from_HalfTensor(xq)
+    else:
+        raise TypeError('xq must be f32 or f16')
+
+    if D is None:
+        D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
+    else:
+        assert D.shape == (nq, k)
+        # interface takes void*, we need to check this
+        assert (D.dtype == torch.float32)
+
+    if I is None:
+        I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
+    else:
+        assert I.shape == (nq, k)
+
+    if I.dtype == torch.int64:
+        I_type = faiss.IndicesDataType_I64
+        I_ptr = swig_ptr_from_IndicesTensor(I)
+    elif I.dtype == I.dtype == torch.int32:
+        I_type = faiss.IndicesDataType_I32
+        I_ptr = swig_ptr_from_IntTensor(I)
+    else:
+        raise TypeError('I must be i64 or i32')
+
+    D_ptr = swig_ptr_from_FloatTensor(D)
+
+    args = faiss.GpuDistanceParams()
+    args.metric = metric
+    args.k = k
+    args.dims = d
+    args.vectors = xb_ptr
+    args.vectorsRowMajor = xb_row_major
+    args.vectorType = xb_type
+    args.numVectors = nb
+    args.queries = xq_ptr
+    args.queriesRowMajor = xq_row_major
+    args.queryType = xq_type
+    args.numQueries = nq
+    args.outDistances = D_ptr
+    args.outIndices = I_ptr
+    args.outIndicesType = I_type
+
+    with using_stream(res):
+        faiss.bfKnn(res, args)
+
+    return D, I
+
+torch_replace_method(faiss_module, 'knn_gpu', torch_replacement_knn_gpu, True, True)

faiss/Index.cpp

@@ -34,11 +34,10 @@ void Index::range_search (idx_t , const float *, float,
   FAISS_THROW_MSG ("range search not implemented");
 }
 
-void Index::assign (idx_t n, const float * x, idx_t * labels, idx_t k)
+void Index::assign (idx_t n, const float * x, idx_t * labels, idx_t k) const
 {
-  float * distances = new float[n * k];
-  ScopeDeleter<float> del(distances);
-  search (n, x, k, distances, labels);
+  std::vector<float> distances(n * k);
+  search (n, x, k, distances.data(), labels);
 }
 
 void Index::add_with_ids(

faiss/Index.h

@@ -133,7 +133,7 @@ struct Index {
      * @param x           input vectors to search, size n * d
      * @param labels      output labels of the NNs, size n*k
      */
-    void assign (idx_t n, const float * x, idx_t * labels, idx_t k = 1);
+    virtual void assign (idx_t n, const float * x, idx_t * labels, idx_t k = 1) const;
 
     /// removes all elements from the database.
     virtual void reset() = 0;

faiss/IndexBinary.cpp

@@ -26,10 +26,9 @@ void IndexBinary::range_search(idx_t, const uint8_t *, int,
   FAISS_THROW_MSG("range search not implemented");
 }
 
-void IndexBinary::assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k) {
-  int *distances = new int[n * k];
-  ScopeDeleter<int> del(distances);
-  search(n, x, k, distances, labels);
+void IndexBinary::assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k) const {
+  std::vector<int> distances(n * k);
+  search(n, x, k, distances.data(), labels);
 }
 
 void IndexBinary::add_with_ids(idx_t, const uint8_t *, const idx_t *) {

faiss/IndexBinary.h

@@ -120,7 +120,7 @@ struct IndexBinary {
    * @param x           input vectors to search, size n * d / 8
    * @param labels      output labels of the NNs, size n*k
    */
-  void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1);
+  void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1) const;
 
   /// Removes all elements from the database.
   virtual void reset() = 0;

faiss/gpu/GpuDistance.cu

@@ -29,7 +29,7 @@ void bfKnnConvert(GpuResourcesProvider* prov, const GpuDistanceParams& args) {
   FAISS_THROW_IF_NOT_MSG(args.vectors,
                          "bfKnn: vectors must be provided (passed null)");
   FAISS_THROW_IF_NOT_MSG(args.numQueries > 0,
-                         "bfKnn: numQUeries must be > 0");
+                         "bfKnn: numQueries must be > 0");
   FAISS_THROW_IF_NOT_MSG(args.queries,
                          "bfKnn: queries must be provided (passed null)");
   FAISS_THROW_IF_NOT_MSG(args.outDistances,

faiss/gpu/GpuIndex.cu

@@ -194,6 +194,38 @@ GpuIndex::addPage_(int n,
   }
 }
 
+void
+GpuIndex::assign(Index::idx_t n,
+                 const float* x,
+                 Index::idx_t* labels,
+                 Index::idx_t k) const {
+  FAISS_THROW_IF_NOT_MSG(this->is_trained, "Index not trained");
+
+  // For now, only support <= max int results
+  FAISS_THROW_IF_NOT_FMT(n <= (Index::idx_t) std::numeric_limits<int>::max(),
+                         "GPU index only supports up to %d indices",
+                         std::numeric_limits<int>::max());
+
+  // Maximum k-selection supported is based on the CUDA SDK
+  FAISS_THROW_IF_NOT_FMT(k <= (Index::idx_t) getMaxKSelection(),
+                         "GPU index only supports k <= %d (requested %d)",
+                         getMaxKSelection(),
+                         (int) k); // select limitation
+
+  DeviceScope scope(config_.device);
+  auto stream = resources_->getDefaultStream(config_.device);
+
+  // We need to create a throw-away buffer for distances, which we don't use but
+  // which we do need for the search call
+  DeviceTensor<float, 2, true> distances(
+    resources_.get(),
+    makeTempAlloc(AllocType::Other, stream),
+    {(int) n, (int) k});
+
+  // Forward to search
+  search(n, x, k, distances.data(), labels);
+}
+
 void
 GpuIndex::search(Index::idx_t n,
                  const float* x,
@@ -476,4 +508,9 @@ GpuIndex::compute_residual_n(Index::idx_t n,
   FAISS_THROW_MSG("compute_residual_n not implemented for this type of index");
 }
 
+std::shared_ptr<GpuResources>
+GpuIndex::getResources() {
+  return resources_;
+}
+
 } } // namespace

faiss/gpu/GpuIndex.h

@@ -39,6 +39,10 @@ class GpuIndex : public faiss::Index {
   /// Returns the device that this index is resident on
   int getDevice() const;
 
+  /// Returns a reference to our GpuResources object that manages memory, stream
+  /// and handle resources on the GPU
+  std::shared_ptr<GpuResources> getResources();
+
   /// Set the minimum data size for searches (in MiB) for which we use
   /// CPU -> GPU paging
   void setMinPagingSize(size_t size);
@@ -58,6 +62,13 @@ class GpuIndex : public faiss::Index {
                     const float* x,
                     const Index::idx_t* ids) override;
 
+  /// `x` and `labels` can be resident on the CPU or any GPU; copies are
+  /// performed as needed
+  void assign(Index::idx_t n,
+              const float* x,
+              Index::idx_t* labels,
+              Index::idx_t k = 1) const override;
+
   /// `x`, `distances` and `labels` can be resident on the CPU or any
   /// GPU; copies are performed as needed
   void search(Index::idx_t n,

faiss/gpu/GpuIndexBinaryFlat.cu

@@ -64,6 +64,11 @@ GpuIndexBinaryFlat::getDevice() const {
   return binaryFlatConfig_.device;
 }
 
+std::shared_ptr<GpuResources>
+GpuIndexBinaryFlat::getResources() {
+  return resources_;
+}
+
 void
 GpuIndexBinaryFlat::copyFrom(const faiss::IndexBinaryFlat* index) {
   DeviceScope scope(binaryFlatConfig_.device);

faiss/gpu/GpuIndexBinaryFlat.h

@@ -41,6 +41,10 @@ class GpuIndexBinaryFlat : public IndexBinary {
   /// Returns the device that this index is resident on
   int getDevice() const;
 
+  /// Returns a reference to our GpuResources object that manages memory, stream
+  /// and handle resources on the GPU
+  std::shared_ptr<GpuResources> getResources();
+
   /// Initialize ourselves from the given CPU index; will overwrite
   /// all data in ourselves
   void copyFrom(const faiss::IndexBinaryFlat* index);

faiss/gpu/GpuIndexIVFFlat.cu

@@ -140,6 +140,11 @@ GpuIndexIVFFlat::reset() {
 
 void
 GpuIndexIVFFlat::train(Index::idx_t n, const float* x) {
+  // For now, only support <= max int results
+  FAISS_THROW_IF_NOT_FMT(n <= (Index::idx_t) std::numeric_limits<int>::max(),
+                         "GPU index only supports up to %d indices",
+                         std::numeric_limits<int>::max());
+
   DeviceScope scope(config_.device);
 
   if (this->is_trained) {
@@ -151,7 +156,14 @@ GpuIndexIVFFlat::train(Index::idx_t n, const float* x) {
 
   FAISS_ASSERT(!index_);
 
-  trainQuantizer_(n, x);
+  // FIXME: GPUize more of this
+  // First, make sure that the data is resident on the CPU, if it is not on the
+  // CPU, as we depend upon parts of the CPU code
+  auto hostData = toHost<float, 2>((float*) x,
+                                   resources_->getDefaultStream(config_.device),
+                                   {(int) n, (int) this->d});
+
+  trainQuantizer_(n, hostData.data());
 
   // The quantizer is now trained; construct the IVF index
   index_.reset(new IVFFlat(resources_.get(),

faiss/gpu/GpuIndexIVFPQ.cu

@@ -276,6 +276,11 @@ GpuIndexIVFPQ::trainResidualQuantizer_(Index::idx_t n, const float* x) {
 
 void
 GpuIndexIVFPQ::train(Index::idx_t n, const float* x) {
+  // For now, only support <= max int results
+  FAISS_THROW_IF_NOT_FMT(n <= (Index::idx_t) std::numeric_limits<int>::max(),
+                         "GPU index only supports up to %d indices",
+                         std::numeric_limits<int>::max());
+
   DeviceScope scope(config_.device);
 
   if (this->is_trained) {

faiss/gpu/GpuIndexIVFScalarQuantizer.cu

@@ -187,6 +187,11 @@ GpuIndexIVFScalarQuantizer::trainResiduals_(Index::idx_t n, const float* x) {
 
 void
 GpuIndexIVFScalarQuantizer::train(Index::idx_t n, const float* x) {
+  // For now, only support <= max int results
+  FAISS_THROW_IF_NOT_FMT(n <= (Index::idx_t) std::numeric_limits<int>::max(),
+                         "GPU index only supports up to %d indices",
+                         std::numeric_limits<int>::max());
+
   DeviceScope scope(config_.device);
 
   if (this->is_trained) {

faiss/gpu/GpuResources.h

@@ -198,6 +198,10 @@ class GpuResources {
   /// given device
   virtual cudaStream_t getDefaultStream(int device) = 0;
 
+  /// Overrides the default stream for a device to the user-supplied stream. The
+  /// resources object does not own this stream (i.e., it will not destroy it).
+  virtual void setDefaultStream(int device, cudaStream_t stream) = 0;
+
   /// Returns the set of alternative streams that we use for the given device
   virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0;
 

faiss/gpu/StandardGpuResources.h

@@ -46,7 +46,7 @@ class StandardGpuResourcesImpl : public GpuResources {
   /// up.
   /// We are guaranteed that all Faiss GPU work is ordered with respect to
   /// this stream upon exit from an index or other Faiss GPU call.
-  void setDefaultStream(int device, cudaStream_t stream);
+  void setDefaultStream(int device, cudaStream_t stream) override;
 
   /// Revert the default stream to the original stream managed by this resources
   /// object, in case someone called `setDefaultStream`.

faiss/gpu/test/test_contrib.py

@@ -8,7 +8,7 @@ import unittest
 import numpy as np
 
 from faiss.contrib import datasets
-from faiss.contrib.exhaustive_search import knn_ground_truth, knn_gpu
+from faiss.contrib.exhaustive_search import knn_ground_truth
 
 from common import get_dataset_2
 
@@ -33,58 +33,3 @@ class TestComputeGT(unittest.TestCase):
 
         np.testing.assert_array_equal(Iref, Inew)
         np.testing.assert_almost_equal(Dref, Dnew, decimal=4)
-
-class TestBfKnnNumpy(unittest.TestCase):
-
-    def test_bf_knn(self):
-        d = 64
-        k = 10
-        xt, xb, xq = get_dataset_2(d, 0, 10000, 100)
-
-        index = faiss.IndexFlatL2(d)
-        index.add(xb)
-        Dref, Iref = index.search(xq, k)
-
-        res = faiss.StandardGpuResources()
-
-        D, I = knn_gpu(res, xb, xq, k)
-
-        np.testing.assert_array_equal(Iref, I)
-        np.testing.assert_almost_equal(Dref, D, decimal=4)
-
-        # Test transpositions
-        xbt = np.ascontiguousarray(xb.T)
-
-        D, I = knn_gpu(res, xbt.T, xq, k)
-
-        np.testing.assert_array_equal(Iref, I)
-        np.testing.assert_almost_equal(Dref, D, decimal=4)
-
-        xqt = np.ascontiguousarray(xq.T)
-
-        D, I = knn_gpu(res, xb, xqt.T, k)
-
-        np.testing.assert_array_equal(Iref, I)
-        np.testing.assert_almost_equal(Dref, D, decimal=4)
-
-        D, I = knn_gpu(res, xbt.T, xqt.T, k)
-
-        np.testing.assert_array_equal(Iref, I)
-        np.testing.assert_almost_equal(Dref, D, decimal=4)
-
-        # Test f16 data types
-        xb16 = xb.astype(np.float16)
-        xq16 = xq.astype(np.float16)
-
-        D, I = knn_gpu(res, xb, xq, k)
-
-        np.testing.assert_array_equal(Iref, I)
-        np.testing.assert_almost_equal(Dref, D, decimal=4)
-
-        # Test i32 indices
-        I32 = np.empty((xq.shape[0], k), dtype=np.int32)
-
-        D, _ = knn_gpu(res, xb, xq, k, I=I32)
-
-        np.testing.assert_array_equal(Iref, I32)
-        np.testing.assert_almost_equal(Dref, D, decimal=4)

faiss/gpu/test/test_pytorch_faiss.py

@@ -1,123 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-#
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-from __future__ import absolute_import, division, print_function, unicode_literals
-
-import numpy as np
-import unittest
-import faiss
-import torch
-
-from faiss.contrib.pytorch_tensors import search_index_pytorch, search_raw_array_pytorch, using_stream
-
-def to_column_major(x):
-    if hasattr(torch, 'contiguous_format'):
-        return x.t().clone(memory_format=torch.contiguous_format).t()
-    else:
-        # was default setting before memory_format was introduced
-        return x.t().clone().t()
-
-class PytorchFaissInterop(unittest.TestCase):
-
-    def test_interop(self):
-        d = 128
-        nq = 100
-        nb = 1000
-        k = 10
-
-        xq = faiss.randn(nq * d, 1234).reshape(nq, d)
-        xb = faiss.randn(nb * d, 1235).reshape(nb, d)
-
-        res = faiss.StandardGpuResources()
-
-        # Let's run on a non-default stream
-        s = torch.cuda.Stream()
-
-        # Torch will run on this stream
-        with torch.cuda.stream(s):
-            # query is pytorch tensor (CPU and GPU)
-            xq_torch_cpu = torch.FloatTensor(xq)
-            xq_torch_gpu = xq_torch_cpu.cuda()
-
-            index = faiss.GpuIndexFlatIP(res, d)
-            index.add(xb)
-
-            # Query with GPU tensor (this will be done on the current pytorch stream)
-            D2, I2 = search_index_pytorch(res, index, xq_torch_gpu, k)
-            Dref, Iref = index.search(xq, k)
-
-            assert np.all(Iref == I2.cpu().numpy())
-
-            # Query with CPU tensor
-            D3, I3 = search_index_pytorch(res, index, xq_torch_cpu, k)
-
-            assert np.all(Iref == I3.numpy())
-
-    def test_raw_array_search(self):
-        d = 32
-        nb = 1024
-        nq = 128
-        k = 10
-
-        # make GT on Faiss CPU
-
-        xq = faiss.randn(nq * d, 1234).reshape(nq, d)
-        xb = faiss.randn(nb * d, 1235).reshape(nb, d)
-
-        index = faiss.IndexFlatL2(d)
-        index.add(xb)
-        gt_D, gt_I = index.search(xq, k)
-
-        # resource object, can be re-used over calls
-        res = faiss.StandardGpuResources()
-
-        # Let's have pytorch use a non-default stream
-        s = torch.cuda.Stream()
-        with torch.cuda.stream(s):
-            for xq_row_major in True, False:
-                for xb_row_major in True, False:
-
-                    # move to pytorch & GPU
-                    xq_t = torch.from_numpy(xq).cuda()
-                    xb_t = torch.from_numpy(xb).cuda()
-
-                    if not xq_row_major:
-                        xq_t = to_column_major(xq_t)
-                        assert not xq_t.is_contiguous()
-
-                    if not xb_row_major:
-                        xb_t = to_column_major(xb_t)
-                        assert not xb_t.is_contiguous()
-
-                    D, I = search_raw_array_pytorch(res, xb_t, xq_t, k)
-
-                    # back to CPU for verification
-                    D = D.cpu().numpy()
-                    I = I.cpu().numpy()
-
-                    assert np.all(I == gt_I)
-                    assert np.all(np.abs(D - gt_D).max() < 1e-4)
-
-                    # test on subset
-                    try:
-                        # This internally uses the current pytorch stream
-                        D, I = search_raw_array_pytorch(res, xb_t, xq_t[60:80], k)
-                    except TypeError:
-                        if not xq_row_major:
-                            # then it is expected
-                            continue
-                        # otherwise it is an error
-                        raise
-
-                    # back to CPU for verification
-                    D = D.cpu().numpy()
-                    I = I.cpu().numpy()
-
-                    assert np.all(I == gt_I[60:80])
-                    assert np.all(np.abs(D - gt_D[60:80]).max() < 1e-4)
-
-
-if __name__ == '__main__':
-    unittest.main()

faiss/gpu/test/torch_test_contrib_gpu.py

@@ -0,0 +1,338 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import faiss
+import torch
+import unittest
+import numpy as np
+import faiss.contrib.torch_utils
+
+def to_column_major_torch(x):
+    if hasattr(torch, 'contiguous_format'):
+        return x.t().clone(memory_format=torch.contiguous_format).t()
+    else:
+        # was default setting before memory_format was introduced
+        return x.t().clone().t()
+
+def to_column_major_numpy(x):
+    return x.T.copy().T
+
+class TestTorchUtilsGPU(unittest.TestCase):
+    # tests add, search
+    def test_lookup(self):
+        cpu_index = faiss.IndexFlatL2(128)
+
+        # Add to CPU index with np
+        xb_torch = torch.rand(10000, 128)
+        cpu_index.add(xb_torch.numpy())
+
+        # Add to CPU index with torch GPU (should fail)
+        xb_torch_gpu = torch.rand(10000, 128, device=torch.device('cuda', 0), dtype=torch.float32)
+        with self.assertRaises(AssertionError):
+            cpu_index.add(xb_torch_gpu)
+
+        # Add to GPU with torch GPU
+        res = faiss.StandardGpuResources()
+        gpu_index = faiss.GpuIndexFlatL2(res, 128)
+        gpu_index.add(xb_torch.cuda())
+
+        # Search with torch CPU
+        xq_torch_cpu = torch.rand(10, 128, dtype=torch.float32)
+        d_torch_cpu, i_torch_cpu = gpu_index.search(xq_torch_cpu, 10)
+
+        # Search with torch GPU
+        xq_torch_gpu = xq_torch_cpu.cuda()
+        d_torch_gpu, i_torch_gpu = gpu_index.search(xq_torch_gpu, 10)
+        self.assertTrue(d_torch_gpu.is_cuda)
+        self.assertTrue(i_torch_gpu.is_cuda)
+
+        # Should be equivalent
+        self.assertTrue(torch.equal(d_torch_cpu.cuda(), d_torch_gpu))
+        self.assertTrue(torch.equal(i_torch_cpu.cuda(), i_torch_gpu))
+
+        # Search with torch GPU using pre-allocated arrays
+        new_d_torch_gpu = torch.zeros(10, 10, device=torch.device('cuda', 0), dtype=torch.float32)
+        new_i_torch_gpu = torch.zeros(10, 10, device=torch.device('cuda', 0), dtype=torch.int64)
+        gpu_index.search(xq_torch_gpu, 10, new_d_torch_gpu, new_i_torch_gpu)
+
+        self.assertTrue(torch.equal(d_torch_cpu.cuda(), new_d_torch_gpu))
+        self.assertTrue(torch.equal(i_torch_cpu.cuda(), new_i_torch_gpu))
+
+        # Search with numpy CPU
+        xq_np_cpu = xq_torch_cpu.numpy()
+        d_np_cpu, i_np_cpu = gpu_index.search(xq_np_cpu, 10)
+        self.assertEqual(type(d_np_cpu), np.ndarray)
+        self.assertEqual(type(i_np_cpu), np.ndarray)
+
+        self.assertTrue(np.array_equal(d_torch_cpu.numpy(), d_np_cpu))
+        self.assertTrue(np.array_equal(i_torch_cpu.numpy(), i_np_cpu))
+
+    # tests train, add_with_ids
+    def test_train_add_with_ids(self):
+        d = 32
+        nlist = 5
+        res = faiss.StandardGpuResources()
+        res.noTempMemory()
+
+        index = faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2)
+        xb = torch.rand(1000, d, device=torch.device('cuda', 0), dtype=torch.float32)
+        index.train(xb)
+
+        ids = torch.arange(1000, 1000 + xb.shape[0], device=torch.device('cuda', 0), dtype=torch.int64)
+
+        # Test add_with_ids with torch gpu
+        index.add_with_ids(xb, ids)
+        _, I = index.search(xb[10:20], 1)
+        self.assertTrue(torch.equal(I.view(10), ids[10:20]))
+
+        # Test add_with_ids with torch cpu
+        index.reset()
+        xb_cpu = xb.cpu()
+        ids_cpu = ids.cpu()
+
+        index.train(xb_cpu)
+        index.add_with_ids(xb_cpu, ids_cpu)
+        _, I = index.search(xb_cpu[10:20], 1)
+        self.assertTrue(torch.equal(I.view(10), ids_cpu[10:20]))
+
+        # Test add_with_ids with numpy
+        index.reset()
+        xb_np = xb.cpu().numpy()
+        ids_np = ids.cpu().numpy()
+
+        index.train(xb_np)
+        index.add_with_ids(xb_np, ids_np)
+        _, I = index.search(xb_np[10:20], 1)
+        self.assertTrue(np.array_equal(I.reshape(10), ids_np[10:20]))
+
+    # tests reconstruct, reconstruct_n
+    def test_reconstruct(self):
+        d = 32
+        res = faiss.StandardGpuResources()
+        res.noTempMemory()
+        index = faiss.GpuIndexFlatL2(res, d)
+
+        xb = torch.rand(100, d, device=torch.device('cuda', 0), dtype=torch.float32)
+        index.add(xb)
+
+        # Test reconstruct with torch gpu (native return)
+        y = index.reconstruct(7)
+        self.assertTrue(y.is_cuda)
+        self.assertTrue(torch.equal(xb[7], y))
+
+        # Test reconstruct with numpy output provided
+        y = np.empty(d, dtype='float32')
+        index.reconstruct(11, y)
+        self.assertTrue(np.array_equal(xb.cpu().numpy()[11], y))
+
+        # Test reconstruct with torch cpu output providesd
+        y = torch.empty(d, dtype=torch.float32)
+        index.reconstruct(12, y)
+        self.assertTrue(torch.equal(xb[12].cpu(), y))
+
+        # Test reconstruct with torch gpu output providesd
+        y = torch.empty(d, device=torch.device('cuda', 0), dtype=torch.float32)
+        index.reconstruct(13, y)
+        self.assertTrue(torch.equal(xb[13], y))
+
+        # Test reconstruct_n with torch gpu (native return)
+        y = index.reconstruct_n(10, 10)
+        self.assertTrue(y.is_cuda)
+        self.assertTrue(torch.equal(xb[10:20], y))
+
+        # Test reconstruct with numpy output provided
+        y = np.empty((10, d), dtype='float32')
+        index.reconstruct_n(20, 10, y)
+        self.assertTrue(np.array_equal(xb.cpu().numpy()[20:30], y))
+
+        # Test reconstruct_n with torch cpu output provided
+        y = torch.empty(10, d, dtype=torch.float32)
+        index.reconstruct_n(40, 10, y)
+        self.assertTrue(torch.equal(xb[40:50].cpu(), y))
+
+        # Test reconstruct_n with torch gpu output provided
+        y = torch.empty(10, d, device=torch.device('cuda', 0), dtype=torch.float32)
+        index.reconstruct_n(50, 10, y)
+        self.assertTrue(torch.equal(xb[50:60], y))
+
+    # tests assign
+    def test_assign(self):
+        d = 32
+        res = faiss.StandardGpuResources()
+        res.noTempMemory()
+
+        index = faiss.GpuIndexFlatL2(res, d)
+        xb = torch.rand(10000, d, device=torch.device('cuda', 0), dtype=torch.float32)
+        index.add(xb)
+
+        index_cpu = faiss.IndexFlatL2(d)
+        index.copyTo(index_cpu)
+
+        # Test assign with native gpu output
+        # both input as gpu torch and input as cpu torch
+        xq = torch.rand(10, d, device=torch.device('cuda', 0), dtype=torch.float32)
+
+        labels = index.assign(xq, 5)
+        labels_cpu = index_cpu.assign(xq.cpu(), 5)
+        self.assertTrue(torch.equal(labels.cpu(), labels_cpu))
+
+        # Test assign with np input
+        labels = index.assign(xq.cpu().numpy(), 5)
+        labels_cpu = index_cpu.assign(xq.cpu().numpy(), 5)
+        self.assertTrue(np.array_equal(labels, labels_cpu))
+
+        # Test assign with numpy output provided
+        labels = np.empty((xq.shape[0], 5), dtype='int64')
+        index.assign(xq.cpu().numpy(), 5, labels)
+        self.assertTrue(np.array_equal(labels, labels_cpu))
+
+        # Test assign with torch cpu output provided
+        labels = torch.empty(xq.shape[0], 5, dtype=torch.int64)
+        index.assign(xq.cpu(), 5, labels)
+        labels_cpu = index_cpu.assign(xq.cpu(), 5)
+        self.assertTrue(torch.equal(labels, labels_cpu))
+
+    # tests remove_ids
+    def test_remove_ids(self):
+        # This is not currently implemented on GPU indices
+        return
+
+    # tests range_search
+    def test_range_search(self):
+        # This is not currently implemented on GPU indices
+        return
+
+    # tests search_and_reconstruct
+    def test_search_and_reconstruct(self):
+        # This is not currently implemented on GPU indices
+        return
+
+    # tests sa_encode, sa_decode
+    def test_sa_encode_decode(self):
+        # This is not currently implemented on GPU indices
+        return
+
+class TestTorchUtilsKnnGpu(unittest.TestCase):
+    def test_knn_gpu(self):
+        torch.manual_seed(10)
+        d = 32
+        nb = 1024
+        nq = 10
+        k = 10
+        res = faiss.StandardGpuResources()
+
+        # make GT on torch cpu and test using IndexFlatL2
+        xb = torch.rand(nb, d, dtype=torch.float32)
+        xq = torch.rand(nq, d, dtype=torch.float32)
+
+        index = faiss.IndexFlatL2(d)
+        index.add(xb)
+        gt_D, gt_I = index.search(xq, k)
+
+        # for the GPU, we'll use a non-default stream
+        s = torch.cuda.Stream()
+        with torch.cuda.stream(s):
+            # test numpy inputs
+            xb_np = xb.numpy()
+            xq_np = xq.numpy()
+
+            for xq_row_major in True, False:
+                for xb_row_major in True, False:
+                    if not xq_row_major:
+                        xq_c = to_column_major_numpy(xq_np)
+                        assert not xq_c.flags.contiguous
+                    else:
+                        xq_c = xq_np
+
+                    if not xb_row_major:
+                        xb_c = to_column_major_numpy(xb_np)
+                        assert not xb_c.flags.contiguous
+                    else:
+                        xb_c = xb_np
+
+                    D, I = faiss.knn_gpu(res, xb_c, xq_c, k)
+
+                    self.assertTrue(torch.equal(torch.from_numpy(I), gt_I))
+                    self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1e-4)
+
+            # test torch (cpu, gpu) inputs
+            for is_cuda in True, False:
+                for xq_row_major in True, False:
+                    for xb_row_major in True, False:
+
+                        if is_cuda:
+                            xq_c = xq.cuda()
+                            xb_c = xb.cuda()
+                        else:
+                            # also test torch cpu tensors
+                            xq_c = xq
+                            xb_c = xb
+
+                        if not xq_row_major:
+                            xq_c = to_column_major_torch(xq)
+                            assert not xq_c.is_contiguous()
+
+                        if not xb_row_major:
+                            xb_c = to_column_major_torch(xb)
+                            assert not xb_c.is_contiguous()
+
+                        D, I = faiss.knn_gpu(res, xb_c, xq_c, k)
+
+                        self.assertTrue(torch.equal(I.cpu(), gt_I))
+                        self.assertLess((D.cpu() - gt_D).abs().max(), 1e-4)
+
+                        # test on subset
+                        try:
+                            # This internally uses the current pytorch stream
+                            D, I = faiss.knn_gpu(res, xb_c, xq_c[6:8], k)
+                        except TypeError:
+                            if not xq_row_major:
+                                # then it is expected
+                                continue
+                            # otherwise it is an error
+                            raise
+
+                        self.assertTrue(torch.equal(I.cpu(), gt_I[6:8]))
+                        self.assertLess((D.cpu() - gt_D[6:8]).abs().max(), 1e-4)
+
+    def test_knn_gpu_datatypes(self):
+        torch.manual_seed(10)
+        d = 10
+        nb = 1024
+        nq = 5
+        k = 10
+        res = faiss.StandardGpuResources()
+
+        # make GT on torch cpu and test using IndexFlatL2
+        xb = torch.rand(nb, d, dtype=torch.float32)
+        xq = torch.rand(nq, d, dtype=torch.float32)
+
+        index = faiss.IndexFlatL2(d)
+        index.add(xb)
+        gt_D, gt_I = index.search(xq, k)
+
+        xb_c = xb.cuda().half()
+        xq_c = xq.cuda().half()
+
+        # use i32 output indices
+        D = torch.zeros(nq, k, device=xb_c.device, dtype=torch.float32)
+        I = torch.zeros(nq, k, device=xb_c.device, dtype=torch.int32)
+
+        faiss.knn_gpu(res, xb_c, xq_c, k, D, I)
+
+        self.assertTrue(torch.equal(I.long().cpu(), gt_I))
+        self.assertLess((D.float().cpu() - gt_D).abs().max(), 1.5e-3)
+
+        # Test using numpy
+        D = np.zeros((nq, k), dtype=np.float32)
+        I = np.zeros((nq, k), dtype=np.int32)
+
+        xb_c = xb.half().numpy()
+        xq_c = xq.half().numpy()
+
+        faiss.knn_gpu(res, xb_c, xq_c, k, D, I)
+
+        self.assertTrue(torch.equal(torch.from_numpy(I).long(), gt_I))
+        self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1.5e-3)

faiss/python/__init__.py

@@ -40,7 +40,6 @@ def replace_method(the_class, name, replacement, ignore_missing=False):
     setattr(the_class, name + '_c', orig_method)
     setattr(the_class, name, replacement)
 
-
 def handle_Clustering():
     def replacement_train(self, x, index, weights=None):
         n, d = x.shape
@@ -99,7 +98,6 @@ handle_Quantizer(ScalarQuantizer)
 def handle_Index(the_class):
 
     def replacement_add(self, x):
-        assert x.flags.contiguous
         n, d = x.shape
         assert d == self.d
         self.add_c(n, swig_ptr(x))
@@ -107,43 +105,68 @@ def handle_Index(the_class):
     def replacement_add_with_ids(self, x, ids):
         n, d = x.shape
         assert d == self.d
+
         assert ids.shape == (n, ), 'not same nb of vectors as ids'
         self.add_with_ids_c(n, swig_ptr(x), swig_ptr(ids))
 
-    def replacement_assign(self, x, k):
+    def replacement_assign(self, x, k, labels=None):
         n, d = x.shape
         assert d == self.d
-        labels = np.empty((n, k), dtype=np.int64)
+
+        if labels is None:
+            labels = np.empty((n, k), dtype=np.int64)
+        else:
+            assert labels.shape == (n, k)
+
         self.assign_c(n, swig_ptr(x), swig_ptr(labels), k)
         return labels
 
     def replacement_train(self, x):
-        assert x.flags.contiguous
         n, d = x.shape
         assert d == self.d
         self.train_c(n, swig_ptr(x))
 
-    def replacement_search(self, x, k):
+    def replacement_search(self, x, k, D=None, I=None):
         n, d = x.shape
         assert d == self.d
-        distances = np.empty((n, k), dtype=np.float32)
-        labels = np.empty((n, k), dtype=np.int64)
-        self.search_c(n, swig_ptr(x),
-                      k, swig_ptr(distances),
-                      swig_ptr(labels))
-        return distances, labels
 
-    def replacement_search_and_reconstruct(self, x, k):
+        if D is None:
+            D = np.empty((n, k), dtype=np.float32)
+        else:
+            assert D.shape == (n, k)
+
+        if I is None:
+            I = np.empty((n, k), dtype=np.int64)
+        else:
+            assert I.shape == (n, k)
+
+        self.search_c(n, swig_ptr(x), k, swig_ptr(D), swig_ptr(I))
+        return D, I
+
+    def replacement_search_and_reconstruct(self, x, k, D=None, I=None, R=None):
         n, d = x.shape
         assert d == self.d
-        distances = np.empty((n, k), dtype=np.float32)
-        labels = np.empty((n, k), dtype=np.int64)
-        recons = np.empty((n, k, d), dtype=np.float32)
+
+        if D is None:
+            D = np.empty((n, k), dtype=np.float32)
+        else:
+            assert D.shape == (n, k)
+
+        if I is None:
+            I = np.empty((n, k), dtype=np.int64)
+        else:
+            assert I.shape == (n, k)
+
+        if R is None:
+            R = np.empty((n, k, d), dtype=np.float32)
+        else:
+            assert R.shape == (n, k, d)
+
         self.search_and_reconstruct_c(n, swig_ptr(x),
-                                      k, swig_ptr(distances),
-                                      swig_ptr(labels),
-                                      swig_ptr(recons))
-        return distances, labels, recons
+                                      k, swig_ptr(D),
+                                      swig_ptr(I),
+                                      swig_ptr(R))
+        return D, I, R
 
     def replacement_remove_ids(self, x):
         if isinstance(x, IDSelector):
@@ -157,13 +180,21 @@ def handle_Index(the_class):
                 sel = IDSelectorBatch(x.size, swig_ptr(x))
         return self.remove_ids_c(sel)
 
-    def replacement_reconstruct(self, key):
-        x = np.empty(self.d, dtype=np.float32)
+    def replacement_reconstruct(self, key, x=None):
+        if x is None:
+            x = np.empty(self.d, dtype=np.float32)
+        else:
+            assert x.shape == (self.d, )
+
         self.reconstruct_c(key, swig_ptr(x))
         return x
 
-    def replacement_reconstruct_n(self, n0, ni):
-        x = np.empty((ni, self.d), dtype=np.float32)
+    def replacement_reconstruct_n(self, n0, ni, x=None):
+        if x is None:
+            x = np.empty((ni, self.d), dtype=np.float32)
+        else:
+            assert x.shape == (ni, self.d)
+
         self.reconstruct_n_c(n0, ni, swig_ptr(x))
         return x
 
@@ -171,11 +202,14 @@ def handle_Index(the_class):
         n = keys.size
         assert keys.shape == (n, )
         assert x.shape == (n, self.d)
+
         self.update_vectors_c(n, swig_ptr(keys), swig_ptr(x))
 
+    # The CPU does not support passed-in output buffers
     def replacement_range_search(self, x, thresh):
         n, d = x.shape
         assert d == self.d
+
         res = RangeSearchResult(n)
         self.range_search_c(n, swig_ptr(x), thresh, res)
         # get pointers and copy them
@@ -185,17 +219,27 @@ def handle_Index(the_class):
         I = rev_swig_ptr(res.labels, nd).copy()
         return lims, D, I
 
-    def replacement_sa_encode(self, x):
+    def replacement_sa_encode(self, x, codes=None):
         n, d = x.shape
         assert d == self.d
-        codes = np.empty((n, self.sa_code_size()), dtype='uint8')
+
+        if codes is None:
+            codes = np.empty((n, self.sa_code_size()), dtype=np.uint8)
+        else:
+            assert codes.shape == (n, self.sa_code_size())
+
         self.sa_encode_c(n, swig_ptr(x), swig_ptr(codes))
         return codes
 
-    def replacement_sa_decode(self, codes):
+    def replacement_sa_decode(self, codes, x=None):
         n, cs = codes.shape
         assert cs == self.sa_code_size()
-        x = np.empty((n, self.d), dtype='float32')
+
+        if x is None:
+            x = np.empty((n, self.d), dtype=np.float32)
+        else:
+            assert x.shape == (n, self.d)
+
         self.sa_decode_c(n, swig_ptr(codes), swig_ptr(x))
         return x
 
@@ -218,7 +262,6 @@ def handle_Index(the_class):
 def handle_IndexBinary(the_class):
 
     def replacement_add(self, x):
-        assert x.flags.contiguous
         n, d = x.shape
         assert d * 8 == self.d
         self.add_c(n, swig_ptr(x))
@@ -230,7 +273,6 @@ def handle_IndexBinary(the_class):
         self.add_with_ids_c(n, swig_ptr(x), swig_ptr(ids))
 
     def replacement_train(self, x):
-        assert x.flags.contiguous
         n, d = x.shape
         assert d * 8 == self.d
         self.train_c(n, swig_ptr(x))
@@ -282,7 +324,6 @@ def handle_IndexBinary(the_class):
 def handle_VectorTransform(the_class):
 
     def apply_method(self, x):
-        assert x.flags.contiguous
         n, d = x.shape
         assert d == self.d_in
         y = np.empty((n, self.d_out), dtype=np.float32)
@@ -297,7 +338,6 @@ def handle_VectorTransform(the_class):
         return y
 
     def replacement_vt_train(self, x):
-        assert x.flags.contiguous
         n, d = x.shape
         assert d == self.d_in
         self.train_c(n, swig_ptr(x))
@@ -503,6 +543,89 @@ def index_cpu_to_gpus_list(index, co=None, gpus=None, ngpu=-1):
     index_gpu = index_cpu_to_gpu_multiple_py(res, index, co, gpus)
     return index_gpu
 
+# allows numpy ndarray usage with bfKnn
+def knn_gpu(res, xb, xq, k, D=None, I=None, metric=METRIC_L2):
+    nq, d = xq.shape
+    if xq.flags.c_contiguous:
+        xq_row_major = True
+    elif xq.flags.f_contiguous:
+        xq = xq.T
+        xq_row_major = False
+    else:
+        raise TypeError('matrix should be row (C) or column-major (Fortran)')
+
+    xq_ptr = swig_ptr(xq)
+
+    if xq.dtype == np.float32:
+        xq_type = DistanceDataType_F32
+    elif xq.dtype == np.float16:
+        xq_type = DistanceDataType_F16
+    else:
+        raise TypeError('xq must be f32 or f16')
+
+    nb, d2 = xb.shape
+    assert d2 == d
+    if xb.flags.c_contiguous:
+        xb_row_major = True
+    elif xb.flags.f_contiguous:
+        xb = xb.T
+        xb_row_major = False
+    else:
+        raise TypeError('matrix should be row (C) or column-major (Fortran)')
+
+    xb_ptr = swig_ptr(xb)
+
+    if xb.dtype == np.float32:
+        xb_type = DistanceDataType_F32
+    elif xb.dtype == np.float16:
+        xb_type = DistanceDataType_F16
+    else:
+        raise TypeError('xb must be f32 or f16')
+
+    if D is None:
+        D = np.empty((nq, k), dtype=np.float32)
+    else:
+        assert D.shape == (nq, k)
+        # interface takes void*, we need to check this
+        assert D.dtype == np.float32
+
+    D_ptr = swig_ptr(D)
+
+    if I is None:
+        I = np.empty((nq, k), dtype=np.int64)
+    else:
+        assert I.shape == (nq, k)
+
+    I_ptr = swig_ptr(I)
+
+    if I.dtype == np.int64:
+        I_type = IndicesDataType_I64
+    elif I.dtype == I.dtype == np.int32:
+        I_type = IndicesDataType_I32
+    else:
+        raise TypeError('I must be i64 or i32')
+
+    args = GpuDistanceParams()
+    args.metric = metric
+    args.k = k
+    args.dims = d
+    args.vectors = xb_ptr
+    args.vectorsRowMajor = xb_row_major
+    args.vectorType = xb_type
+    args.numVectors = nb
+    args.queries = xq_ptr
+    args.queriesRowMajor = xq_row_major
+    args.queryType = xq_type
+    args.numQueries = nq
+    args.outDistances = D_ptr
+    args.outIndices = I_ptr
+    args.outIndicesType = I_type
+
+    # no stream synchronization needed, inputs and outputs are guaranteed to
+    # be on the CPU (numpy arrays)
+    bfKnn(res, args)
+
+    return D, I
 
 ###########################################
 # numpy array / std::vector conversions

faiss/python/swigfaiss.swig

@@ -184,6 +184,7 @@ namespace std {
 %include <std_string.i>
 %include <std_pair.i>
 %include <std_map.i>
+%include <std_shared_ptr.i>
 
 // primitive array types
 %template(FloatVector) std::vector<float>;
@@ -238,6 +239,9 @@ void gpu_sync_all_devices();
 
 #ifdef GPU_WRAPPER
 
+%shared_ptr(faiss::gpu::GpuResources);
+%shared_ptr(faiss::gpu::StandardGpuResourcesImpl);
+
 %{
 
 #include <faiss/gpu/StandardGpuResources.h>
@@ -293,11 +297,11 @@ typedef CUstream_st* cudaStream_t;
 %inline %{
 
 // interop between pytorch exposed cudaStream_t and faiss
-cudaStream_t cast_integer_to_cudastream_t(long x) {
+cudaStream_t cast_integer_to_cudastream_t(long long x) {
   return (cudaStream_t) x;
 }
 
-long cast_cudastream_t_to_integer(cudaStream_t x) {
+long long cast_cudastream_t_to_integer(cudaStream_t x) {
   return (long) x;
 }
 
@@ -799,18 +803,26 @@ void *memcpy(void *dest, const void *src, size_t n);
  *******************************************************************/
 
 %inline %{
-float * cast_integer_to_float_ptr (long x) {
+uint8_t * cast_integer_to_uint8_ptr (long long x) {
+    return (uint8_t*)x;
+}
+
+float * cast_integer_to_float_ptr (long long x) {
     return (float*)x;
 }
 
-long * cast_integer_to_long_ptr (long x) {
-    return (long*)x;
+faiss::Index::idx_t* cast_integer_to_idx_t_ptr (long long x) {
+    return (faiss::Index::idx_t*)x;
 }
 
-int * cast_integer_to_int_ptr (long x) {
+int * cast_integer_to_int_ptr (long long x) {
     return (int*)x;
 }
 
+void * cast_integer_to_void_ptr (long long x) {
+    return (void*)x;
+}
+
 %}
 
 

tests/torch_test_contrib.py

@@ -0,0 +1,344 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import faiss
+import torch
+import unittest
+import numpy as np
+import faiss.contrib.torch_utils
+
+class TestTorchUtilsCPU(unittest.TestCase):
+    # tests add, search
+    def test_lookup(self):
+        d = 128
+        index = faiss.IndexFlatL2(d)
+
+        # Add to CPU index with torch CPU
+        xb_torch = torch.rand(10000, d)
+        index.add(xb_torch)
+
+        # Test reconstruct
+        y_torch = index.reconstruct(10)
+        self.assertTrue(torch.equal(y_torch, xb_torch[10]))
+
+        # Add to CPU index with numpy CPU
+        xb_np = torch.rand(500, d).numpy()
+        index.add(xb_np)
+        self.assertEqual(index.ntotal, 10500)
+
+        y_np = np.zeros(d, dtype=np.float32)
+        index.reconstruct(10100, y_np)
+        self.assertTrue(np.array_equal(y_np, xb_np[100]))
+
+        # Search with np cpu
+        xq_torch = torch.rand(10, d, dtype=torch.float32)
+        d_np, I_np = index.search(xq_torch.numpy(), 5)
+
+        # Search with torch cpu
+        d_torch, I_torch = index.search(xq_torch, 5)
+
+        # The two should be equivalent
+        self.assertTrue(np.array_equal(d_np, d_torch.numpy()))
+        self.assertTrue(np.array_equal(I_np, I_torch.numpy()))
+
+        # Search with np cpu using pre-allocated arrays
+        d_np_input = np.zeros((10, 5), dtype=np.float32)
+        I_np_input = np.zeros((10, 5), dtype=np.int64)
+        index.search(xq_torch.numpy(), 5, d_np_input, I_np_input)
+
+        self.assertTrue(np.array_equal(d_np, d_np_input))
+        self.assertTrue(np.array_equal(I_np, I_np_input))
+
+        # Search with torch cpu using pre-allocated arrays
+        d_torch_input = torch.zeros(10, 5, dtype=torch.float32)
+        I_torch_input = torch.zeros(10, 5, dtype=torch.int64)
+        index.search(xq_torch, 5, d_torch_input, I_torch_input)
+
+        self.assertTrue(np.array_equal(d_torch_input.numpy(), d_np))
+        self.assertTrue(np.array_equal(I_torch_input.numpy(), I_np))
+
+    # tests train, add_with_ids
+    def test_train_add_with_ids(self):
+        d = 32
+        nlist = 5
+
+        quantizer = faiss.IndexFlatL2(d)
+        index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
+        xb = torch.rand(1000, d, dtype=torch.float32)
+        index.train(xb)
+
+        # Test add_with_ids with torch cpu
+        ids = torch.arange(1000, 1000 + xb.shape[0], dtype=torch.int64)
+        index.add_with_ids(xb, ids)
+        _, I = index.search(xb[10:20], 1)
+        self.assertTrue(torch.equal(I.view(10), ids[10:20]))
+
+        # Test add_with_ids with numpy
+        index.reset()
+        index.train(xb.numpy())
+        index.add_with_ids(xb.numpy(), ids.numpy())
+        _, I = index.search(xb.numpy()[10:20], 1)
+        self.assertTrue(np.array_equal(I.reshape(10), ids.numpy()[10:20]))
+
+    # tests reconstruct, reconstruct_n
+    def test_reconstruct(self):
+        d = 32
+        index = faiss.IndexFlatL2(d)
+
+        xb = torch.rand(100, d, dtype=torch.float32)
+        index.add(xb)
+
+        # Test reconstruct with torch cpu (native return)
+        y = index.reconstruct(7)
+        self.assertTrue(torch.equal(xb[7], y))
+
+        # Test reconstruct with numpy output provided
+        y = np.empty(d, dtype=np.float32)
+        index.reconstruct(11, y)
+        self.assertTrue(np.array_equal(xb.numpy()[11], y))
+
+        # Test reconstruct with torch cpu output providesd
+        y = torch.empty(d, dtype=torch.float32)
+        index.reconstruct(12, y)
+        self.assertTrue(torch.equal(xb[12], y))
+
+        # Test reconstruct_n with torch cpu (native return)
+        y = index.reconstruct_n(10, 10)
+        self.assertTrue(torch.equal(xb[10:20], y))
+
+        # Test reconstruct with numpy output provided
+        y = np.empty((10, d), dtype=np.float32)
+        index.reconstruct_n(20, 10, y)
+        self.assertTrue(np.array_equal(xb.cpu().numpy()[20:30], y))
+
+        # Test reconstruct_n with torch cpu output provided
+        y = torch.empty(10, d, dtype=torch.float32)
+        index.reconstruct_n(40, 10, y)
+        self.assertTrue(torch.equal(xb[40:50].cpu(), y))
+
+    # tests assign
+    def test_assign(self):
+        d = 32
+        index = faiss.IndexFlatL2(d)
+        xb = torch.rand(1000, d, dtype=torch.float32)
+        index.add(xb)
+
+        index_ref = faiss.IndexFlatL2(d)
+        index_ref.add(xb.numpy())
+
+        # Test assign with native cpu output
+        xq = torch.rand(10, d, dtype=torch.float32)
+        labels = index.assign(xq, 5)
+        labels_ref = index_ref.assign(xq.cpu(), 5)
+
+        self.assertTrue(torch.equal(labels, labels_ref))
+
+        # Test assign with np input
+        labels = index.assign(xq.numpy(), 5)
+        labels_ref = index_ref.assign(xq.numpy(), 5)
+        self.assertTrue(np.array_equal(labels, labels_ref))
+
+        # Test assign with numpy output provided
+        labels = np.empty((xq.shape[0], 5), dtype='int64')
+        index.assign(xq.numpy(), 5, labels)
+        self.assertTrue(np.array_equal(labels, labels_ref))
+
+        # Test assign with torch cpu output provided
+        labels = torch.empty(xq.shape[0], 5, dtype=torch.int64)
+        index.assign(xq, 5, labels)
+        labels_ref = index_ref.assign(xq, 5)
+        self.assertTrue(torch.equal(labels, labels_ref))
+
+    # tests remove_ids
+    def test_remove_ids(self):
+        # only implemented for cpu index + numpy at the moment
+        d = 32
+        quantizer = faiss.IndexFlatL2(d)
+        index = faiss.IndexIVFFlat(quantizer, d, 5)
+        index.make_direct_map()
+        index.set_direct_map_type(faiss.DirectMap.Hashtable)
+
+        xb = torch.rand(1000, d, dtype=torch.float32)
+        ids = torch.arange(1000, 1000 + xb.shape[0], dtype=torch.int64)
+        index.train(xb)
+        index.add_with_ids(xb, ids)
+
+        ids_remove = np.array([1010], dtype=np.int64)
+        index.remove_ids(ids_remove)
+
+        # We should find this
+        y = index.reconstruct(1011)
+        self.assertTrue(np.array_equal(xb[11].numpy(), y))
+
+        # We should not find this
+        with self.assertRaises(RuntimeError):
+            y = index.reconstruct(1010)
+
+        # Torch not yet supported
+        ids_remove = torch.tensor([1012], dtype=torch.int64)
+        with self.assertRaises(AssertionError):
+            index.remove_ids(ids_remove)
+
+    # tests update_vectors
+    def test_update_vectors(self):
+        d = 32
+        quantizer_np = faiss.IndexFlatL2(d)
+        index_np = faiss.IndexIVFFlat(quantizer_np, d, 5)
+        index_np.make_direct_map()
+        index_np.set_direct_map_type(faiss.DirectMap.Hashtable)
+
+        quantizer_torch = faiss.IndexFlatL2(d)
+        index_torch = faiss.IndexIVFFlat(quantizer_torch, d, 5)
+        index_torch.make_direct_map()
+        index_torch.set_direct_map_type(faiss.DirectMap.Hashtable)
+
+        xb = torch.rand(1000, d, dtype=torch.float32)
+        ids = torch.arange(1000, 1000 + xb.shape[0], dtype=torch.int64)
+
+        index_np.train(xb.numpy())
+        index_np.add_with_ids(xb.numpy(), ids.numpy())
+
+        index_torch.train(xb)
+        index_torch.add_with_ids(xb, ids)
+
+        xb_up = torch.rand(10, d, dtype=torch.float32)
+        ids_up = ids[0:10]
+
+        index_np.update_vectors(ids_up.numpy(), xb_up.numpy())
+        index_torch.update_vectors(ids_up, xb_up)
+
+        xq = torch.rand(10, d, dtype=torch.float32)
+
+        D_np, I_np = index_np.search(xq.numpy(), 5)
+        D_torch, I_torch = index_torch.search(xq, 5)
+
+        self.assertTrue(np.array_equal(D_np, D_torch.numpy()))
+        self.assertTrue(np.array_equal(I_np, I_torch.numpy()))
+
+    # tests range_search
+    def test_range_search(self):
+        torch.manual_seed(10)
+        d = 32
+        index = faiss.IndexFlatL2(d)
+        xb = torch.rand(100, d, dtype=torch.float32)
+        index.add(xb)
+
+        # torch cpu as ground truth
+        thresh = 2.9
+        xq = torch.rand(10, d, dtype=torch.float32)
+        lims, D, I = index.range_search(xq, thresh)
+
+        # compare against np
+        lims_np, D_np, I_np = index.range_search(xq.numpy(), thresh)
+
+        self.assertTrue(np.array_equal(lims.numpy(), lims_np))
+        self.assertTrue(np.array_equal(D.numpy(), D_np))
+        self.assertTrue(np.array_equal(I.numpy(), I_np))
+
+    # tests search_and_reconstruct
+    def test_search_and_reconstruct(self):
+        d = 32
+        nlist = 10
+        M = 4
+        k = 5
+        quantizer = faiss.IndexFlatL2(d)
+        index = faiss.IndexIVFPQ(quantizer, d, nlist, M, 4)
+
+        xb = torch.rand(1000, d, dtype=torch.float32)
+        index.train(xb)
+
+        # different set
+        xb = torch.rand(500, d, dtype=torch.float32)
+        index.add(xb)
+
+        # torch cpu as ground truth
+        xq = torch.rand(10, d, dtype=torch.float32)
+        D, I, R = index.search_and_reconstruct(xq, k)
+
+        # compare against numpy
+        D_np, I_np, R_np = index.search_and_reconstruct(xq.numpy(), k)
+
+        self.assertTrue(np.array_equal(D.numpy(), D_np))
+        self.assertTrue(np.array_equal(I.numpy(), I_np))
+        self.assertTrue(np.array_equal(R.numpy(), R_np))
+
+        # numpy input values
+        D_input = np.zeros((xq.shape[0], k), dtype=np.float32)
+        I_input = np.zeros((xq.shape[0], k), dtype=np.int64)
+        R_input = np.zeros((xq.shape[0], k, d), dtype=np.float32)
+
+        index.search_and_reconstruct(xq.numpy(), k, D_input, I_input, R_input)
+
+        self.assertTrue(np.array_equal(D.numpy(), D_input))
+        self.assertTrue(np.array_equal(I.numpy(), I_input))
+        self.assertTrue(np.array_equal(R.numpy(), R_input))
+
+        # torch input values
+        D_input = torch.zeros(xq.shape[0], k, dtype=torch.float32)
+        I_input = torch.zeros(xq.shape[0], k, dtype=torch.int64)
+        R_input = torch.zeros(xq.shape[0], k, d, dtype=torch.float32)
+
+        index.search_and_reconstruct(xq, k, D_input, I_input, R_input)
+
+        self.assertTrue(torch.equal(D, D_input))
+        self.assertTrue(torch.equal(I, I_input))
+        self.assertTrue(torch.equal(R, R_input))
+
+    # tests sa_encode, sa_decode
+    def test_sa_encode_decode(self):
+        d = 16
+        index = faiss.IndexScalarQuantizer(d, faiss.ScalarQuantizer.QT_8bit)
+
+        xb = torch.rand(1000, d, dtype=torch.float32)
+        index.train(xb)
+
+        # torch cpu as ground truth
+        nq = 10
+        xq = torch.rand(nq, d, dtype=torch.float32)
+        encoded_torch = index.sa_encode(xq)
+
+        # numpy cpu
+        encoded_np = index.sa_encode(xq.numpy())
+
+        self.assertTrue(np.array_equal(encoded_torch.numpy(), encoded_np))
+
+        decoded_torch = index.sa_decode(encoded_torch)
+        decoded_np = index.sa_decode(encoded_np)
+
+        self.assertTrue(torch.equal(decoded_torch, torch.from_numpy(decoded_np)))
+
+        # torch cpu as output parameter
+        encoded_torch_param = torch.zeros(nq, d, dtype=torch.uint8)
+        index.sa_encode(xq, encoded_torch_param)
+
+        self.assertTrue(torch.equal(encoded_torch, encoded_torch))
+
+        decoded_torch_param = torch.zeros(nq, d, dtype=torch.float32)
+        index.sa_decode(encoded_torch, decoded_torch_param)
+
+        self.assertTrue(torch.equal(decoded_torch, decoded_torch_param))
+
+        # np as output parameter
+        encoded_np_param = np.zeros((nq, d), dtype=np.uint8)
+        index.sa_encode(xq.numpy(), encoded_np_param)
+
+        self.assertTrue(np.array_equal(encoded_torch.numpy(), encoded_np_param))
+
+        decoded_np_param = np.zeros((nq, d), dtype=np.float32)
+        index.sa_decode(encoded_np_param, decoded_np_param)
+
+        self.assertTrue(np.array_equal(decoded_np, decoded_np_param))
+
+    def test_non_contiguous(self):
+        d = 128
+        index = faiss.IndexFlatL2(d)
+
+        xb = torch.rand(d, 100).transpose(0, 1)
+
+        with self.assertRaises(AssertionError):
+            index.add(xb)
+
+        with self.assertRaises(ValueError):
+            index.add(xb.numpy())