Revert D80734790 (facebookresearch#4563)

Summary:
Pull Request resolved: facebookresearch#4563

This diff reverts D80734790
(The context such as a Sandcastle job, Task, SEV, etc. was not provided.)

Depends on D80734790

Reviewed By: limqiying

Differential Revision: D81147721

fbshipit-source-id: 33031f34660155a82f30cc2327a932e3f47ff141

Author: generatedunixname89002005232357

contrib/torch_utils.py

@@ -80,15 +80,6 @@ def swig_ptr_from_IndicesTensor(x):
     return faiss.cast_integer_to_idx_t_ptr(
         x.untyped_storage().data_ptr() + x.storage_offset() * 8)
 
-
-def swig_ptr_from_Int64Tensor(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_int64_ptr(
-        x.untyped_storage().data_ptr() + x.storage_offset() * 8)
-
-
 ##################################################################
 # utilities
 ##################################################################
@@ -168,7 +159,7 @@ def torch_replacement_add(self, x, numeric_type = faiss.Float32):
             # CPU torch
             self.add_ex(n, x_ptr, numeric_type)
 
-    def torch_replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32, ids_type = faiss.Int64):
+    def torch_replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32):
         if type(x) is np.ndarray:
             # forward to faiss __init__.py base method
             return self.add_with_ids_numpy(x, ids)
@@ -185,20 +176,17 @@ def torch_replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32, i
 
         assert type(ids) is torch.Tensor
         assert ids.shape == (n, ), 'not same number of vectors as ids'
-        if ids_type == faiss.Int64:
-            ids_ptr = swig_ptr_from_Int64Tensor(ids)
-        else:
-            raise ValueError("ids type for add_with_ids must be faiss.Int64")
+        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_ex(n, x_ptr, numeric_type, ids_ptr, ids_type)
+                self.add_with_ids_ex(n, x_ptr, numeric_type, ids_ptr)
         else:
             # CPU torch
-            self.add_with_ids_ex(n, x_ptr, numeric_type, ids_ptr, ids_type)
+            self.add_with_ids_ex(n, x_ptr, numeric_type, ids_ptr)
 
     def torch_replacement_assign(self, x, k, labels=None):
         if type(x) is np.ndarray:
@@ -272,18 +260,12 @@ def search_methods_common(x, k, D, I, numeric_type=faiss.Float32):
 
         if I is None:
             I = torch.empty(n, k, device=x.device, dtype=torch.int64)
-            I_ptr = swig_ptr_from_Int64Tensor(I)
-            I_type = faiss.Int64
         else:
             assert type(I) is torch.Tensor
             assert I.shape == (n, k)
-            if I.dtype == torch.int64:
-                I_ptr = swig_ptr_from_Int64Tensor(I)
-                I_type = faiss.Int64
-            else:
-                raise ValueError("labels for search should be int64 type")
+        I_ptr = swig_ptr_from_IndicesTensor(I)
 
-        return x_ptr, D_ptr, I_ptr, D, I, I_type
+        return x_ptr, D_ptr, I_ptr, D, I
 
     def torch_replacement_search(self, x, k, D=None, I=None, numeric_type=faiss.Float32):
         if type(x) is np.ndarray:
@@ -294,17 +276,17 @@ def torch_replacement_search(self, x, k, D=None, I=None, numeric_type=faiss.Floa
         n, d = x.shape
         assert d == self.d
 
-        x_ptr, D_ptr, I_ptr, D, I, I_type = search_methods_common(x, k, D, I)
+        x_ptr, D_ptr, I_ptr, D, I = search_methods_common(x, k, D, 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_ex(n, x_ptr, numeric_type, k, D_ptr, I_ptr, I_type)
+                self.search_ex(n, x_ptr, numeric_type, k, D_ptr, I_ptr)
         else:
             # CPU torch
-            self.search_ex(n, x_ptr, numeric_type, k, D_ptr, I_ptr, I_type)
+            self.search_ex(n, x_ptr, numeric_type, k, D_ptr, I_ptr)
 
         return D, I
 
@@ -317,7 +299,7 @@ def torch_replacement_search_and_reconstruct(self, x, k, D=None, I=None, R=None)
         n, d = x.shape
         assert d == self.d
 
-        x_ptr, D_ptr, I_ptr, D, I, _ = search_methods_common(x, k, D, I)
+        x_ptr, D_ptr, I_ptr, D, I = search_methods_common(x, k, D, I)
 
         if R is None:
             R = torch.empty(n, k, d, device=x.device, dtype=torch.float32)
@@ -347,7 +329,7 @@ def torch_replacement_search_preassigned(self, x, k, Iq, Dq, *, D=None, I=None):
         n, d = x.shape
         assert d == self.d
 
-        x_ptr, D_ptr, I_ptr, D, I, _ = search_methods_common(x, k, D, I)
+        x_ptr, D_ptr, I_ptr, D, I = search_methods_common(x, k, D, I)
 
         assert Iq.shape == (n, self.nprobe)
         Iq = Iq.contiguous()

faiss/Index.h

@@ -13,7 +13,6 @@
 #include <faiss/MetricType.h>
 #include <faiss/impl/FaissAssert.h>
 
-#include <cstdint>
 #include <cstdio>
 #include <sstream>
 
@@ -60,28 +59,22 @@ struct DistanceComputer;
 enum NumericType {
     Float32,
     Float16,
-    Int64,
     UInt8,
     Int8,
-    NONE, // corresponding to nullptr
 };
 
 inline size_t get_numeric_type_size(NumericType numeric_type) {
     switch (numeric_type) {
-        case NumericType::Int64:
-            return 8;
         case NumericType::Float32:
             return 4;
         case NumericType::Float16:
             return 2;
         case NumericType::UInt8:
         case NumericType::Int8:
             return 1;
-        case NumericType::NONE:
-            return 0;
         default:
             FAISS_THROW_MSG(
-                    "Unknown Numeric Type. Only supports Float32, Float16, Int64, UInt8, Int8");
+                    "Unknown Numeric Type. Only supports Float32, Float16");
     }
 }
 
@@ -176,17 +169,11 @@ struct Index {
             idx_t n,
             const void* x,
             NumericType numeric_type,
-            const void* xids,
-            NumericType xids_type) {
-        if (numeric_type == NumericType::Float32 &&
-            xids_type == NumericType::Int64) {
-            add_with_ids(
-                    n,
-                    static_cast<const float*>(x),
-                    static_cast<const idx_t*>(xids));
+            const idx_t* xids) {
+        if (numeric_type == NumericType::Float32) {
+            add_with_ids(n, static_cast<const float*>(x), xids);
         } else {
-            FAISS_THROW_MSG(
-                    "Index::add_with_ids: unsupported numeric type or xids type");
+            FAISS_THROW_MSG("Index::add_with_ids: unsupported numeric type");
         }
     }
 
@@ -215,20 +202,17 @@ struct Index {
             NumericType numeric_type,
             idx_t k,
             float* distances,
-            void* labels,
-            NumericType labels_type,
+            idx_t* labels,
             const SearchParameters* params = nullptr) const {
-        if (numeric_type == NumericType::Float32 &&
-            labels_type == NumericType::Int64) {
+        if (numeric_type == NumericType::Float32) {
             search(n,
                    static_cast<const float*>(x),
                    k,
                    distances,
-                   static_cast<idx_t*>(labels),
+                   labels,
                    params);
         } else {
-            FAISS_THROW_MSG(
-                    "Index::search: unsupported numeric type or label type");
+            FAISS_THROW_MSG("Index::search: unsupported numeric type");
         }
     }
 

faiss/IndexBinary.h

@@ -86,17 +86,12 @@ struct IndexBinary {
             idx_t n,
             const void* x,
             NumericType numeric_type,
-            const void* xids,
-            NumericType xids_type) {
-        if (numeric_type == NumericType::UInt8 &&
-            xids_type == NumericType::Int64) {
-            add_with_ids(
-                    n,
-                    static_cast<const uint8_t*>(x),
-                    static_cast<const idx_t*>(xids));
+            const idx_t* xids) {
+        if (numeric_type == NumericType::UInt8) {
+            add_with_ids(n, static_cast<const uint8_t*>(x), xids);
         } else {
             FAISS_THROW_MSG(
-                    "IndexBinary::add_with_ids: unsupported numeric type or xids type");
+                    "IndexBinary::add_with_ids: unsupported numeric type");
         }
     };
 
@@ -122,20 +117,17 @@ struct IndexBinary {
             NumericType numeric_type,
             idx_t k,
             int32_t* distances,
-            void* labels,
-            NumericType labels_type,
+            idx_t* labels,
             const SearchParameters* params = nullptr) const {
-        if (numeric_type == NumericType::UInt8 &&
-            labels_type == NumericType::Int64) {
+        if (numeric_type == NumericType::UInt8) {
             search(n,
                    static_cast<const uint8_t*>(x),
                    k,
                    distances,
-                   static_cast<idx_t*>(labels),
+                   labels,
                    params);
         } else {
-            FAISS_THROW_MSG(
-                    "IndexBinary::search: unsupported numeric type or label type");
+            FAISS_THROW_MSG("IndexBinary::search: unsupported numeric type");
         }
     };
 

faiss/IndexIDMap.cpp

@@ -12,7 +12,6 @@
 #include <cinttypes>
 #include <cstdint>
 #include <cstdio>
-#include "Index.h"
 #include "faiss/Index.h"
 
 #include <faiss/impl/AuxIndexStructures.h>
@@ -109,14 +108,10 @@ void IndexIDMapTemplate<IndexT>::add_with_ids_ex(
         idx_t n,
         const void* x,
         NumericType numeric_type,
-        const void* xids,
-        NumericType xids_type) {
-    FAISS_THROW_IF_NOT_MSG(
-            xids_type == NumericType::Int64,
-            "IndexIDMapTemplate::add_with_ids_ex only supports int64 as xids type");
+        const idx_t* xids) {
     index->add_ex(n, x, numeric_type);
     for (idx_t i = 0; i < n; i++) {
-        id_map.push_back(static_cast<const idx_t*>(xids)[i]);
+        id_map.push_back(xids[i]);
     }
     this->ntotal = index->ntotal;
 }
@@ -130,8 +125,7 @@ void IndexIDMapTemplate<IndexT>::add_with_ids(
             n,
             static_cast<const void*>(x),
             component_t_to_numeric<typename IndexT::component_t>(),
-            static_cast<const void*>(xids),
-            NumericType::Int64);
+            xids);
 }
 
 template <typename IndexT>
@@ -179,8 +173,7 @@ void IndexIDMapTemplate<IndexT>::search_ex(
         NumericType numeric_type,
         idx_t k,
         typename IndexT::distance_t* distances,
-        void* labels,
-        NumericType labels_type,
+        idx_t* labels,
         const SearchParameters* params) const {
     IDSelectorTranslated this_idtrans(this->id_map, nullptr);
     ScopedSelChange sel_change;
@@ -201,13 +194,8 @@ void IndexIDMapTemplate<IndexT>::search_ex(
             sel_change.set(params_non_const, &this_idtrans);
         }
     }
-    index->search_ex(
-            n, x, numeric_type, k, distances, labels, labels_type, params);
-
-    FAISS_THROW_IF_NOT_MSG(
-            labels_type == NumericType::Int64,
-            "IndexIDMapTemplate::search_ex only supports int64 as labels type");
-    idx_t* li = static_cast<idx_t*>(labels);
+    index->search_ex(n, x, numeric_type, k, distances, labels, params);
+    idx_t* li = labels;
 #pragma omp parallel for
     for (idx_t i = 0; i < n * k; i++) {
         li[i] = li[i] < 0 ? li[i] : id_map[li[i]];
@@ -228,8 +216,7 @@ void IndexIDMapTemplate<IndexT>::search(
             component_t_to_numeric<typename IndexT::component_t>(),
             k,
             distances,
-            static_cast<void*>(labels),
-            NumericType::Int64,
+            labels,
             params);
 }
 
@@ -319,11 +306,9 @@ void IndexIDMap2Template<IndexT>::add_with_ids_ex(
         idx_t n,
         const void* x,
         NumericType numeric_type,
-        const void* xids,
-        NumericType xids_type) {
+        const idx_t* xids) {
     size_t prev_ntotal = this->ntotal;
-    IndexIDMapTemplate<IndexT>::add_with_ids_ex(
-            n, x, numeric_type, xids, xids_type);
+    IndexIDMapTemplate<IndexT>::add_with_ids_ex(n, x, numeric_type, xids);
     for (size_t i = prev_ntotal; i < this->ntotal; i++) {
         rev_map[this->id_map[i]] = i;
     }
@@ -338,8 +323,7 @@ void IndexIDMap2Template<IndexT>::add_with_ids(
             n,
             static_cast<const void*>(x),
             component_t_to_numeric<typename IndexT::component_t>(),
-            static_cast<const void*>(xids),
-            NumericType::Int64);
+            xids);
 }
 
 template <typename IndexT>

faiss/IndexIDMap.h

@@ -35,8 +35,7 @@ struct IndexIDMapTemplate : IndexT {
             idx_t n,
             const void* x,
             NumericType numeric_type,
-            const void* xids,
-            NumericType xids_type) override;
+            const idx_t* xids) override;
 
     /// this will fail. Use add_with_ids
     void add(idx_t n, const component_t* x) override;
@@ -55,8 +54,7 @@ struct IndexIDMapTemplate : IndexT {
             NumericType numeric_type,
             idx_t k,
             distance_t* distances,
-            void* labels,
-            NumericType labels_type,
+            idx_t* labels,
             const SearchParameters* params = nullptr) const override;
 
     void train(idx_t n, const component_t* x) override;
@@ -110,8 +108,7 @@ struct IndexIDMap2Template : IndexIDMapTemplate<IndexT> {
             idx_t n,
             const void* x,
             NumericType numeric_type,
-            const void* xids,
-            NumericType xids_type) override;
+            const idx_t* xids) override;
 
     size_t remove_ids(const IDSelector& sel) override;