Enable clang-format for .cuh (CUDA headers) (#1726)

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

The diff enables clang-format for CUDA headers and applies it for fbsource

Reviewed By: zertosh

Differential Revision: D26695628

fbshipit-source-id: 30e53bfd6ad8aedd93c1b18076c5bd0a104a893f

faiss/gpu/impl/BinaryDistance.cuh

@@ -5,17 +5,20 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // Performs brute-force k-NN comparison between `vecs` and `query`, where they
 // are encoded as binary vectors
-void runBinaryDistance(Tensor<unsigned char, 2, true>& vecs,
-                       Tensor<unsigned char, 2, true>& query,
-                       Tensor<int, 2, true>& outK,
-                       Tensor<int, 2, true>& outV,
-                       int k, cudaStream_t stream);
+void runBinaryDistance(
+        Tensor<unsigned char, 2, true>& vecs,
+        Tensor<unsigned char, 2, true>& query,
+        Tensor<int, 2, true>& outK,
+        Tensor<int, 2, true>& outV,
+        int k,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/BinaryFlatIndex.cuh

@@ -5,61 +5,61 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 #include <faiss/gpu/utils/DeviceVector.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
 /// Holder of GPU resources for a particular flat index
 class BinaryFlatIndex {
- public:
-  BinaryFlatIndex(GpuResources* res,
-                  int dim,
-                  MemorySpace space);
+   public:
+    BinaryFlatIndex(GpuResources* res, int dim, MemorySpace space);
 
-  /// Returns the number of vectors we contain
-  int getSize() const;
+    /// Returns the number of vectors we contain
+    int getSize() const;
 
-  int getDim() const;
+    int getDim() const;
 
-  /// Reserve storage that can contain at least this many vectors
-  void reserve(size_t numVecs, cudaStream_t stream);
+    /// Reserve storage that can contain at least this many vectors
+    void reserve(size_t numVecs, cudaStream_t stream);
 
-  /// Returns a reference to our vectors currently in use
-  Tensor<unsigned char, 2, true>& getVectorsRef();
+    /// Returns a reference to our vectors currently in use
+    Tensor<unsigned char, 2, true>& getVectorsRef();
 
-  void query(Tensor<unsigned char, 2, true>& vecs,
-             int k,
-             Tensor<int, 2, true>& outDistances,
-             Tensor<int, 2, true>& outIndices);
+    void query(
+            Tensor<unsigned char, 2, true>& vecs,
+            int k,
+            Tensor<int, 2, true>& outDistances,
+            Tensor<int, 2, true>& outIndices);
 
-  /// Add vectors to ourselves; the pointer passed can be on the host
-  /// or the device
-  void add(const unsigned char* data, int numVecs, cudaStream_t stream);
+    /// Add vectors to ourselves; the pointer passed can be on the host
+    /// or the device
+    void add(const unsigned char* data, int numVecs, cudaStream_t stream);
 
-  /// Free all storage
-  void reset();
+    /// Free all storage
+    void reset();
 
- private:
-  /// Collection of GPU resources that we use
-  GpuResources* resources_;
+   private:
+    /// Collection of GPU resources that we use
+    GpuResources* resources_;
 
-  /// Dimensionality of our vectors
-  const int dim_;
+    /// Dimensionality of our vectors
+    const int dim_;
 
-  /// How many vectors we have
-  int num_;
+    /// How many vectors we have
+    int num_;
 
-  /// The underlying expandable storage
-  DeviceVector<char> rawData_;
+    /// The underlying expandable storage
+    DeviceVector<char> rawData_;
 
-  /// Vectors currently in rawData_
-  DeviceTensor<unsigned char, 2, true> vectors_;
+    /// Vectors currently in rawData_
+    DeviceTensor<unsigned char, 2, true> vectors_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/BroadcastSum.cuh

@@ -5,41 +5,48 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // output[x][i] += input[i] for all x
-void runSumAlongColumns(Tensor<float, 1, true>& input,
-                        Tensor<float, 2, true>& output,
-                        cudaStream_t stream);
+void runSumAlongColumns(
+        Tensor<float, 1, true>& input,
+        Tensor<float, 2, true>& output,
+        cudaStream_t stream);
 
-void runSumAlongColumns(Tensor<half, 1, true>& input,
-                        Tensor<half, 2, true>& output,
-                        cudaStream_t stream);
+void runSumAlongColumns(
+        Tensor<half, 1, true>& input,
+        Tensor<half, 2, true>& output,
+        cudaStream_t stream);
 
 // output[x][i] = input[i] for all x
-void runAssignAlongColumns(Tensor<float, 1, true>& input,
-                           Tensor<float, 2, true>& output,
-                           cudaStream_t stream);
+void runAssignAlongColumns(
+        Tensor<float, 1, true>& input,
+        Tensor<float, 2, true>& output,
+        cudaStream_t stream);
 
-void runAssignAlongColumns(Tensor<half, 1, true>& input,
-                           Tensor<half, 2, true>& output,
-                           cudaStream_t stream);
+void runAssignAlongColumns(
+        Tensor<half, 1, true>& input,
+        Tensor<half, 2, true>& output,
+        cudaStream_t stream);
 
 // output[i][x] += input[i] for all x
 // If zeroClamp, output[i][x] = max(output[i][x] + input[i], 0) for all x
-void runSumAlongRows(Tensor<float, 1, true>& input,
-                     Tensor<float, 2, true>& output,
-                     bool zeroClamp,
-                     cudaStream_t stream);
+void runSumAlongRows(
+        Tensor<float, 1, true>& input,
+        Tensor<float, 2, true>& output,
+        bool zeroClamp,
+        cudaStream_t stream);
 
-void runSumAlongRows(Tensor<half, 1, true>& input,
-                     Tensor<half, 2, true>& output,
-                     bool zeroClamp,
-                     cudaStream_t stream);
+void runSumAlongRows(
+        Tensor<half, 1, true>& input,
+        Tensor<half, 2, true>& output,
+        bool zeroClamp,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/Distance.cuh

@@ -5,64 +5,68 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
+#include <faiss/gpu/impl/GeneralDistance.cuh>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 #include <faiss/gpu/utils/Float16.cuh>
-#include <faiss/gpu/impl/GeneralDistance.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
 /// Calculates brute-force L2 distance between `vectors` and
 /// `queries`, returning the k closest results seen
-void runL2Distance(GpuResources* resources,
-                   Tensor<float, 2, true>& vectors,
-                   bool vectorsRowMajor,
-                   // can be optionally pre-computed; nullptr if we
-                   // have to compute it upon the call
-                   Tensor<float, 1, true>* vectorNorms,
-                   Tensor<float, 2, true>& queries,
-                   bool queriesRowMajor,
-                   int k,
-                   Tensor<float, 2, true>& outDistances,
-                   Tensor<int, 2, true>& outIndices,
-                   // Do we care about `outDistances`? If not, we can
-                   // take shortcuts.
-                   bool ignoreOutDistances = false);
+void runL2Distance(
+        GpuResources* resources,
+        Tensor<float, 2, true>& vectors,
+        bool vectorsRowMajor,
+        // can be optionally pre-computed; nullptr if we
+        // have to compute it upon the call
+        Tensor<float, 1, true>* vectorNorms,
+        Tensor<float, 2, true>& queries,
+        bool queriesRowMajor,
+        int k,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<int, 2, true>& outIndices,
+        // Do we care about `outDistances`? If not, we can
+        // take shortcuts.
+        bool ignoreOutDistances = false);
 
 /// Calculates brute-force inner product distance between `vectors`
 /// and `queries`, returning the k closest results seen
-void runIPDistance(GpuResources* resources,
-                   Tensor<float, 2, true>& vectors,
-                   bool vectorsRowMajor,
-                   Tensor<float, 2, true>& queries,
-                   bool queriesRowMajor,
-                   int k,
-                   Tensor<float, 2, true>& outDistances,
-                   Tensor<int, 2, true>& outIndices);
+void runIPDistance(
+        GpuResources* resources,
+        Tensor<float, 2, true>& vectors,
+        bool vectorsRowMajor,
+        Tensor<float, 2, true>& queries,
+        bool queriesRowMajor,
+        int k,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<int, 2, true>& outIndices);
 
-void runIPDistance(GpuResources* resources,
-                   Tensor<half, 2, true>& vectors,
-                   bool vectorsRowMajor,
-                   Tensor<half, 2, true>& queries,
-                   bool queriesRowMajor,
-                   int k,
-                   Tensor<float, 2, true>& outDistances,
-                   Tensor<int, 2, true>& outIndices);
+void runIPDistance(
+        GpuResources* resources,
+        Tensor<half, 2, true>& vectors,
+        bool vectorsRowMajor,
+        Tensor<half, 2, true>& queries,
+        bool queriesRowMajor,
+        int k,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<int, 2, true>& outIndices);
 
-void runL2Distance(GpuResources* resources,
-                   Tensor<half, 2, true>& vectors,
-                   bool vectorsRowMajor,
-                   Tensor<float, 1, true>* vectorNorms,
-                   Tensor<half, 2, true>& queries,
-                   bool queriesRowMajor,
-                   int k,
-                   Tensor<float, 2, true>& outDistances,
-                   Tensor<int, 2, true>& outIndices,
-                   bool ignoreOutDistances = false);
+void runL2Distance(
+        GpuResources* resources,
+        Tensor<half, 2, true>& vectors,
+        bool vectorsRowMajor,
+        Tensor<float, 1, true>* vectorNorms,
+        Tensor<half, 2, true>& queries,
+        bool queriesRowMajor,
+        int k,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<int, 2, true>& outIndices,
+        bool ignoreOutDistances = false);
 
 //
 // General distance implementation, assumes that all arguments are on the
@@ -70,131 +74,136 @@ void runL2Distance(GpuResources* resources,
 // based on metric type.
 //
 template <typename T>
-void bfKnnOnDevice(GpuResources* resources,
-                   int device,
-                   cudaStream_t stream,
-                   Tensor<T, 2, true>& vectors,
-                   bool vectorsRowMajor,
-                   Tensor<float, 1, true>* vectorNorms,
-                   Tensor<T, 2, true>& queries,
-                   bool queriesRowMajor,
-                   int k,
-                   faiss::MetricType metric,
-                   float metricArg,
-                   Tensor<float, 2, true>& outDistances,
-                   Tensor<int, 2, true>& outIndices,
-                   bool ignoreOutDistances) {
-  // We are guaranteed that all data arguments are resident on our preferred
-  // `device` here, and are ordered wrt `stream`
+void bfKnnOnDevice(
+        GpuResources* resources,
+        int device,
+        cudaStream_t stream,
+        Tensor<T, 2, true>& vectors,
+        bool vectorsRowMajor,
+        Tensor<float, 1, true>* vectorNorms,
+        Tensor<T, 2, true>& queries,
+        bool queriesRowMajor,
+        int k,
+        faiss::MetricType metric,
+        float metricArg,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<int, 2, true>& outIndices,
+        bool ignoreOutDistances) {
+    // We are guaranteed that all data arguments are resident on our preferred
+    // `device` here, and are ordered wrt `stream`
 
-  // L2 and IP are specialized to use GEMM and an optimized L2 + selection or
-  // pure k-selection kernel.
-  if ((metric == faiss::MetricType::METRIC_L2) ||
-      (metric == faiss::MetricType::METRIC_Lp &&
-       metricArg == 2)) {
-    runL2Distance(resources,
-                  vectors,
-                  vectorsRowMajor,
-                  vectorNorms,
-                  queries,
-                  queriesRowMajor,
-                  k,
-                  outDistances,
-                  outIndices);
-  } else if (metric == faiss::MetricType::METRIC_INNER_PRODUCT) {
-    runIPDistance(resources,
-                  vectors,
-                  vectorsRowMajor,
-                  queries,
-                  queriesRowMajor,
-                  k,
-                  outDistances,
-                  outIndices);
-  } else {
-    //
-    // General pairwise distance kernel
-    //
-    // The general distance kernel does not have specializations for
-    // transpositions (NN, NT, TN); instead, the transposition is just handled
-    // upon data load for now, which could result in poor data loading behavior
-    // for NT / TN. This can be fixed at a later date if desired, but efficiency
-    // is low versus GEMM anyways.
-    //
-
-    Tensor<T, 2> tVectorsDimInnermost =
-      vectorsRowMajor ?
-      vectors.transposeInnermost(1) :
-      vectors.transposeInnermost(0);
-    Tensor<T, 2> tQueriesDimInnermost =
-      queriesRowMajor ?
-      queries.transposeInnermost(1) :
-      queries.transposeInnermost(0);
-
-    if ((metric == faiss::MetricType::METRIC_L1) ||
-        (metric == faiss::MetricType::METRIC_Lp &&
-         metricArg == 1)) {
-      runGeneralDistance(resources,
-                         tVectorsDimInnermost,
-                         tQueriesDimInnermost,
-                         k,
-                         L1Distance(),
-                         outDistances,
-                         outIndices);
-    } else if (metric == faiss::MetricType::METRIC_Lp &&
-               metricArg == -1) {
-      // A way to test L2 distance
-      runGeneralDistance(resources,
-                         tVectorsDimInnermost,
-                         tQueriesDimInnermost,
-                         k,
-                         L2Distance(),
-                         outDistances,
-                         outIndices);
-    } else if (metric == faiss::MetricType::METRIC_Lp) {
-      runGeneralDistance(resources,
-                         tVectorsDimInnermost,
-                         tQueriesDimInnermost,
-                         k,
-                         LpDistance(metricArg),
-                         outDistances,
-                         outIndices);
-    } else if (metric == faiss::MetricType::METRIC_Linf) {
-      runGeneralDistance(resources,
-                         tVectorsDimInnermost,
-                         tQueriesDimInnermost,
-                         k,
-                         LinfDistance(),
-                         outDistances,
-                         outIndices);
-    } else if (metric == faiss::MetricType::METRIC_Canberra) {
-      runGeneralDistance(resources,
-                         tVectorsDimInnermost,
-                         tQueriesDimInnermost,
-                         k,
-                         CanberraDistance(),
-                         outDistances,
-                         outIndices);
-    } else if (metric == faiss::MetricType::METRIC_BrayCurtis) {
-      runGeneralDistance(resources,
-                         tVectorsDimInnermost,
-                         tQueriesDimInnermost,
-                         k,
-                         BrayCurtisDistance(),
-                         outDistances,
-                         outIndices);
-    } else if (metric == faiss::MetricType::METRIC_JensenShannon) {
-      runGeneralDistance(resources,
-                         tVectorsDimInnermost,
-                         tQueriesDimInnermost,
-                         k,
-                         JensenShannonDistance(),
-                         outDistances,
-                         outIndices);
+    // L2 and IP are specialized to use GEMM and an optimized L2 + selection or
+    // pure k-selection kernel.
+    if ((metric == faiss::MetricType::METRIC_L2) ||
+        (metric == faiss::MetricType::METRIC_Lp && metricArg == 2)) {
+        runL2Distance(
+                resources,
+                vectors,
+                vectorsRowMajor,
+                vectorNorms,
+                queries,
+                queriesRowMajor,
+                k,
+                outDistances,
+                outIndices);
+    } else if (metric == faiss::MetricType::METRIC_INNER_PRODUCT) {
+        runIPDistance(
+                resources,
+                vectors,
+                vectorsRowMajor,
+                queries,
+                queriesRowMajor,
+                k,
+                outDistances,
+                outIndices);
     } else {
-      FAISS_THROW_FMT("unsupported metric type %d", metric);
+        //
+        // General pairwise distance kernel
+        //
+        // The general distance kernel does not have specializations for
+        // transpositions (NN, NT, TN); instead, the transposition is just
+        // handled upon data load for now, which could result in poor data
+        // loading behavior for NT / TN. This can be fixed at a later date if
+        // desired, but efficiency is low versus GEMM anyways.
+        //
+
+        Tensor<T, 2> tVectorsDimInnermost = vectorsRowMajor
+                ? vectors.transposeInnermost(1)
+                : vectors.transposeInnermost(0);
+        Tensor<T, 2> tQueriesDimInnermost = queriesRowMajor
+                ? queries.transposeInnermost(1)
+                : queries.transposeInnermost(0);
+
+        if ((metric == faiss::MetricType::METRIC_L1) ||
+            (metric == faiss::MetricType::METRIC_Lp && metricArg == 1)) {
+            runGeneralDistance(
+                    resources,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    L1Distance(),
+                    outDistances,
+                    outIndices);
+        } else if (metric == faiss::MetricType::METRIC_Lp && metricArg == -1) {
+            // A way to test L2 distance
+            runGeneralDistance(
+                    resources,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    L2Distance(),
+                    outDistances,
+                    outIndices);
+        } else if (metric == faiss::MetricType::METRIC_Lp) {
+            runGeneralDistance(
+                    resources,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    LpDistance(metricArg),
+                    outDistances,
+                    outIndices);
+        } else if (metric == faiss::MetricType::METRIC_Linf) {
+            runGeneralDistance(
+                    resources,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    LinfDistance(),
+                    outDistances,
+                    outIndices);
+        } else if (metric == faiss::MetricType::METRIC_Canberra) {
+            runGeneralDistance(
+                    resources,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    CanberraDistance(),
+                    outDistances,
+                    outIndices);
+        } else if (metric == faiss::MetricType::METRIC_BrayCurtis) {
+            runGeneralDistance(
+                    resources,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    BrayCurtisDistance(),
+                    outDistances,
+                    outIndices);
+        } else if (metric == faiss::MetricType::METRIC_JensenShannon) {
+            runGeneralDistance(
+                    resources,
+                    tVectorsDimInnermost,
+                    tQueriesDimInnermost,
+                    k,
+                    JensenShannonDistance(),
+                    outDistances,
+                    outIndices);
+        } else {
+            FAISS_THROW_FMT("unsupported metric type %d", metric);
+        }
     }
-  }
 }
 
-
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/DistanceUtils.cuh

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Tensor.cuh>
@@ -14,330 +13,331 @@
 // Shared utilities for brute-force distance calculations
 //
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 struct IPDistance {
-  __host__ __device__ IPDistance() : dist(0) {}
+    __host__ __device__ IPDistance() : dist(0) {}
 
-  static constexpr bool kDirection = true; // maximize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = -std::numeric_limits<float>::max();
+    static constexpr bool kDirection = true; // maximize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = -std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    dist += a * b;
-  }
+    __host__ __device__ void handle(float a, float b) {
+        dist += a * b;
+    }
 
-  __host__ __device__ float reduce() {
-    return dist;
-  }
+    __host__ __device__ float reduce() {
+        return dist;
+    }
 
-  __host__ __device__ void combine(const IPDistance& v) {
-    dist += v.dist;
-  }
+    __host__ __device__ void combine(const IPDistance& v) {
+        dist += v.dist;
+    }
 
-  __host__ __device__ IPDistance zero() const {
-    return IPDistance();
-  }
+    __host__ __device__ IPDistance zero() const {
+        return IPDistance();
+    }
 
-  float dist;
+    float dist;
 };
 
 struct L1Distance {
-  __host__ __device__ L1Distance() : dist(0) {}
+    __host__ __device__ L1Distance() : dist(0) {}
 
-  static constexpr bool kDirection = false; // minimize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = std::numeric_limits<float>::max();
+    static constexpr bool kDirection = false; // minimize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    dist += fabsf(a - b);
-  }
+    __host__ __device__ void handle(float a, float b) {
+        dist += fabsf(a - b);
+    }
 
-  __host__ __device__ float reduce() {
-    return dist;
-  }
+    __host__ __device__ float reduce() {
+        return dist;
+    }
 
-  __host__ __device__ void combine(const L1Distance& v) {
-    dist += v.dist;
-  }
+    __host__ __device__ void combine(const L1Distance& v) {
+        dist += v.dist;
+    }
 
-  __host__ __device__ L1Distance zero() const {
-    return L1Distance();
-  }
+    __host__ __device__ L1Distance zero() const {
+        return L1Distance();
+    }
 
-  float dist;
+    float dist;
 };
 
 struct L2Distance {
-  __host__ __device__ L2Distance() : dist(0) {}
+    __host__ __device__ L2Distance() : dist(0) {}
 
-  static constexpr bool kDirection = false; // minimize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = std::numeric_limits<float>::max();
+    static constexpr bool kDirection = false; // minimize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    float v = a - b;
-    dist += v * v;
-  }
+    __host__ __device__ void handle(float a, float b) {
+        float v = a - b;
+        dist += v * v;
+    }
 
-  __host__ __device__ float reduce() {
-    return dist;
-  }
+    __host__ __device__ float reduce() {
+        return dist;
+    }
 
-  __host__ __device__ void combine(const L2Distance& v) {
-    dist += v.dist;
-  }
+    __host__ __device__ void combine(const L2Distance& v) {
+        dist += v.dist;
+    }
 
-  __host__ __device__ L2Distance zero() const {
-    return L2Distance();
-  }
+    __host__ __device__ L2Distance zero() const {
+        return L2Distance();
+    }
 
-  float dist;
+    float dist;
 };
 
 struct LpDistance {
-  __host__ __device__ LpDistance()
-      : p(2), dist(0) {}
+    __host__ __device__ LpDistance() : p(2), dist(0) {}
 
-  __host__ __device__ LpDistance(float arg)
-      : p(arg), dist(0) {}
+    __host__ __device__ LpDistance(float arg) : p(arg), dist(0) {}
 
-  __host__ __device__ LpDistance(const LpDistance& v)
-      : p(v.p), dist(v.dist) {}
+    __host__ __device__ LpDistance(const LpDistance& v)
+            : p(v.p), dist(v.dist) {}
 
-  __host__ __device__ LpDistance& operator=(const LpDistance& v) {
-    p = v.p;
-    dist = v.dist;
-    return *this;
-  }
+    __host__ __device__ LpDistance& operator=(const LpDistance& v) {
+        p = v.p;
+        dist = v.dist;
+        return *this;
+    }
 
-  static constexpr bool kDirection = false; // minimize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = std::numeric_limits<float>::max();
+    static constexpr bool kDirection = false; // minimize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    dist += powf(fabsf(a - b), p);
-  }
+    __host__ __device__ void handle(float a, float b) {
+        dist += powf(fabsf(a - b), p);
+    }
 
-  __host__ __device__ float reduce() {
-    return dist;
-  }
+    __host__ __device__ float reduce() {
+        return dist;
+    }
 
-  __host__ __device__ void combine(const LpDistance& v) {
-    dist += v.dist;
-  }
+    __host__ __device__ void combine(const LpDistance& v) {
+        dist += v.dist;
+    }
 
-  __host__ __device__ LpDistance zero() const {
-    return LpDistance(p);
-  }
+    __host__ __device__ LpDistance zero() const {
+        return LpDistance(p);
+    }
 
-  float p;
-  float dist;
+    float p;
+    float dist;
 };
 
 struct LinfDistance {
-  __host__ __device__ LinfDistance() : dist(0) {}
+    __host__ __device__ LinfDistance() : dist(0) {}
 
-  static constexpr bool kDirection = false; // minimize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = std::numeric_limits<float>::max();
+    static constexpr bool kDirection = false; // minimize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    dist = fmaxf(dist, fabsf(a - b));
-  }
+    __host__ __device__ void handle(float a, float b) {
+        dist = fmaxf(dist, fabsf(a - b));
+    }
 
-  __host__ __device__ float reduce() {
-    return dist;
-  }
+    __host__ __device__ float reduce() {
+        return dist;
+    }
 
-  __host__ __device__ void combine(const LinfDistance& v) {
-    dist = fmaxf(dist, v.dist);
-  }
+    __host__ __device__ void combine(const LinfDistance& v) {
+        dist = fmaxf(dist, v.dist);
+    }
 
-  __host__ __device__ LinfDistance zero() const {
-    return LinfDistance();
-  }
+    __host__ __device__ LinfDistance zero() const {
+        return LinfDistance();
+    }
 
-  float dist;
+    float dist;
 };
 
 struct CanberraDistance {
-  __host__ __device__ CanberraDistance() : dist(0) {}
+    __host__ __device__ CanberraDistance() : dist(0) {}
 
-  static constexpr bool kDirection = false; // minimize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = std::numeric_limits<float>::max();
+    static constexpr bool kDirection = false; // minimize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    float denom = fabsf(a) + fabsf(b);
-    dist += fabsf(a - b) / denom;
-  }
+    __host__ __device__ void handle(float a, float b) {
+        float denom = fabsf(a) + fabsf(b);
+        dist += fabsf(a - b) / denom;
+    }
 
-  __host__ __device__ float reduce() {
-    return dist;
-  }
+    __host__ __device__ float reduce() {
+        return dist;
+    }
 
-  __host__ __device__ void combine(const CanberraDistance& v) {
-    dist += v.dist;
-  }
+    __host__ __device__ void combine(const CanberraDistance& v) {
+        dist += v.dist;
+    }
 
-  __host__ __device__ CanberraDistance zero() const {
-    return CanberraDistance();
-  }
+    __host__ __device__ CanberraDistance zero() const {
+        return CanberraDistance();
+    }
 
-  float dist;
+    float dist;
 };
 
 struct BrayCurtisDistance {
-  __host__ __device__ BrayCurtisDistance()
-      : numerator(0), denominator(0) {}
+    __host__ __device__ BrayCurtisDistance() : numerator(0), denominator(0) {}
 
-  static constexpr bool kDirection = false; // minimize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = std::numeric_limits<float>::max();
+    static constexpr bool kDirection = false; // minimize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    numerator += fabsf(a - b);
-    denominator += fabsf(a + b);
-  }
+    __host__ __device__ void handle(float a, float b) {
+        numerator += fabsf(a - b);
+        denominator += fabsf(a + b);
+    }
 
-  __host__ __device__ float reduce() {
-    return (numerator / denominator);
-  }
+    __host__ __device__ float reduce() {
+        return (numerator / denominator);
+    }
 
-  __host__ __device__ void combine(const BrayCurtisDistance& v) {
-    numerator += v.numerator;
-    denominator += v.denominator;
-  }
+    __host__ __device__ void combine(const BrayCurtisDistance& v) {
+        numerator += v.numerator;
+        denominator += v.denominator;
+    }
 
-  __host__ __device__ BrayCurtisDistance zero() const {
-    return BrayCurtisDistance();
-  }
+    __host__ __device__ BrayCurtisDistance zero() const {
+        return BrayCurtisDistance();
+    }
 
-  float numerator;
-  float denominator;
+    float numerator;
+    float denominator;
 };
 
 struct JensenShannonDistance {
-  __host__ __device__ JensenShannonDistance()
-      : dist(0) {}
+    __host__ __device__ JensenShannonDistance() : dist(0) {}
 
-  static constexpr bool kDirection = false; // minimize
-  static constexpr float kIdentityData = 0;
-  static constexpr float kMaxDistance = std::numeric_limits<float>::max();
+    static constexpr bool kDirection = false; // minimize
+    static constexpr float kIdentityData = 0;
+    static constexpr float kMaxDistance = std::numeric_limits<float>::max();
 
-  __host__ __device__ void handle(float a, float b) {
-    float m = 0.5f * (a + b);
+    __host__ __device__ void handle(float a, float b) {
+        float m = 0.5f * (a + b);
 
-    float x = m / a;
-    float y = m / b;
+        float x = m / a;
+        float y = m / b;
 
-    float kl1 = -a * log(x);
-    float kl2 = -b * log(y);
+        float kl1 = -a * log(x);
+        float kl2 = -b * log(y);
 
-    dist += kl1 + kl2;
-  }
+        dist += kl1 + kl2;
+    }
 
-  __host__ __device__ float reduce() {
-    return 0.5 * dist;
-  }
+    __host__ __device__ float reduce() {
+        return 0.5 * dist;
+    }
 
-  __host__ __device__ void combine(const JensenShannonDistance& v) {
-    dist += v.dist;
-  }
+    __host__ __device__ void combine(const JensenShannonDistance& v) {
+        dist += v.dist;
+    }
 
-  __host__ __device__ JensenShannonDistance zero() const {
-    return JensenShannonDistance();
-  }
+    __host__ __device__ JensenShannonDistance zero() const {
+        return JensenShannonDistance();
+    }
 
-  float dist;
+    float dist;
 };
 
 template <typename T, bool InnerContig>
-Tensor<T, 2, InnerContig> sliceCentroids(Tensor<T, 2, InnerContig>& centroids,
-                                         bool centroidsRowMajor,
-                                         int startCentroid,
-                                         int num) {
-  // Row major is (num, dim)
-  // Col major is (dim, num)
-  if (startCentroid == 0 &&
-      num == centroids.getSize(centroidsRowMajor ? 0 : 1)) {
-    return centroids;
-  }
+Tensor<T, 2, InnerContig> sliceCentroids(
+        Tensor<T, 2, InnerContig>& centroids,
+        bool centroidsRowMajor,
+        int startCentroid,
+        int num) {
+    // Row major is (num, dim)
+    // Col major is (dim, num)
+    if (startCentroid == 0 &&
+        num == centroids.getSize(centroidsRowMajor ? 0 : 1)) {
+        return centroids;
+    }
 
-  return centroids.narrow(centroidsRowMajor ? 0 : 1, startCentroid, num);
+    return centroids.narrow(centroidsRowMajor ? 0 : 1, startCentroid, num);
 }
 
 // For each chunk of k indices, increment the index by chunk * increment
 template <typename T>
-__global__ void incrementIndex(Tensor<T, 2, true> indices,
-                               int k,
-                               int increment) {
-  for (int i = threadIdx.x; i < k; i += blockDim.x) {
-    indices[blockIdx.y][blockIdx.x * k + i] += blockIdx.x * increment;
-  }
+__global__ void incrementIndex(
+        Tensor<T, 2, true> indices,
+        int k,
+        int increment) {
+    for (int i = threadIdx.x; i < k; i += blockDim.x) {
+        indices[blockIdx.y][blockIdx.x * k + i] += blockIdx.x * increment;
+    }
 }
 
 // Used to update result indices in distance computation where the number of
 // centroids is high, and is tiled
 template <typename T>
-void runIncrementIndex(Tensor<T, 2, true>& indices,
-                       int k,
-                       int increment,
-                       cudaStream_t stream) {
-  dim3 grid(indices.getSize(1) / k, indices.getSize(0));
-  int block = std::min(k, 512);
+void runIncrementIndex(
+        Tensor<T, 2, true>& indices,
+        int k,
+        int increment,
+        cudaStream_t stream) {
+    dim3 grid(indices.getSize(1) / k, indices.getSize(0));
+    int block = std::min(k, 512);
 
-  // should be exact
-  FAISS_ASSERT(grid.x * k == indices.getSize(1));
+    // should be exact
+    FAISS_ASSERT(grid.x * k == indices.getSize(1));
 
-  incrementIndex<<<grid, block, 0, stream>>>(indices, k, increment);
+    incrementIndex<<<grid, block, 0, stream>>>(indices, k, increment);
 }
 
 // If the inner size (dim) of the vectors is small, we want a larger query tile
 // size, like 1024
-inline void chooseTileSize(int numQueries,
-                           int numCentroids,
-                           int dim,
-                           int elementSize,
-                           size_t tempMemAvailable,
-                           int& tileRows,
-                           int& tileCols) {
-  // The matrix multiplication should be large enough to be efficient, but if it
-  // is too large, we seem to lose efficiency as opposed to double-streaming.
-  // Each tile size here defines 1/2 of the memory use due to double streaming.
-  // We ignore available temporary memory, as that is adjusted independently by
-  // the user and can thus meet these requirements (or not).
-  // For <= 4 GB GPUs, prefer 512 MB of usage.
-  // For <= 8 GB GPUs, prefer 768 MB of usage.
-  // Otherwise, prefer 1 GB of usage.
-  auto totalMem = getCurrentDeviceProperties().totalGlobalMem;
+inline void chooseTileSize(
+        int numQueries,
+        int numCentroids,
+        int dim,
+        int elementSize,
+        size_t tempMemAvailable,
+        int& tileRows,
+        int& tileCols) {
+    // The matrix multiplication should be large enough to be efficient, but if
+    // it is too large, we seem to lose efficiency as opposed to
+    // double-streaming. Each tile size here defines 1/2 of the memory use due
+    // to double streaming. We ignore available temporary memory, as that is
+    // adjusted independently by the user and can thus meet these requirements
+    // (or not). For <= 4 GB GPUs, prefer 512 MB of usage. For <= 8 GB GPUs,
+    // prefer 768 MB of usage. Otherwise, prefer 1 GB of usage.
+    auto totalMem = getCurrentDeviceProperties().totalGlobalMem;
 
-  int targetUsage = 0;
+    int targetUsage = 0;
 
-  if (totalMem <= ((size_t) 4) * 1024 * 1024 * 1024) {
-    targetUsage = 512 * 1024 * 1024;
-  } else if (totalMem <= ((size_t) 8) * 1024 * 1024 * 1024) {
-    targetUsage = 768 * 1024 * 1024;
-  } else {
-    targetUsage = 1024 * 1024 * 1024;
-  }
+    if (totalMem <= ((size_t)4) * 1024 * 1024 * 1024) {
+        targetUsage = 512 * 1024 * 1024;
+    } else if (totalMem <= ((size_t)8) * 1024 * 1024 * 1024) {
+        targetUsage = 768 * 1024 * 1024;
+    } else {
+        targetUsage = 1024 * 1024 * 1024;
+    }
 
-  targetUsage /= 2 * elementSize;
+    targetUsage /= 2 * elementSize;
 
-  // 512 seems to be a batch size sweetspot for float32.
-  // If we are on float16, increase to 512.
-  // If the k size (vec dim) of the matrix multiplication is small (<= 32),
-  // increase to 1024.
-  int preferredTileRows = 512;
-  if (dim <= 32) {
-    preferredTileRows = 1024;
-  }
+    // 512 seems to be a batch size sweetspot for float32.
+    // If we are on float16, increase to 512.
+    // If the k size (vec dim) of the matrix multiplication is small (<= 32),
+    // increase to 1024.
+    int preferredTileRows = 512;
+    if (dim <= 32) {
+        preferredTileRows = 1024;
+    }
 
-  tileRows = std::min(preferredTileRows, numQueries);
+    tileRows = std::min(preferredTileRows, numQueries);
 
-  // tileCols is the remainder size
-  tileCols = std::min(targetUsage / preferredTileRows, numCentroids);
+    // tileCols is the remainder size
+    tileCols = std::min(targetUsage / preferredTileRows, numCentroids);
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/FlatIndex.cuh

@@ -5,127 +5,135 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/MetricType.h>
+#include <faiss/gpu/GpuResources.h>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 #include <faiss/gpu/utils/DeviceVector.cuh>
-#include <faiss/gpu/GpuResources.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
 /// Holder of GPU resources for a particular flat index
 class FlatIndex {
- public:
-  FlatIndex(GpuResources* res,
+   public:
+    FlatIndex(
+            GpuResources* res,
             int dim,
             bool useFloat16,
             bool storeTransposed,
             MemorySpace space);
 
-  /// Whether or not this flat index primarily stores data in float16
-  bool getUseFloat16() const;
+    /// Whether or not this flat index primarily stores data in float16
+    bool getUseFloat16() const;
 
-  /// Returns the number of vectors we contain
-  int getSize() const;
+    /// Returns the number of vectors we contain
+    int getSize() const;
 
-  /// Returns the dimensionality of the vectors
-  int getDim() const;
+    /// Returns the dimensionality of the vectors
+    int getDim() const;
 
-  /// Reserve storage that can contain at least this many vectors
-  void reserve(size_t numVecs, cudaStream_t stream);
+    /// Reserve storage that can contain at least this many vectors
+    void reserve(size_t numVecs, cudaStream_t stream);
 
-  /// Returns the vectors based on the type desired; the FlatIndex must be of
-  /// the same type (float16 or float32) to not assert
-  template <typename T>
-  Tensor<T, 2, true>& getVectorsRef();
+    /// Returns the vectors based on the type desired; the FlatIndex must be of
+    /// the same type (float16 or float32) to not assert
+    template <typename T>
+    Tensor<T, 2, true>& getVectorsRef();
 
-  /// Returns a reference to our vectors currently in use
-  Tensor<float, 2, true>& getVectorsFloat32Ref();
+    /// Returns a reference to our vectors currently in use
+    Tensor<float, 2, true>& getVectorsFloat32Ref();
 
-  /// Returns a reference to our vectors currently in use (useFloat16 mode)
-  Tensor<half, 2, true>& getVectorsFloat16Ref();
+    /// Returns a reference to our vectors currently in use (useFloat16 mode)
+    Tensor<half, 2, true>& getVectorsFloat16Ref();
 
-  /// Performs a copy of the vectors on the given device, converting
-  /// as needed from float16
-  DeviceTensor<float, 2, true> getVectorsFloat32Copy(cudaStream_t stream);
+    /// Performs a copy of the vectors on the given device, converting
+    /// as needed from float16
+    DeviceTensor<float, 2, true> getVectorsFloat32Copy(cudaStream_t stream);
 
-  /// Returns only a subset of the vectors
-  DeviceTensor<float, 2, true> getVectorsFloat32Copy(int from,
-                                                     int num,
-                                                     cudaStream_t stream);
+    /// Returns only a subset of the vectors
+    DeviceTensor<float, 2, true> getVectorsFloat32Copy(
+            int from,
+            int num,
+            cudaStream_t stream);
 
-  void query(Tensor<float, 2, true>& vecs,
-             int k,
-             faiss::MetricType metric,
-             float metricArg,
-             Tensor<float, 2, true>& outDistances,
-             Tensor<int, 2, true>& outIndices,
-             bool exactDistance);
+    void query(
+            Tensor<float, 2, true>& vecs,
+            int k,
+            faiss::MetricType metric,
+            float metricArg,
+            Tensor<float, 2, true>& outDistances,
+            Tensor<int, 2, true>& outIndices,
+            bool exactDistance);
 
-  void query(Tensor<half, 2, true>& vecs,
-             int k,
-             faiss::MetricType metric,
-             float metricArg,
-             Tensor<float, 2, true>& outDistances,
-             Tensor<int, 2, true>& outIndices,
-             bool exactDistance);
+    void query(
+            Tensor<half, 2, true>& vecs,
+            int k,
+            faiss::MetricType metric,
+            float metricArg,
+            Tensor<float, 2, true>& outDistances,
+            Tensor<int, 2, true>& outIndices,
+            bool exactDistance);
 
-  /// Compute residual for set of vectors
-  void computeResidual(Tensor<float, 2, true>& vecs,
-                       Tensor<int, 1, true>& listIds,
-                       Tensor<float, 2, true>& residuals);
+    /// Compute residual for set of vectors
+    void computeResidual(
+            Tensor<float, 2, true>& vecs,
+            Tensor<int, 1, true>& listIds,
+            Tensor<float, 2, true>& residuals);
 
-  /// Gather vectors given the set of IDs
-  void reconstruct(Tensor<int, 1, true>& listIds,
-                   Tensor<float, 2, true>& vecs);
+    /// Gather vectors given the set of IDs
+    void reconstruct(
+            Tensor<int, 1, true>& listIds,
+            Tensor<float, 2, true>& vecs);
 
-  void reconstruct(Tensor<int, 2, true>& listIds,
-                   Tensor<float, 3, true>& vecs);
+    void reconstruct(
+            Tensor<int, 2, true>& listIds,
+            Tensor<float, 3, true>& vecs);
 
-  /// Add vectors to ourselves; the pointer passed can be on the host
-  /// or the device
-  void add(const float* data, int numVecs, cudaStream_t stream);
+    /// Add vectors to ourselves; the pointer passed can be on the host
+    /// or the device
+    void add(const float* data, int numVecs, cudaStream_t stream);
 
-  /// Free all storage
-  void reset();
+    /// Free all storage
+    void reset();
 
- private:
-  /// Collection of GPU resources that we use
-  GpuResources* resources_;
+   private:
+    /// Collection of GPU resources that we use
+    GpuResources* resources_;
 
-  /// Dimensionality of our vectors
-  const int dim_;
+    /// Dimensionality of our vectors
+    const int dim_;
 
-  /// Float16 data format
-  const bool useFloat16_;
+    /// Float16 data format
+    const bool useFloat16_;
 
-  /// Store vectors in transposed layout for speed; makes addition to
-  /// the index slower
-  const bool storeTransposed_;
+    /// Store vectors in transposed layout for speed; makes addition to
+    /// the index slower
+    const bool storeTransposed_;
 
-  /// Memory space for our allocations
-  MemorySpace space_;
+    /// Memory space for our allocations
+    MemorySpace space_;
 
-  /// How many vectors we have
-  int num_;
+    /// How many vectors we have
+    int num_;
 
-  /// The underlying expandable storage
-  DeviceVector<char> rawData_;
+    /// The underlying expandable storage
+    DeviceVector<char> rawData_;
 
-  /// Vectors currently in rawData_
-  DeviceTensor<float, 2, true> vectors_;
-  DeviceTensor<float, 2, true> vectorsTransposed_;
+    /// Vectors currently in rawData_
+    DeviceTensor<float, 2, true> vectors_;
+    DeviceTensor<float, 2, true> vectorsTransposed_;
 
-  /// Vectors currently in rawData_, float16 form
-  DeviceTensor<half, 2, true> vectorsHalf_;
-  DeviceTensor<half, 2, true> vectorsHalfTransposed_;
+    /// Vectors currently in rawData_, float16 form
+    DeviceTensor<half, 2, true> vectorsHalf_;
+    DeviceTensor<half, 2, true> vectorsHalfTransposed_;
 
-  /// Precomputed L2 norms
-  DeviceTensor<float, 1, true> norms_;
+    /// Precomputed L2 norms
+    DeviceTensor<float, 1, true> norms_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/GeneralDistance.cuh

@@ -5,51 +5,51 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #include <faiss/MetricType.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/gpu/impl/DistanceUtils.cuh>
-#include <faiss/gpu/utils/ConversionOperators.cuh>
-#include <faiss/gpu/utils/DeviceTensor.cuh>
-#include <faiss/gpu/utils/DeviceDefs.cuh>
-#include <faiss/gpu/utils/DeviceUtils.h>
-#include <faiss/gpu/utils/Select.cuh>
 #include <faiss/gpu/utils/BlockSelectKernel.cuh>
+#include <faiss/gpu/utils/ConversionOperators.cuh>
+#include <faiss/gpu/utils/DeviceDefs.cuh>
+#include <faiss/gpu/utils/DeviceTensor.cuh>
+#include <faiss/gpu/utils/Select.cuh>
 
-#include <memory>
-#include <algorithm>
-#include <thrust/fill.h>
-#include <thrust/for_each.h>
 #include <thrust/device_ptr.h>
 #include <thrust/execution_policy.h>
+#include <thrust/fill.h>
+#include <thrust/for_each.h>
+#include <algorithm>
 #include <iostream>
+#include <memory>
 
 //
 // Kernels for non-L2 / inner product distances
 //
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // Reduction tree operator
 template <typename DistanceOp, int N>
 struct ReduceDistanceOp {
-  __device__ static DistanceOp reduce(DistanceOp ops[N]) {
-    DistanceOp vals[N/2];
+    __device__ static DistanceOp reduce(DistanceOp ops[N]) {
+        DistanceOp vals[N / 2];
 #pragma unroll
-    for (int i = 0; i < N / 2; ++i) {
-      vals[i] = ops[i * 2];
-      vals[i].combine(ops[i * 2 + 1]);
-    }
+        for (int i = 0; i < N / 2; ++i) {
+            vals[i] = ops[i * 2];
+            vals[i].combine(ops[i * 2 + 1]);
+        }
 
-    return ReduceDistanceOp<DistanceOp, N/2>::reduce(vals);
-  }
+        return ReduceDistanceOp<DistanceOp, N / 2>::reduce(vals);
+    }
 };
 
 template <typename DistanceOp>
 struct ReduceDistanceOp<DistanceOp, 1> {
-  __device__ static DistanceOp reduce(DistanceOp ops[1]) {
-    return ops[0];
-  }
+    __device__ static DistanceOp reduce(DistanceOp ops[1]) {
+        return ops[0];
+    }
 };
 
 // Implements a pairwise reduction tree
@@ -58,371 +58,401 @@ inline __device__ DistanceOp
 reduce(const DistanceOp& in,
        const T queryTile[kWarpSize][DimMultiple * kWarpSize + 1],
        const T vecTile[kWarpSize][DimMultiple * kWarpSize + 1]) {
-  DistanceOp accs[Unroll];
+    DistanceOp accs[Unroll];
 #pragma unroll
-  for (int i = 0; i < Unroll; ++i) {
-    accs[i] = in.zero();
-  }
-
-  auto vecTileBase = vecTile[threadIdx.x];
-  auto queryTileBase = queryTile[threadIdx.y];
-
-#pragma unroll
-  for (int i = 0; i < Unroll; ++i) {
-#pragma unroll
-    for (int j = 0; j < (kWarpSize * DimMultiple / Unroll); ++j) {
-      int idx = i * (kWarpSize * DimMultiple / Unroll) + j;
-      accs[i].handle(ConvertTo<float>::to(queryTileBase[idx]),
-                     ConvertTo<float>::to(vecTileBase[idx]));
+    for (int i = 0; i < Unroll; ++i) {
+        accs[i] = in.zero();
     }
-  }
 
-  return ReduceDistanceOp<DistanceOp, Unroll>::reduce(accs);
+    auto vecTileBase = vecTile[threadIdx.x];
+    auto queryTileBase = queryTile[threadIdx.y];
+
+#pragma unroll
+    for (int i = 0; i < Unroll; ++i) {
+#pragma unroll
+        for (int j = 0; j < (kWarpSize * DimMultiple / Unroll); ++j) {
+            int idx = i * (kWarpSize * DimMultiple / Unroll) + j;
+            accs[i].handle(
+                    ConvertTo<float>::to(queryTileBase[idx]),
+                    ConvertTo<float>::to(vecTileBase[idx]));
+        }
+    }
+
+    return ReduceDistanceOp<DistanceOp, Unroll>::reduce(accs);
 }
 
 // Our general distance matrix "multiplication" kernel
 template <typename T, typename DistanceOp, bool InnerContig>
-__launch_bounds__(kWarpSize * kWarpSize)
-__global__ void
-generalDistance(Tensor<T, 2, InnerContig> query, // m x k
-                Tensor<T, 2, InnerContig> vec, // n x k
-                DistanceOp op,
-                Tensor<float, 2, true> out) { // m x n
-  constexpr int kDimMultiple = 1;
+__launch_bounds__(kWarpSize* kWarpSize) __global__ void generalDistance(
+        Tensor<T, 2, InnerContig> query, // m x k
+        Tensor<T, 2, InnerContig> vec,   // n x k
+        DistanceOp op,
+        Tensor<float, 2, true> out) { // m x n
+    constexpr int kDimMultiple = 1;
 
-  __shared__ T queryTile[kWarpSize][kWarpSize * kDimMultiple + 1];
-  __shared__ T vecTile[kWarpSize][kWarpSize * kDimMultiple + 1];
+    __shared__ T queryTile[kWarpSize][kWarpSize * kDimMultiple + 1];
+    __shared__ T vecTile[kWarpSize][kWarpSize * kDimMultiple + 1];
 
-  // block y -> query
-  // block x -> vector
+    // block y -> query
+    // block x -> vector
 
-  int queryBlock = blockIdx.y * kWarpSize;
-  int queryThread = queryBlock + threadIdx.y;
+    int queryBlock = blockIdx.y * kWarpSize;
+    int queryThread = queryBlock + threadIdx.y;
 
-  int vecBlock = blockIdx.x * kWarpSize;
-  int vecThreadLoad = vecBlock + threadIdx.y;
-  int vecThreadSave = vecBlock + threadIdx.x;
+    int vecBlock = blockIdx.x * kWarpSize;
+    int vecThreadLoad = vecBlock + threadIdx.y;
+    int vecThreadSave = vecBlock + threadIdx.x;
 
-  DistanceOp acc = op.zero();
+    DistanceOp acc = op.zero();
 
-  auto queryTileBase = queryTile[threadIdx.y];
-  auto vecTileBase = vecTile[threadIdx.y];
+    auto queryTileBase = queryTile[threadIdx.y];
+    auto vecTileBase = vecTile[threadIdx.y];
 
-  auto queryBase = query[queryThread];
-  auto vecBase = vec[vecThreadLoad];
+    auto queryBase = query[queryThread];
+    auto vecBase = vec[vecThreadLoad];
 
-  if ((blockIdx.x != (gridDim.x - 1)) && (blockIdx.y != (gridDim.y - 1))) {
-    //
-    // Interior tile
-    //
-    int limit = utils::roundDown(query.getSize(1), kWarpSize * kDimMultiple);
+    if ((blockIdx.x != (gridDim.x - 1)) && (blockIdx.y != (gridDim.y - 1))) {
+        //
+        // Interior tile
+        //
+        int limit =
+                utils::roundDown(query.getSize(1), kWarpSize * kDimMultiple);
 
-    for (int k = threadIdx.x; k < limit; k += kWarpSize * kDimMultiple) {
-      // Load query tile
+        for (int k = threadIdx.x; k < limit; k += kWarpSize * kDimMultiple) {
+            // Load query tile
 #pragma unroll
-      for (int i = 0; i < kDimMultiple; ++i) {
-        queryTileBase[threadIdx.x + i * kWarpSize] =
-          queryBase[k + i * kWarpSize];
-        vecTileBase[threadIdx.x + i * kWarpSize] =
-          vecBase[k + i * kWarpSize];
-      }
+            for (int i = 0; i < kDimMultiple; ++i) {
+                queryTileBase[threadIdx.x + i * kWarpSize] =
+                        queryBase[k + i * kWarpSize];
+                vecTileBase[threadIdx.x + i * kWarpSize] =
+                        vecBase[k + i * kWarpSize];
+            }
 
-      __syncthreads();
+            __syncthreads();
 
-      // thread (y, x) does (query y, vec x)
-      acc.combine(
-        reduce<T, 8, kDimMultiple, DistanceOp>(op, queryTile, vecTile));
+            // thread (y, x) does (query y, vec x)
+            acc.combine(reduce<T, 8, kDimMultiple, DistanceOp>(
+                    op, queryTile, vecTile));
 
-      __syncthreads();
-    }
+            __syncthreads();
+        }
 
-    // Handle remainder
-    if (limit < query.getSize(1)) {
+        // Handle remainder
+        if (limit < query.getSize(1)) {
 #pragma unroll
-      for (int i = 0; i < kDimMultiple; ++i) {
-        int k = limit + threadIdx.x + i * kWarpSize;
-        bool kInBounds = k < query.getSize(1);
+            for (int i = 0; i < kDimMultiple; ++i) {
+                int k = limit + threadIdx.x + i * kWarpSize;
+                bool kInBounds = k < query.getSize(1);
 
-        queryTileBase[threadIdx.x + i * kWarpSize] =
-                             kInBounds ?
-                             queryBase[k] : ConvertTo<T>::to(0);
+                queryTileBase[threadIdx.x + i * kWarpSize] =
+                        kInBounds ? queryBase[k] : ConvertTo<T>::to(0);
 
-        vecTileBase[threadIdx.x + i * kWarpSize] =
-                             kInBounds ?
-                             vecBase[k] : ConvertTo<T>::to(0);
-      }
+                vecTileBase[threadIdx.x + i * kWarpSize] =
+                        kInBounds ? vecBase[k] : ConvertTo<T>::to(0);
+            }
 
-      __syncthreads();
+            __syncthreads();
 
-      int remainder = query.getSize(1) - limit;
+            int remainder = query.getSize(1) - limit;
 
-      // thread (y, x) does (query y, vec x)
+            // thread (y, x) does (query y, vec x)
 #pragma unroll
-      for (int i = 0; i < remainder; ++i) {
-        acc.handle(ConvertTo<float>::to(queryTileBase[i]),
-                   ConvertTo<float>::to(vecTile[threadIdx.x][i]));
-      }
-    }
+            for (int i = 0; i < remainder; ++i) {
+                acc.handle(
+                        ConvertTo<float>::to(queryTileBase[i]),
+                        ConvertTo<float>::to(vecTile[threadIdx.x][i]));
+            }
+        }
 
-    // Write out results
-    out[queryThread][vecThreadSave] = acc.reduce();
-  } else {
-    //
-    // Otherwise, we're an exterior tile
-    //
+        // Write out results
+        out[queryThread][vecThreadSave] = acc.reduce();
+    } else {
+        //
+        // Otherwise, we're an exterior tile
+        //
 
-    bool queryThreadInBounds = queryThread < query.getSize(0);
-    bool vecThreadInBoundsLoad = vecThreadLoad < vec.getSize(0);
-    bool vecThreadInBoundsSave = vecThreadSave < vec.getSize(0);
-    int limit = utils::roundDown(query.getSize(1), kWarpSize);
+        bool queryThreadInBounds = queryThread < query.getSize(0);
+        bool vecThreadInBoundsLoad = vecThreadLoad < vec.getSize(0);
+        bool vecThreadInBoundsSave = vecThreadSave < vec.getSize(0);
+        int limit = utils::roundDown(query.getSize(1), kWarpSize);
 
-    for (int k = threadIdx.x; k < limit; k += kWarpSize) {
-      // Load query tile
-      queryTileBase[threadIdx.x] =
-        queryThreadInBounds ?
-        queryBase[k] : ConvertTo<T>::to(0);
+        for (int k = threadIdx.x; k < limit; k += kWarpSize) {
+            // Load query tile
+            queryTileBase[threadIdx.x] =
+                    queryThreadInBounds ? queryBase[k] : ConvertTo<T>::to(0);
 
-      vecTileBase[threadIdx.x] =
-        vecThreadInBoundsLoad ?
-        vecBase[k] : ConvertTo<T>::to(0);
+            vecTileBase[threadIdx.x] =
+                    vecThreadInBoundsLoad ? vecBase[k] : ConvertTo<T>::to(0);
 
-      __syncthreads();
+            __syncthreads();
 
-      // thread (y, x) does (query y, vec x)
+            // thread (y, x) does (query y, vec x)
 #pragma unroll
-      for (int i = 0; i < kWarpSize; ++i) {
-        acc.handle(ConvertTo<float>::to(queryTileBase[i]),
-                   ConvertTo<float>::to(vecTile[threadIdx.x][i]));
-      }
+            for (int i = 0; i < kWarpSize; ++i) {
+                acc.handle(
+                        ConvertTo<float>::to(queryTileBase[i]),
+                        ConvertTo<float>::to(vecTile[threadIdx.x][i]));
+            }
 
-      __syncthreads();
+            __syncthreads();
+        }
+
+        // Handle remainder
+        if (limit < query.getSize(1)) {
+            int k = limit + threadIdx.x;
+            bool kInBounds = k < query.getSize(1);
+
+            // Load query tile
+            queryTileBase[threadIdx.x] = queryThreadInBounds && kInBounds
+                    ? queryBase[k]
+                    : ConvertTo<T>::to(0);
+
+            vecTileBase[threadIdx.x] = vecThreadInBoundsLoad && kInBounds
+                    ? vecBase[k]
+                    : ConvertTo<T>::to(0);
+
+            __syncthreads();
+
+            int remainder = query.getSize(1) - limit;
+
+            // thread (y, x) does (query y, vec x)
+            for (int i = 0; i < remainder; ++i) {
+                acc.handle(
+                        ConvertTo<float>::to(queryTileBase[i]),
+                        ConvertTo<float>::to(vecTile[threadIdx.x][i]));
+            }
+        }
+
+        // Write out results
+        if (queryThreadInBounds && vecThreadInBoundsSave) {
+            out[queryThread][vecThreadSave] = acc.reduce();
+        }
     }
-
-    // Handle remainder
-    if (limit < query.getSize(1)) {
-      int k = limit + threadIdx.x;
-      bool kInBounds = k < query.getSize(1);
-
-      // Load query tile
-      queryTileBase[threadIdx.x] =
-                           queryThreadInBounds && kInBounds ?
-                           queryBase[k] : ConvertTo<T>::to(0);
-
-      vecTileBase[threadIdx.x] =
-                           vecThreadInBoundsLoad && kInBounds ?
-                           vecBase[k] : ConvertTo<T>::to(0);
-
-      __syncthreads();
-
-      int remainder = query.getSize(1) - limit;
-
-      // thread (y, x) does (query y, vec x)
-      for (int i = 0; i < remainder; ++i) {
-        acc.handle(ConvertTo<float>::to(queryTileBase[i]),
-                   ConvertTo<float>::to(vecTile[threadIdx.x][i]));
-      }
-    }
-
-    // Write out results
-    if (queryThreadInBounds && vecThreadInBoundsSave) {
-      out[queryThread][vecThreadSave] = acc.reduce();
-    }
-  }
 }
 
-
 template <typename T, typename DistanceOp, bool InnerContig>
-void runGeneralDistanceKernel(Tensor<T, 2, InnerContig>& vecs,
-                              Tensor<T, 2, InnerContig>& query,
-                              Tensor<float, 2, true>& out,
-                              const DistanceOp& op,
-                              cudaStream_t stream) {
-  FAISS_ASSERT(vecs.getSize(1) == query.getSize(1));
-  FAISS_ASSERT(out.getSize(0) == query.getSize(0));
-  FAISS_ASSERT(out.getSize(1) == vecs.getSize(0));
+void runGeneralDistanceKernel(
+        Tensor<T, 2, InnerContig>& vecs,
+        Tensor<T, 2, InnerContig>& query,
+        Tensor<float, 2, true>& out,
+        const DistanceOp& op,
+        cudaStream_t stream) {
+    FAISS_ASSERT(vecs.getSize(1) == query.getSize(1));
+    FAISS_ASSERT(out.getSize(0) == query.getSize(0));
+    FAISS_ASSERT(out.getSize(1) == vecs.getSize(0));
 
-  dim3 grid(utils::divUp(vecs.getSize(0), kWarpSize),
+    dim3 grid(
+            utils::divUp(vecs.getSize(0), kWarpSize),
             utils::divUp(query.getSize(0), kWarpSize));
-  dim3 block(kWarpSize, kWarpSize);
+    dim3 block(kWarpSize, kWarpSize);
 
-  generalDistance<<<grid, block, 0, stream>>>(query, vecs, op, out);
+    generalDistance<<<grid, block, 0, stream>>>(query, vecs, op, out);
 }
 
 template <typename T, typename DistanceOp, bool InnerContig>
-void runGeneralDistance(GpuResources* res,
-                        Tensor<T, 2, InnerContig>& centroids,
-                        Tensor<T, 2, InnerContig>& queries,
-                        int k,
-                        const DistanceOp& op,
-                        Tensor<float, 2, true>& outDistances,
-                        Tensor<int, 2, true>& outIndices) {
-  // The # of centroids in `centroids` based on memory layout
-  auto numCentroids = centroids.getSize(0);
+void runGeneralDistance(
+        GpuResources* res,
+        Tensor<T, 2, InnerContig>& centroids,
+        Tensor<T, 2, InnerContig>& queries,
+        int k,
+        const DistanceOp& op,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<int, 2, true>& outIndices) {
+    // The # of centroids in `centroids` based on memory layout
+    auto numCentroids = centroids.getSize(0);
 
-  // The # of queries in `queries` based on memory layout
-  auto numQueries = queries.getSize(0);
+    // The # of queries in `queries` based on memory layout
+    auto numQueries = queries.getSize(0);
 
-  // The dimensions of the vectors to consider
-  auto dim = queries.getSize(1);
-  FAISS_ASSERT((numQueries == 0 || numCentroids == 0) ||
-               dim == centroids.getSize(1));
+    // The dimensions of the vectors to consider
+    auto dim = queries.getSize(1);
+    FAISS_ASSERT(
+            (numQueries == 0 || numCentroids == 0) ||
+            dim == centroids.getSize(1));
 
-  FAISS_ASSERT(outDistances.getSize(0) == numQueries);
-  FAISS_ASSERT(outIndices.getSize(0) == numQueries);
-  FAISS_ASSERT(outDistances.getSize(1) == k);
-  FAISS_ASSERT(outIndices.getSize(1) == k);
+    FAISS_ASSERT(outDistances.getSize(0) == numQueries);
+    FAISS_ASSERT(outIndices.getSize(0) == numQueries);
+    FAISS_ASSERT(outDistances.getSize(1) == k);
+    FAISS_ASSERT(outIndices.getSize(1) == k);
 
-  auto defaultStream = res->getDefaultStreamCurrentDevice();
+    auto defaultStream = res->getDefaultStreamCurrentDevice();
 
-  // If we're quering against a 0 sized set, just return empty results
-  if (centroids.numElements() == 0) {
-    thrust::fill(thrust::cuda::par.on(defaultStream),
-                 outDistances.data(), outDistances.end(),
-                 Limits<float>::getMax());
+    // If we're quering against a 0 sized set, just return empty results
+    if (centroids.numElements() == 0) {
+        thrust::fill(
+                thrust::cuda::par.on(defaultStream),
+                outDistances.data(),
+                outDistances.end(),
+                Limits<float>::getMax());
 
-    thrust::fill(thrust::cuda::par.on(defaultStream),
-                 outIndices.data(), outIndices.end(),
-                 -1);
+        thrust::fill(
+                thrust::cuda::par.on(defaultStream),
+                outIndices.data(),
+                outIndices.end(),
+                -1);
 
-    return;
-  }
-
-  // By default, aim to use up to 512 MB of memory for the processing, with both
-  // number of queries and number of centroids being at least 512.
-  int tileRows = 0;
-  int tileCols = 0;
-  chooseTileSize(numQueries,
-                 numCentroids,
-                 dim,
-                 sizeof(T),
-                 res->getTempMemoryAvailableCurrentDevice(),
-                 tileRows,
-                 tileCols);
-
-  int numColTiles = utils::divUp(numCentroids, tileCols);
-
-  // We can have any number of vectors to query against, even less than k, in
-  // which case we'll return -1 for the index
-  FAISS_ASSERT(k <= GPU_MAX_SELECTION_K); // select limitation
-
-  // Temporary output memory space we'll use
-  DeviceTensor<float, 2, true> distanceBuf1(
-    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, tileCols});
-  DeviceTensor<float, 2, true> distanceBuf2(
-    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, tileCols});
-  DeviceTensor<float, 2, true>* distanceBufs[2] =
-    {&distanceBuf1, &distanceBuf2};
-
-  DeviceTensor<float, 2, true> outDistanceBuf1(
-    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
-  DeviceTensor<float, 2, true> outDistanceBuf2(
-    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
-  DeviceTensor<float, 2, true>* outDistanceBufs[2] =
-    {&outDistanceBuf1, &outDistanceBuf2};
-
-  DeviceTensor<int, 2, true> outIndexBuf1(
-    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
-  DeviceTensor<int, 2, true> outIndexBuf2(
-    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
-  DeviceTensor<int, 2, true>* outIndexBufs[2] =
-    {&outIndexBuf1, &outIndexBuf2};
-
-  auto streams = res->getAlternateStreamsCurrentDevice();
-  streamWait(streams, {defaultStream});
-
-  int curStream = 0;
-  bool interrupt = false;
-
-  // Tile over the input queries
-  for (int i = 0; i < numQueries; i += tileRows) {
-    if (interrupt || InterruptCallback::is_interrupted()) {
-      interrupt = true;
-      break;
+        return;
     }
 
-    int curQuerySize = std::min(tileRows, numQueries - i);
+    // By default, aim to use up to 512 MB of memory for the processing, with
+    // both number of queries and number of centroids being at least 512.
+    int tileRows = 0;
+    int tileCols = 0;
+    chooseTileSize(
+            numQueries,
+            numCentroids,
+            dim,
+            sizeof(T),
+            res->getTempMemoryAvailableCurrentDevice(),
+            tileRows,
+            tileCols);
 
-    auto outDistanceView =
-      outDistances.narrow(0, i, curQuerySize);
-    auto outIndexView =
-      outIndices.narrow(0, i, curQuerySize);
+    int numColTiles = utils::divUp(numCentroids, tileCols);
 
-    auto queryView =
-      queries.narrow(0, i, curQuerySize);
+    // We can have any number of vectors to query against, even less than k, in
+    // which case we'll return -1 for the index
+    FAISS_ASSERT(k <= GPU_MAX_SELECTION_K); // select limitation
 
-    auto outDistanceBufRowView =
-      outDistanceBufs[curStream]->narrow(0, 0, curQuerySize);
-    auto outIndexBufRowView =
-      outIndexBufs[curStream]->narrow(0, 0, curQuerySize);
+    // Temporary output memory space we'll use
+    DeviceTensor<float, 2, true> distanceBuf1(
+            res,
+            makeTempAlloc(AllocType::Other, defaultStream),
+            {tileRows, tileCols});
+    DeviceTensor<float, 2, true> distanceBuf2(
+            res,
+            makeTempAlloc(AllocType::Other, defaultStream),
+            {tileRows, tileCols});
+    DeviceTensor<float, 2, true>* distanceBufs[2] = {
+            &distanceBuf1, &distanceBuf2};
 
-    // Tile over the centroids
-    for (int j = 0; j < numCentroids; j += tileCols) {
-      if (InterruptCallback::is_interrupted()) {
-        interrupt = true;
-        break;
-      }
+    DeviceTensor<float, 2, true> outDistanceBuf1(
+            res,
+            makeTempAlloc(AllocType::Other, defaultStream),
+            {tileRows, numColTiles * k});
+    DeviceTensor<float, 2, true> outDistanceBuf2(
+            res,
+            makeTempAlloc(AllocType::Other, defaultStream),
+            {tileRows, numColTiles * k});
+    DeviceTensor<float, 2, true>* outDistanceBufs[2] = {
+            &outDistanceBuf1, &outDistanceBuf2};
 
-      int curCentroidSize = std::min(tileCols, numCentroids - j);
-      int curColTile = j / tileCols;
+    DeviceTensor<int, 2, true> outIndexBuf1(
+            res,
+            makeTempAlloc(AllocType::Other, defaultStream),
+            {tileRows, numColTiles * k});
+    DeviceTensor<int, 2, true> outIndexBuf2(
+            res,
+            makeTempAlloc(AllocType::Other, defaultStream),
+            {tileRows, numColTiles * k});
+    DeviceTensor<int, 2, true>* outIndexBufs[2] = {
+            &outIndexBuf1, &outIndexBuf2};
 
-      auto centroidsView =
-        sliceCentroids(centroids, true, j, curCentroidSize);
+    auto streams = res->getAlternateStreamsCurrentDevice();
+    streamWait(streams, {defaultStream});
 
-      auto distanceBufView = distanceBufs[curStream]->
-        narrow(0, 0, curQuerySize).narrow(1, 0, curCentroidSize);
+    int curStream = 0;
+    bool interrupt = false;
 
-      auto outDistanceBufColView =
-        outDistanceBufRowView.narrow(1, k * curColTile, k);
-      auto outIndexBufColView =
-        outIndexBufRowView.narrow(1, k * curColTile, k);
+    // Tile over the input queries
+    for (int i = 0; i < numQueries; i += tileRows) {
+        if (interrupt || InterruptCallback::is_interrupted()) {
+            interrupt = true;
+            break;
+        }
 
-      runGeneralDistanceKernel(centroidsView,
-                               queryView,
-                               distanceBufView,
-                               op,
-                               streams[curStream]);
+        int curQuerySize = std::min(tileRows, numQueries - i);
 
-      // For IP, just k-select the output for this tile
-      if (tileCols == numCentroids) {
-        // Write into the final output
-        runBlockSelect(distanceBufView,
-                       outDistanceView,
-                       outIndexView,
-                       DistanceOp::kDirection, k, streams[curStream]);
-      } else {
-        // Write into the intermediate output
-        runBlockSelect(distanceBufView,
-                       outDistanceBufColView,
-                       outIndexBufColView,
-                       DistanceOp::kDirection, k, streams[curStream]);
-      }
+        auto outDistanceView = outDistances.narrow(0, i, curQuerySize);
+        auto outIndexView = outIndices.narrow(0, i, curQuerySize);
+
+        auto queryView = queries.narrow(0, i, curQuerySize);
+
+        auto outDistanceBufRowView =
+                outDistanceBufs[curStream]->narrow(0, 0, curQuerySize);
+        auto outIndexBufRowView =
+                outIndexBufs[curStream]->narrow(0, 0, curQuerySize);
+
+        // Tile over the centroids
+        for (int j = 0; j < numCentroids; j += tileCols) {
+            if (InterruptCallback::is_interrupted()) {
+                interrupt = true;
+                break;
+            }
+
+            int curCentroidSize = std::min(tileCols, numCentroids - j);
+            int curColTile = j / tileCols;
+
+            auto centroidsView =
+                    sliceCentroids(centroids, true, j, curCentroidSize);
+
+            auto distanceBufView = distanceBufs[curStream]
+                                           ->narrow(0, 0, curQuerySize)
+                                           .narrow(1, 0, curCentroidSize);
+
+            auto outDistanceBufColView =
+                    outDistanceBufRowView.narrow(1, k * curColTile, k);
+            auto outIndexBufColView =
+                    outIndexBufRowView.narrow(1, k * curColTile, k);
+
+            runGeneralDistanceKernel(
+                    centroidsView,
+                    queryView,
+                    distanceBufView,
+                    op,
+                    streams[curStream]);
+
+            // For IP, just k-select the output for this tile
+            if (tileCols == numCentroids) {
+                // Write into the final output
+                runBlockSelect(
+                        distanceBufView,
+                        outDistanceView,
+                        outIndexView,
+                        DistanceOp::kDirection,
+                        k,
+                        streams[curStream]);
+            } else {
+                // Write into the intermediate output
+                runBlockSelect(
+                        distanceBufView,
+                        outDistanceBufColView,
+                        outIndexBufColView,
+                        DistanceOp::kDirection,
+                        k,
+                        streams[curStream]);
+            }
+        }
+
+        // As we're finished with processing a full set of centroids, perform
+        // the final k-selection
+        if (tileCols != numCentroids) {
+            // The indices are tile-relative; for each tile of k, we need to add
+            // tileCols to the index
+            runIncrementIndex(
+                    outIndexBufRowView, k, tileCols, streams[curStream]);
+
+            runBlockSelectPair(
+                    outDistanceBufRowView,
+                    outIndexBufRowView,
+                    outDistanceView,
+                    outIndexView,
+                    DistanceOp::kDirection,
+                    k,
+                    streams[curStream]);
+        }
+
+        curStream = (curStream + 1) % 2;
     }
 
-    // As we're finished with processing a full set of centroids, perform the
-    // final k-selection
-    if (tileCols != numCentroids) {
-      // The indices are tile-relative; for each tile of k, we need to add
-      // tileCols to the index
-      runIncrementIndex(outIndexBufRowView, k, tileCols, streams[curStream]);
+    // Have the desired ordering stream wait on the multi-stream
+    streamWait({defaultStream}, streams);
 
-      runBlockSelectPair(outDistanceBufRowView,
-                         outIndexBufRowView,
-                         outDistanceView,
-                         outIndexView,
-                         DistanceOp::kDirection, k, streams[curStream]);
+    if (interrupt) {
+        FAISS_THROW_MSG("interrupted");
     }
 
-    curStream = (curStream + 1) % 2;
-  }
-
-  // Have the desired ordering stream wait on the multi-stream
-  streamWait({defaultStream}, streams);
-
-  if (interrupt) {
-    FAISS_THROW_MSG("interrupted");
-  }
-
-  CUDA_TEST_ERROR();
+    CUDA_TEST_ERROR();
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/IVFAppend.cuh

@@ -5,72 +5,79 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/gpu/impl/GpuScalarQuantizer.cuh>
 #include <faiss/gpu/GpuIndicesOptions.h>
-#include <faiss/gpu/utils/Tensor.cuh>
 #include <thrust/device_vector.h>
+#include <faiss/gpu/impl/GpuScalarQuantizer.cuh>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// Append user indices to IVF lists
-void runIVFIndicesAppend(Tensor<int, 1, true>& listIds,
-                         Tensor<int, 1, true>& listOffset,
-                         Tensor<Index::idx_t, 1, true>& indices,
-                         IndicesOptions opt,
-                         thrust::device_vector<void*>& listIndices,
-                         cudaStream_t stream);
+void runIVFIndicesAppend(
+        Tensor<int, 1, true>& listIds,
+        Tensor<int, 1, true>& listOffset,
+        Tensor<Index::idx_t, 1, true>& indices,
+        IndicesOptions opt,
+        thrust::device_vector<void*>& listIndices,
+        cudaStream_t stream);
 
 /// Update device-side list pointers in a batch
-void runUpdateListPointers(Tensor<int, 1, true>& listIds,
-                           Tensor<int, 1, true>& newListLength,
-                           Tensor<void*, 1, true>& newCodePointers,
-                           Tensor<void*, 1, true>& newIndexPointers,
-                           thrust::device_vector<int>& listLengths,
-                           thrust::device_vector<void*>& listCodes,
-                           thrust::device_vector<void*>& listIndices,
-                           cudaStream_t stream);
+void runUpdateListPointers(
+        Tensor<int, 1, true>& listIds,
+        Tensor<int, 1, true>& newListLength,
+        Tensor<void*, 1, true>& newCodePointers,
+        Tensor<void*, 1, true>& newIndexPointers,
+        thrust::device_vector<int>& listLengths,
+        thrust::device_vector<void*>& listCodes,
+        thrust::device_vector<void*>& listIndices,
+        cudaStream_t stream);
 
 /// Append PQ codes to IVF lists (non-interleaved format)
-void runIVFPQAppend(Tensor<int, 1, true>& listIds,
-                    Tensor<int, 1, true>& listOffset,
-                    Tensor<uint8_t, 2, true>& encodings,
-                    thrust::device_vector<void*>& listCodes,
-                    cudaStream_t stream);
+void runIVFPQAppend(
+        Tensor<int, 1, true>& listIds,
+        Tensor<int, 1, true>& listOffset,
+        Tensor<uint8_t, 2, true>& encodings,
+        thrust::device_vector<void*>& listCodes,
+        cudaStream_t stream);
 
 /// Append PQ codes to IVF lists (interleaved format)
-void runIVFPQInterleavedAppend(Tensor<int, 1, true>& listIds,
-                               Tensor<int, 1, true>& listOffset,
-                               Tensor<int, 1, true>& uniqueLists,
-                               Tensor<int, 1, true>& vectorsByUniqueList,
-                               Tensor<int, 1, true>& uniqueListVectorStart,
-                               Tensor<int, 1, true>& uniqueListStartOffset,
-                               int bitsPerCode,
-                               Tensor<uint8_t, 2, true>& encodings,
-                               thrust::device_vector<void*>& listCodes,
-                               cudaStream_t stream);
+void runIVFPQInterleavedAppend(
+        Tensor<int, 1, true>& listIds,
+        Tensor<int, 1, true>& listOffset,
+        Tensor<int, 1, true>& uniqueLists,
+        Tensor<int, 1, true>& vectorsByUniqueList,
+        Tensor<int, 1, true>& uniqueListVectorStart,
+        Tensor<int, 1, true>& uniqueListStartOffset,
+        int bitsPerCode,
+        Tensor<uint8_t, 2, true>& encodings,
+        thrust::device_vector<void*>& listCodes,
+        cudaStream_t stream);
 
 /// Append SQ codes to IVF lists (non-interleaved, old format)
-void runIVFFlatAppend(Tensor<int, 1, true>& listIds,
-                      Tensor<int, 1, true>& listOffset,
-                      Tensor<float, 2, true>& vecs,
-                      GpuScalarQuantizer* scalarQ,
-                      thrust::device_vector<void*>& listData,
-                      cudaStream_t stream);
+void runIVFFlatAppend(
+        Tensor<int, 1, true>& listIds,
+        Tensor<int, 1, true>& listOffset,
+        Tensor<float, 2, true>& vecs,
+        GpuScalarQuantizer* scalarQ,
+        thrust::device_vector<void*>& listData,
+        cudaStream_t stream);
 
 /// Append SQ codes to IVF lists (interleaved)
-void runIVFFlatInterleavedAppend(Tensor<int, 1, true>& listIds,
-                                 Tensor<int, 1, true>& listOffset,
-                                 Tensor<int, 1, true>& uniqueLists,
-                                 Tensor<int, 1, true>& vectorsByUniqueList,
-                                 Tensor<int, 1, true>& uniqueListVectorStart,
-                                 Tensor<int, 1, true>& uniqueListStartOffset,
-                                 Tensor<float, 2, true>& vecs,
-                                 GpuScalarQuantizer* scalarQ,
-                                 thrust::device_vector<void*>& listData,
-                                 GpuResources* res,
-                                 cudaStream_t stream);
+void runIVFFlatInterleavedAppend(
+        Tensor<int, 1, true>& listIds,
+        Tensor<int, 1, true>& listOffset,
+        Tensor<int, 1, true>& uniqueLists,
+        Tensor<int, 1, true>& vectorsByUniqueList,
+        Tensor<int, 1, true>& uniqueListVectorStart,
+        Tensor<int, 1, true>& uniqueListStartOffset,
+        Tensor<float, 2, true>& vecs,
+        GpuScalarQuantizer* scalarQ,
+        thrust::device_vector<void*>& listData,
+        GpuResources* res,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/IVFBase.cuh

@@ -5,205 +5,215 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/Index.h>
 #include <faiss/MetricType.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
-#include <faiss/gpu/utils/DeviceVector.cuh>
-#include <faiss/gpu/utils/DeviceTensor.cuh>
-#include <memory>
 #include <thrust/device_vector.h>
+#include <faiss/gpu/utils/DeviceTensor.cuh>
+#include <faiss/gpu/utils/DeviceVector.cuh>
+#include <memory>
 #include <vector>
 
-namespace faiss { struct InvertedLists; }
+namespace faiss {
+struct InvertedLists;
+}
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 class FlatIndex;
 
 /// Base inverted list functionality for IVFFlat and IVFPQ
 class IVFBase {
- public:
-  IVFBase(GpuResources* resources,
-          faiss::MetricType metric,
-          float metricArg,
-          /// We do not own this reference
-          FlatIndex* quantizer,
-          bool interleavedLayout,
-          IndicesOptions indicesOptions,
-          MemorySpace space);
+   public:
+    IVFBase(GpuResources* resources,
+            faiss::MetricType metric,
+            float metricArg,
+            /// We do not own this reference
+            FlatIndex* quantizer,
+            bool interleavedLayout,
+            IndicesOptions indicesOptions,
+            MemorySpace space);
 
-  virtual ~IVFBase();
+    virtual ~IVFBase();
 
-  /// Reserve GPU memory in our inverted lists for this number of vectors
-  void reserveMemory(size_t numVecs);
+    /// Reserve GPU memory in our inverted lists for this number of vectors
+    void reserveMemory(size_t numVecs);
 
-  /// Clear out all inverted lists, but retain the coarse quantizer
-  /// and the product quantizer info
-  void reset();
+    /// Clear out all inverted lists, but retain the coarse quantizer
+    /// and the product quantizer info
+    void reset();
 
-  /// Return the number of dimensions we are indexing
-  int getDim() const;
+    /// Return the number of dimensions we are indexing
+    int getDim() const;
 
-  /// After adding vectors, one can call this to reclaim device memory
-  /// to exactly the amount needed. Returns space reclaimed in bytes
-  size_t reclaimMemory();
+    /// After adding vectors, one can call this to reclaim device memory
+    /// to exactly the amount needed. Returns space reclaimed in bytes
+    size_t reclaimMemory();
 
-  /// Returns the number of inverted lists
-  size_t getNumLists() const;
+    /// Returns the number of inverted lists
+    size_t getNumLists() const;
 
-  /// For debugging purposes, return the list length of a particular
-  /// list
-  int getListLength(int listId) const;
+    /// For debugging purposes, return the list length of a particular
+    /// list
+    int getListLength(int listId) const;
 
-  /// Return the list indices of a particular list back to the CPU
-  std::vector<Index::idx_t> getListIndices(int listId) const;
+    /// Return the list indices of a particular list back to the CPU
+    std::vector<Index::idx_t> getListIndices(int listId) const;
 
-  /// Return the encoded vectors of a particular list back to the CPU
-  std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat) const;
+    /// Return the encoded vectors of a particular list back to the CPU
+    std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat) const;
 
-  /// Copy all inverted lists from a CPU representation to ourselves
-  void copyInvertedListsFrom(const InvertedLists* ivf);
+    /// Copy all inverted lists from a CPU representation to ourselves
+    void copyInvertedListsFrom(const InvertedLists* ivf);
 
-  /// Copy all inverted lists from ourselves to a CPU representation
-  void copyInvertedListsTo(InvertedLists* ivf);
+    /// Copy all inverted lists from ourselves to a CPU representation
+    void copyInvertedListsTo(InvertedLists* ivf);
 
-  /// Classify and encode/add vectors to our IVF lists.
-  /// The input data must be on our current device.
-  /// Returns the number of vectors successfully added. Vectors may
-  /// not be able to be added because they contain NaNs.
-  int addVectors(Tensor<float, 2, true>& vecs,
-                 Tensor<Index::idx_t, 1, true>& indices);
+    /// Classify and encode/add vectors to our IVF lists.
+    /// The input data must be on our current device.
+    /// Returns the number of vectors successfully added. Vectors may
+    /// not be able to be added because they contain NaNs.
+    int addVectors(
+            Tensor<float, 2, true>& vecs,
+            Tensor<Index::idx_t, 1, true>& indices);
 
- protected:
-  /// Adds a set of codes and indices to a list, with the representation coming
-  /// from the CPU equivalent
-  void addEncodedVectorsToList_(int listId,
-                                // resident on the host
-                                const void* codes,
-                                // resident on the host
-                                const Index::idx_t* indices,
-                                size_t numVecs);
+   protected:
+    /// Adds a set of codes and indices to a list, with the representation
+    /// coming from the CPU equivalent
+    void addEncodedVectorsToList_(
+            int listId,
+            // resident on the host
+            const void* codes,
+            // resident on the host
+            const Index::idx_t* indices,
+            size_t numVecs);
 
-  /// Returns the number of bytes in which an IVF list containing numVecs
-  /// vectors is encoded on the device. Note that due to padding this is not the
-  /// same as the encoding size for a subset of vectors in an IVF list; this is
-  /// the size for an entire IVF list
-  virtual size_t getGpuVectorsEncodingSize_(int numVecs) const = 0;
-  virtual size_t getCpuVectorsEncodingSize_(int numVecs) const = 0;
+    /// Returns the number of bytes in which an IVF list containing numVecs
+    /// vectors is encoded on the device. Note that due to padding this is not
+    /// the same as the encoding size for a subset of vectors in an IVF list;
+    /// this is the size for an entire IVF list
+    virtual size_t getGpuVectorsEncodingSize_(int numVecs) const = 0;
+    virtual size_t getCpuVectorsEncodingSize_(int numVecs) const = 0;
 
-  /// Translate to our preferred GPU encoding
-  virtual std::vector<uint8_t> translateCodesToGpu_(std::vector<uint8_t> codes,
-                                                    size_t numVecs) const = 0;
+    /// Translate to our preferred GPU encoding
+    virtual std::vector<uint8_t> translateCodesToGpu_(
+            std::vector<uint8_t> codes,
+            size_t numVecs) const = 0;
 
-  /// Translate from our preferred GPU encoding
-  virtual std::vector<uint8_t> translateCodesFromGpu_(std::vector<uint8_t> codes,
-                                                      size_t numVecs) const = 0;
+    /// Translate from our preferred GPU encoding
+    virtual std::vector<uint8_t> translateCodesFromGpu_(
+            std::vector<uint8_t> codes,
+            size_t numVecs) const = 0;
 
-  /// Append vectors to our on-device lists
-  virtual void appendVectors_(Tensor<float, 2, true>& vecs,
-                              Tensor<Index::idx_t, 1, true>& indices,
-                              Tensor<int, 1, true>& uniqueLists,
-                              Tensor<int, 1, true>& vectorsByUniqueList,
-                              Tensor<int, 1, true>& uniqueListVectorStart,
-                              Tensor<int, 1, true>& uniqueListStartOffset,
-                              Tensor<int, 1, true>& listIds,
-                              Tensor<int, 1, true>& listOffset,
-                              cudaStream_t stream) = 0;
+    /// Append vectors to our on-device lists
+    virtual void appendVectors_(
+            Tensor<float, 2, true>& vecs,
+            Tensor<Index::idx_t, 1, true>& indices,
+            Tensor<int, 1, true>& uniqueLists,
+            Tensor<int, 1, true>& vectorsByUniqueList,
+            Tensor<int, 1, true>& uniqueListVectorStart,
+            Tensor<int, 1, true>& uniqueListStartOffset,
+            Tensor<int, 1, true>& listIds,
+            Tensor<int, 1, true>& listOffset,
+            cudaStream_t stream) = 0;
 
-  /// Reclaim memory consumed on the device for our inverted lists
-  /// `exact` means we trim exactly to the memory needed
-  size_t reclaimMemory_(bool exact);
+    /// Reclaim memory consumed on the device for our inverted lists
+    /// `exact` means we trim exactly to the memory needed
+    size_t reclaimMemory_(bool exact);
 
-  /// Update all device-side list pointer and size information
-  void updateDeviceListInfo_(cudaStream_t stream);
+    /// Update all device-side list pointer and size information
+    void updateDeviceListInfo_(cudaStream_t stream);
 
-  /// For a set of list IDs, update device-side list pointer and size
-  /// information
-  void updateDeviceListInfo_(const std::vector<int>& listIds,
-                             cudaStream_t stream);
+    /// For a set of list IDs, update device-side list pointer and size
+    /// information
+    void updateDeviceListInfo_(
+            const std::vector<int>& listIds,
+            cudaStream_t stream);
 
-  /// Shared function to copy indices from CPU to GPU
-  void addIndicesFromCpu_(int listId,
-                          const Index::idx_t* indices,
-                          size_t numVecs);
+    /// Shared function to copy indices from CPU to GPU
+    void addIndicesFromCpu_(
+            int listId,
+            const Index::idx_t* indices,
+            size_t numVecs);
 
- protected:
-  /// Collection of GPU resources that we use
-  GpuResources* resources_;
+   protected:
+    /// Collection of GPU resources that we use
+    GpuResources* resources_;
 
-  /// Metric type of the index
-  faiss::MetricType metric_;
+    /// Metric type of the index
+    faiss::MetricType metric_;
 
-  /// Metric arg
-  float metricArg_;
+    /// Metric arg
+    float metricArg_;
 
-  /// Quantizer object
-  FlatIndex* quantizer_;
+    /// Quantizer object
+    FlatIndex* quantizer_;
 
-  /// Expected dimensionality of the vectors
-  const int dim_;
+    /// Expected dimensionality of the vectors
+    const int dim_;
 
-  /// Number of inverted lists we maintain
-  const int numLists_;
+    /// Number of inverted lists we maintain
+    const int numLists_;
 
-  /// Whether or not our index uses an interleaved by 32 layout:
-  /// The default memory layout is [vector][PQ/SQ component]:
-  /// (v0 d0) (v0 d1) ... (v0 dD-1) (v1 d0) (v1 d1) ...
-  ///
-  /// The interleaved by 32 memory layout is:
-  /// [vector / 32][PQ/SQ component][vector % 32] with padding:
-  /// (v0 d0) (v1 d0) ... (v31 d0) (v0 d1) (v1 d1) ... (v31 dD-1) (v32 d0) (v33
-  /// d0) ...
-  /// so the list length is always a multiple of num quantizers * 32
-  bool interleavedLayout_;
+    /// Whether or not our index uses an interleaved by 32 layout:
+    /// The default memory layout is [vector][PQ/SQ component]:
+    /// (v0 d0) (v0 d1) ... (v0 dD-1) (v1 d0) (v1 d1) ...
+    ///
+    /// The interleaved by 32 memory layout is:
+    /// [vector / 32][PQ/SQ component][vector % 32] with padding:
+    /// (v0 d0) (v1 d0) ... (v31 d0) (v0 d1) (v1 d1) ... (v31 dD-1) (v32 d0)
+    /// (v33 d0) ... so the list length is always a multiple of num quantizers *
+    /// 32
+    bool interleavedLayout_;
 
-  /// How are user indices stored on the GPU?
-  const IndicesOptions indicesOptions_;
+    /// How are user indices stored on the GPU?
+    const IndicesOptions indicesOptions_;
 
-  /// What memory space our inverted list storage is in
-  const MemorySpace space_;
+    /// What memory space our inverted list storage is in
+    const MemorySpace space_;
 
-  /// Device representation of all inverted list data
-  /// id -> data
-  thrust::device_vector<void*> deviceListDataPointers_;
+    /// Device representation of all inverted list data
+    /// id -> data
+    thrust::device_vector<void*> deviceListDataPointers_;
 
-  /// Device representation of all inverted list index pointers
-  /// id -> data
-  thrust::device_vector<void*> deviceListIndexPointers_;
+    /// Device representation of all inverted list index pointers
+    /// id -> data
+    thrust::device_vector<void*> deviceListIndexPointers_;
 
-  /// Device representation of all inverted list lengths
-  /// id -> length in number of vectors
-  thrust::device_vector<int> deviceListLengths_;
+    /// Device representation of all inverted list lengths
+    /// id -> length in number of vectors
+    thrust::device_vector<int> deviceListLengths_;
 
-  /// Maximum list length seen
-  int maxListLength_;
+    /// Maximum list length seen
+    int maxListLength_;
 
-  struct DeviceIVFList {
-    DeviceIVFList(GpuResources* res, const AllocInfo& info);
+    struct DeviceIVFList {
+        DeviceIVFList(GpuResources* res, const AllocInfo& info);
 
-    /// The on-device memory for this particular IVF list
-    DeviceVector<uint8_t> data;
+        /// The on-device memory for this particular IVF list
+        DeviceVector<uint8_t> data;
 
-    /// The number of vectors encoded in this list, which may be unrelated to
-    /// the above allocated data size
-    int numVecs;
-  };
+        /// The number of vectors encoded in this list, which may be unrelated
+        /// to the above allocated data size
+        int numVecs;
+    };
 
-  /// Device memory for each separate list, as managed by the host.
-  /// Device memory as stored in DeviceVector is stored as unique_ptr
-  /// since deviceList*Pointers_ must remain valid despite
-  /// resizing (and potential re-allocation) of deviceList*_
-  std::vector<std::unique_ptr<DeviceIVFList>> deviceListData_;
-  std::vector<std::unique_ptr<DeviceIVFList>> deviceListIndices_;
+    /// Device memory for each separate list, as managed by the host.
+    /// Device memory as stored in DeviceVector is stored as unique_ptr
+    /// since deviceList*Pointers_ must remain valid despite
+    /// resizing (and potential re-allocation) of deviceList*_
+    std::vector<std::unique_ptr<DeviceIVFList>> deviceListData_;
+    std::vector<std::unique_ptr<DeviceIVFList>> deviceListIndices_;
 
-  /// If we are storing indices on the CPU (indicesOptions_ is
-  /// INDICES_CPU), then this maintains a CPU-side map of what
-  /// (inverted list id, offset) maps to which user index
-  std::vector<std::vector<Index::idx_t>> listOffsetToUserIndex_;
+    /// If we are storing indices on the CPU (indicesOptions_ is
+    /// INDICES_CPU), then this maintains a CPU-side map of what
+    /// (inverted list id, offset) maps to which user index
+    std::vector<std::vector<Index::idx_t>> listOffsetToUserIndex_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/IVFFlat.cuh

@@ -5,72 +5,77 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/gpu/impl/IVFBase.cuh>
 #include <faiss/gpu/impl/GpuScalarQuantizer.cuh>
+#include <faiss/gpu/impl/IVFBase.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class IVFFlat : public IVFBase {
- public:
-  /// Construct from a quantizer that has elemen
-  IVFFlat(GpuResources* resources,
-          /// We do not own this reference
-          FlatIndex* quantizer,
-          faiss::MetricType metric,
-          float metricArg,
-          bool useResidual,
-          /// Optional ScalarQuantizer
-          faiss::ScalarQuantizer* scalarQ,
-          bool interleavedLayout,
-          IndicesOptions indicesOptions,
-          MemorySpace space);
+   public:
+    /// Construct from a quantizer that has elemen
+    IVFFlat(GpuResources* resources,
+            /// We do not own this reference
+            FlatIndex* quantizer,
+            faiss::MetricType metric,
+            float metricArg,
+            bool useResidual,
+            /// Optional ScalarQuantizer
+            faiss::ScalarQuantizer* scalarQ,
+            bool interleavedLayout,
+            IndicesOptions indicesOptions,
+            MemorySpace space);
 
-  ~IVFFlat() override;
+    ~IVFFlat() override;
 
-  /// Find the approximate k nearest neigbors for `queries` against
-  /// our database
-  void query(Tensor<float, 2, true>& queries,
-             int nprobe,
-             int k,
-             Tensor<float, 2, true>& outDistances,
-             Tensor<Index::idx_t, 2, true>& outIndices);
+    /// Find the approximate k nearest neigbors for `queries` against
+    /// our database
+    void query(
+            Tensor<float, 2, true>& queries,
+            int nprobe,
+            int k,
+            Tensor<float, 2, true>& outDistances,
+            Tensor<Index::idx_t, 2, true>& outIndices);
 
- protected:
-  /// Returns the number of bytes in which an IVF list containing numVecs
-  /// vectors is encoded on the device. Note that due to padding this is not the
-  /// same as the encoding size for a subset of vectors in an IVF list; this is
-  /// the size for an entire IVF list
-  size_t getGpuVectorsEncodingSize_(int numVecs) const override;
-  size_t getCpuVectorsEncodingSize_(int numVecs) const override;
+   protected:
+    /// Returns the number of bytes in which an IVF list containing numVecs
+    /// vectors is encoded on the device. Note that due to padding this is not
+    /// the same as the encoding size for a subset of vectors in an IVF list;
+    /// this is the size for an entire IVF list
+    size_t getGpuVectorsEncodingSize_(int numVecs) const override;
+    size_t getCpuVectorsEncodingSize_(int numVecs) const override;
 
-  /// Translate to our preferred GPU encoding
-  std::vector<uint8_t> translateCodesToGpu_(std::vector<uint8_t> codes,
-                                            size_t numVecs) const override;
+    /// Translate to our preferred GPU encoding
+    std::vector<uint8_t> translateCodesToGpu_(
+            std::vector<uint8_t> codes,
+            size_t numVecs) const override;
 
-  /// Translate from our preferred GPU encoding
-  std::vector<uint8_t> translateCodesFromGpu_(std::vector<uint8_t> codes,
-                                              size_t numVecs) const override;
+    /// Translate from our preferred GPU encoding
+    std::vector<uint8_t> translateCodesFromGpu_(
+            std::vector<uint8_t> codes,
+            size_t numVecs) const override;
 
-  /// Encode the vectors that we're adding and append to our IVF lists
-  void appendVectors_(Tensor<float, 2, true>& vecs,
-                      Tensor<Index::idx_t, 1, true>& indices,
-                      Tensor<int, 1, true>& uniqueLists,
-                      Tensor<int, 1, true>& vectorsByUniqueList,
-                      Tensor<int, 1, true>& uniqueListVectorStart,
-                      Tensor<int, 1, true>& uniqueListStartOffset,
-                      Tensor<int, 1, true>& listIds,
-                      Tensor<int, 1, true>& listOffset,
-                      cudaStream_t stream) override;
+    /// Encode the vectors that we're adding and append to our IVF lists
+    void appendVectors_(
+            Tensor<float, 2, true>& vecs,
+            Tensor<Index::idx_t, 1, true>& indices,
+            Tensor<int, 1, true>& uniqueLists,
+            Tensor<int, 1, true>& vectorsByUniqueList,
+            Tensor<int, 1, true>& uniqueListVectorStart,
+            Tensor<int, 1, true>& uniqueListStartOffset,
+            Tensor<int, 1, true>& listIds,
+            Tensor<int, 1, true>& listOffset,
+            cudaStream_t stream) override;
 
- protected:
-  /// Do we encode the residual from a coarse quantizer or not?
-  bool useResidual_;
+   protected:
+    /// Do we encode the residual from a coarse quantizer or not?
+    bool useResidual_;
 
-  /// Scalar quantizer for encoded vectors, if any
-  std::unique_ptr<GpuScalarQuantizer> scalarQ_;
+    /// Scalar quantizer for encoded vectors, if any
+    std::unique_ptr<GpuScalarQuantizer> scalarQ_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/IVFFlatScan.cuh

@@ -5,35 +5,37 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/MetricType.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
+#include <thrust/device_vector.h>
 #include <faiss/gpu/impl/GpuScalarQuantizer.cuh>
 #include <faiss/gpu/utils/Tensor.cuh>
-#include <thrust/device_vector.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
-void runIVFFlatScan(Tensor<float, 2, true>& queries,
-                    Tensor<int, 2, true>& listIds,
-                    thrust::device_vector<void*>& listData,
-                    thrust::device_vector<void*>& listIndices,
-                    IndicesOptions indicesOptions,
-                    thrust::device_vector<int>& listLengths,
-                    int maxListLength,
-                    int k,
-                    faiss::MetricType metric,
-                    bool useResidual,
-                    Tensor<float, 3, true>& residualBase,
-                    GpuScalarQuantizer* scalarQ,
-                    // output
-                    Tensor<float, 2, true>& outDistances,
-                    // output
-                    Tensor<Index::idx_t, 2, true>& outIndices,
-                    GpuResources* res);
+void runIVFFlatScan(
+        Tensor<float, 2, true>& queries,
+        Tensor<int, 2, true>& listIds,
+        thrust::device_vector<void*>& listData,
+        thrust::device_vector<void*>& listIndices,
+        IndicesOptions indicesOptions,
+        thrust::device_vector<int>& listLengths,
+        int maxListLength,
+        int k,
+        faiss::MetricType metric,
+        bool useResidual,
+        Tensor<float, 3, true>& residualBase,
+        GpuScalarQuantizer* scalarQ,
+        // output
+        Tensor<float, 2, true>& outDistances,
+        // output
+        Tensor<Index::idx_t, 2, true>& outIndices,
+        GpuResources* res);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/IVFInterleaved.cuh

@@ -10,193 +10,202 @@
 #include <faiss/MetricType.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
 #include <faiss/gpu/GpuResources.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/gpu/utils/StaticUtils.h>
+#include <thrust/device_vector.h>
 #include <faiss/gpu/impl/DistanceUtils.cuh>
 #include <faiss/gpu/impl/GpuScalarQuantizer.cuh>
 #include <faiss/gpu/utils/Comparators.cuh>
 #include <faiss/gpu/utils/ConversionOperators.cuh>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
-#include <faiss/gpu/utils/DeviceUtils.h>
 #include <faiss/gpu/utils/Float16.cuh>
 #include <faiss/gpu/utils/MathOperators.cuh>
 #include <faiss/gpu/utils/PtxUtils.cuh>
 #include <faiss/gpu/utils/Select.cuh>
-#include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/gpu/utils/WarpPackedBits.cuh>
-#include <thrust/device_vector.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// First pass kernel to perform scanning of IVF lists to produce top-k
 /// candidates
-template <typename Codec,
-          typename Metric,
-          int ThreadsPerBlock,
-          int NumWarpQ,
-          int NumThreadQ,
-          bool Residual>
-__global__ void
-ivfInterleavedScan(Tensor<float, 2, true> queries,
-                   Tensor<float, 3, true> residualBase,
-                   Tensor<int, 2, true> listIds,
-                   void** allListData,
-                   int* listLengths,
-                   Codec codec,
-                   Metric metric,
-                   int k,
-                   // [query][probe][k]
-                   Tensor<float, 3, true> distanceOut,
-                   Tensor<int, 3, true> indicesOut) {
-  extern __shared__ float smem[];
+template <
+        typename Codec,
+        typename Metric,
+        int ThreadsPerBlock,
+        int NumWarpQ,
+        int NumThreadQ,
+        bool Residual>
+__global__ void ivfInterleavedScan(
+        Tensor<float, 2, true> queries,
+        Tensor<float, 3, true> residualBase,
+        Tensor<int, 2, true> listIds,
+        void** allListData,
+        int* listLengths,
+        Codec codec,
+        Metric metric,
+        int k,
+        // [query][probe][k]
+        Tensor<float, 3, true> distanceOut,
+        Tensor<int, 3, true> indicesOut) {
+    extern __shared__ float smem[];
 
-  constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
+    constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
 
-  int queryId = blockIdx.y;
-  int probeId = blockIdx.x;
-  int listId = listIds[queryId][probeId];
+    int queryId = blockIdx.y;
+    int probeId = blockIdx.x;
+    int listId = listIds[queryId][probeId];
 
-  // Safety guard in case NaNs in input cause no list ID to be generated, or we
-  // have more nprobe than nlist
-  if (listId == -1) {
-    return;
-  }
+    // Safety guard in case NaNs in input cause no list ID to be generated, or
+    // we have more nprobe than nlist
+    if (listId == -1) {
+        return;
+    }
 
-  int dim = queries.getSize(1);
+    int dim = queries.getSize(1);
 
-  // FIXME: some issue with getLaneId() and CUDA 10.1 and P4 GPUs?
-  int laneId = threadIdx.x % kWarpSize;
-  int warpId = threadIdx.x / kWarpSize;
+    // FIXME: some issue with getLaneId() and CUDA 10.1 and P4 GPUs?
+    int laneId = threadIdx.x % kWarpSize;
+    int warpId = threadIdx.x / kWarpSize;
 
-  using EncodeT = typename Codec::EncodeT;
+    using EncodeT = typename Codec::EncodeT;
 
-  auto query = queries[queryId].data();
-  auto vecsBase = (EncodeT*) allListData[listId];
-  int numVecs = listLengths[listId];
-  auto residualBaseSlice = residualBase[queryId][probeId].data();
+    auto query = queries[queryId].data();
+    auto vecsBase = (EncodeT*)allListData[listId];
+    int numVecs = listLengths[listId];
+    auto residualBaseSlice = residualBase[queryId][probeId].data();
 
-  constexpr auto kInit = Metric::kDirection ? kFloatMin : kFloatMax;
+    constexpr auto kInit = Metric::kDirection ? kFloatMin : kFloatMax;
 
-  __shared__ float smemK[kNumWarps * NumWarpQ];
-  __shared__ int smemV[kNumWarps * NumWarpQ];
+    __shared__ float smemK[kNumWarps * NumWarpQ];
+    __shared__ int smemV[kNumWarps * NumWarpQ];
 
-  BlockSelect<float, int, Metric::kDirection, Comparator<float>,
-              NumWarpQ, NumThreadQ, ThreadsPerBlock>
-    heap(kInit, -1, smemK, smemV, k);
+    BlockSelect<
+            float,
+            int,
+            Metric::kDirection,
+            Comparator<float>,
+            NumWarpQ,
+            NumThreadQ,
+            ThreadsPerBlock>
+            heap(kInit, -1, smemK, smemV, k);
 
-  // The codec might be dependent upon data that we need to reference or store
-  // in shared memory
-  codec.initKernel(smem, dim);
-  __syncthreads();
+    // The codec might be dependent upon data that we need to reference or store
+    // in shared memory
+    codec.initKernel(smem, dim);
+    __syncthreads();
 
-  // How many vector blocks of 32 are in this list?
-  int numBlocks = utils::divUp(numVecs, 32);
+    // How many vector blocks of 32 are in this list?
+    int numBlocks = utils::divUp(numVecs, 32);
 
-  // Number of EncodeT words per each dimension of block of 32 vecs
-  constexpr int bytesPerVectorBlockDim =
-    Codec::kEncodeBits * 32 / 8;
-  constexpr int wordsPerVectorBlockDim =
-    bytesPerVectorBlockDim / sizeof(EncodeT);
-  int wordsPerVectorBlock = wordsPerVectorBlockDim * dim;
+    // Number of EncodeT words per each dimension of block of 32 vecs
+    constexpr int bytesPerVectorBlockDim = Codec::kEncodeBits * 32 / 8;
+    constexpr int wordsPerVectorBlockDim =
+            bytesPerVectorBlockDim / sizeof(EncodeT);
+    int wordsPerVectorBlock = wordsPerVectorBlockDim * dim;
 
-  int dimBlocks = utils::roundDown(dim, kWarpSize);
+    int dimBlocks = utils::roundDown(dim, kWarpSize);
 
-  for (int block = warpId; block < numBlocks; block += kNumWarps) {
-    // We're handling a new vector
-    Metric dist = metric.zero();
+    for (int block = warpId; block < numBlocks; block += kNumWarps) {
+        // We're handling a new vector
+        Metric dist = metric.zero();
 
-    // This is the vector a given lane/thread handles
-    int vec = block * kWarpSize + laneId;
-    bool valid = vec < numVecs;
+        // This is the vector a given lane/thread handles
+        int vec = block * kWarpSize + laneId;
+        bool valid = vec < numVecs;
 
-    // This is where this warp begins reading data
-    EncodeT* data = vecsBase + block * wordsPerVectorBlock;
+        // This is where this warp begins reading data
+        EncodeT* data = vecsBase + block * wordsPerVectorBlock;
 
-    // whole blocks
-    for (int dBase = 0; dBase < dimBlocks; dBase += kWarpSize) {
-      int loadDim = dBase + laneId;
-      float queryReg = query[loadDim];
-      float residualReg = Residual ? residualBaseSlice[loadDim] : 0;
+        // whole blocks
+        for (int dBase = 0; dBase < dimBlocks; dBase += kWarpSize) {
+            int loadDim = dBase + laneId;
+            float queryReg = query[loadDim];
+            float residualReg = Residual ? residualBaseSlice[loadDim] : 0;
 
-      constexpr int kUnroll = 4;
+            constexpr int kUnroll = 4;
 
 #pragma unroll
-      for (int i = 0; i < kWarpSize / kUnroll;
-           ++i, data += kUnroll * wordsPerVectorBlockDim) {
-        EncodeT encV[kUnroll];
- #pragma unroll
-        for (int j = 0; j < kUnroll; ++j) {
-          encV[j] = WarpPackedBits<EncodeT, Codec::kEncodeBits>::
-            read(laneId, data + j * wordsPerVectorBlockDim);
-        }
-
- #pragma unroll
-        for (int j = 0; j < kUnroll; ++j) {
-          encV[j] = WarpPackedBits<EncodeT, Codec::kEncodeBits>::
-            postRead(laneId, encV[j]);
-        }
-
-        float decV[kUnroll];
- #pragma unroll
-        for (int j = 0; j < kUnroll; ++j) {
-          int d = i * kUnroll + j;
-          decV[j] = codec.decodeNew(dBase + d, encV[j]);
-        }
-
-        if (Residual) {
- #pragma unroll
-          for (int j = 0; j < kUnroll; ++j) {
-            int d = i * kUnroll + j;
-            decV[j] += SHFL_SYNC(residualReg, d, kWarpSize);
-          }
-        }
+            for (int i = 0; i < kWarpSize / kUnroll;
+                 ++i, data += kUnroll * wordsPerVectorBlockDim) {
+                EncodeT encV[kUnroll];
+#pragma unroll
+                for (int j = 0; j < kUnroll; ++j) {
+                    encV[j] = WarpPackedBits<EncodeT, Codec::kEncodeBits>::read(
+                            laneId, data + j * wordsPerVectorBlockDim);
+                }
 
 #pragma unroll
-        for (int j = 0; j < kUnroll; ++j) {
-          int d = i * kUnroll + j;
-          float q = SHFL_SYNC(queryReg, d, kWarpSize);
-          dist.handle(q, decV[j]);
+                for (int j = 0; j < kUnroll; ++j) {
+                    encV[j] = WarpPackedBits<EncodeT, Codec::kEncodeBits>::
+                            postRead(laneId, encV[j]);
+                }
+
+                float decV[kUnroll];
+#pragma unroll
+                for (int j = 0; j < kUnroll; ++j) {
+                    int d = i * kUnroll + j;
+                    decV[j] = codec.decodeNew(dBase + d, encV[j]);
+                }
+
+                if (Residual) {
+#pragma unroll
+                    for (int j = 0; j < kUnroll; ++j) {
+                        int d = i * kUnroll + j;
+                        decV[j] += SHFL_SYNC(residualReg, d, kWarpSize);
+                    }
+                }
+
+#pragma unroll
+                for (int j = 0; j < kUnroll; ++j) {
+                    int d = i * kUnroll + j;
+                    float q = SHFL_SYNC(queryReg, d, kWarpSize);
+                    dist.handle(q, decV[j]);
+                }
+            }
         }
-      }
+
+        // remainder
+        int loadDim = dimBlocks + laneId;
+        bool loadDimInBounds = loadDim < dim;
+
+        float queryReg = loadDimInBounds ? query[loadDim] : 0;
+        float residualReg =
+                Residual && loadDimInBounds ? residualBaseSlice[loadDim] : 0;
+
+        for (int d = 0; d < dim - dimBlocks;
+             ++d, data += wordsPerVectorBlockDim) {
+            float q = SHFL_SYNC(queryReg, d, kWarpSize);
+
+            EncodeT enc = WarpPackedBits<EncodeT, Codec::kEncodeBits>::read(
+                    laneId, data);
+            enc = WarpPackedBits<EncodeT, Codec::kEncodeBits>::postRead(
+                    laneId, enc);
+            float dec = codec.decodeNew(dimBlocks + d, enc);
+            if (Residual) {
+                dec += SHFL_SYNC(residualReg, d, kWarpSize);
+            }
+
+            dist.handle(q, dec);
+        }
+
+        if (valid) {
+            heap.addThreadQ(dist.reduce(), vec);
+        }
+
+        heap.checkThreadQ();
     }
 
-    // remainder
-    int loadDim = dimBlocks + laneId;
-    bool loadDimInBounds = loadDim < dim;
+    heap.reduce();
 
-    float queryReg = loadDimInBounds ? query[loadDim] : 0;
-    float residualReg = Residual && loadDimInBounds ? residualBaseSlice[loadDim] : 0;
+    auto distanceOutBase = distanceOut[queryId][probeId].data();
+    auto indicesOutBase = indicesOut[queryId][probeId].data();
 
-    for (int d = 0; d < dim - dimBlocks; ++d, data += wordsPerVectorBlockDim) {
-      float q = SHFL_SYNC(queryReg, d, kWarpSize);
-
-      EncodeT enc =
-        WarpPackedBits<EncodeT, Codec::kEncodeBits>::read(laneId, data);
-      enc =
-        WarpPackedBits<EncodeT, Codec::kEncodeBits>::postRead(laneId, enc);
-      float dec = codec.decodeNew(dimBlocks + d, enc);
-      if (Residual) {
-        dec += SHFL_SYNC(residualReg, d, kWarpSize);
-      }
-
-      dist.handle(q, dec);
+    for (int i = threadIdx.x; i < k; i += blockDim.x) {
+        distanceOutBase[i] = smemK[i];
+        indicesOutBase[i] = smemV[i];
     }
-
-    if (valid) {
-      heap.addThreadQ(dist.reduce(), vec);
-    }
-
-    heap.checkThreadQ();
-  }
-
-  heap.reduce();
-
-  auto distanceOutBase = distanceOut[queryId][probeId].data();
-  auto indicesOutBase = indicesOut[queryId][probeId].data();
-
-  for (int i = threadIdx.x; i < k; i += blockDim.x) {
-    distanceOutBase[i] = smemK[i];
-    indicesOutBase[i] = smemV[i];
-  }
 }
 
 //
@@ -205,179 +214,230 @@ ivfInterleavedScan(Tensor<float, 2, true> queries,
 //
 
 #define IVFINT_RUN(CODEC_TYPE, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q) \
-  do {                                                                  \
-    dim3 grid(nprobe, nq);                                              \
-    if (useResidual) {                                                  \
-      ivfInterleavedScan<CODEC_TYPE, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q, true> \
-        <<<grid, THREADS, codec.getSmemSize(dim), stream>>>(            \
-        queries,                                                        \
-        residualBase,                                                   \
-        listIds,                                                        \
-        listData.data().get(),                                          \
-        listLengths.data().get(),                                       \
-        codec,                                                          \
-        metric,                                                         \
-        k,                                                              \
-        distanceTemp,                                                   \
-        indicesTemp);                                                   \
-     } else {                                                           \
-      ivfInterleavedScan<CODEC_TYPE, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q, false> \
-        <<<grid, THREADS, codec.getSmemSize(dim), stream>>>(            \
-          queries,                                                      \
-          residualBase,                                                 \
-          listIds,                                                      \
-          listData.data().get(),                                        \
-          listLengths.data().get(),                                     \
-          codec,                                                        \
-          metric,                                                       \
-          k,                                                            \
-          distanceTemp,                                                 \
-          indicesTemp);                                                 \
-    }                                                                   \
-                                                                        \
-    runIVFInterleavedScan2(distanceTemp,                                \
-                           indicesTemp,                                 \
-                           listIds,                                     \
-                           k,                                           \
-                           listIndices,                                 \
-                           indicesOptions,                              \
-                           METRIC_TYPE::kDirection,                     \
-                           outDistances,                                \
-                           outIndices,                                  \
-                           stream);                                     \
-  } while (0);
+    do {                                                                       \
+        dim3 grid(nprobe, nq);                                                 \
+        if (useResidual) {                                                     \
+            ivfInterleavedScan<                                                \
+                    CODEC_TYPE,                                                \
+                    METRIC_TYPE,                                               \
+                    THREADS,                                                   \
+                    NUM_WARP_Q,                                                \
+                    NUM_THREAD_Q,                                              \
+                    true><<<grid, THREADS, codec.getSmemSize(dim), stream>>>(  \
+                    queries,                                                   \
+                    residualBase,                                              \
+                    listIds,                                                   \
+                    listData.data().get(),                                     \
+                    listLengths.data().get(),                                  \
+                    codec,                                                     \
+                    metric,                                                    \
+                    k,                                                         \
+                    distanceTemp,                                              \
+                    indicesTemp);                                              \
+        } else {                                                               \
+            ivfInterleavedScan<                                                \
+                    CODEC_TYPE,                                                \
+                    METRIC_TYPE,                                               \
+                    THREADS,                                                   \
+                    NUM_WARP_Q,                                                \
+                    NUM_THREAD_Q,                                              \
+                    false><<<grid, THREADS, codec.getSmemSize(dim), stream>>>( \
+                    queries,                                                   \
+                    residualBase,                                              \
+                    listIds,                                                   \
+                    listData.data().get(),                                     \
+                    listLengths.data().get(),                                  \
+                    codec,                                                     \
+                    metric,                                                    \
+                    k,                                                         \
+                    distanceTemp,                                              \
+                    indicesTemp);                                              \
+        }                                                                      \
+                                                                               \
+        runIVFInterleavedScan2(                                                \
+                distanceTemp,                                                  \
+                indicesTemp,                                                   \
+                listIds,                                                       \
+                k,                                                             \
+                listIndices,                                                   \
+                indicesOptions,                                                \
+                METRIC_TYPE::kDirection,                                       \
+                outDistances,                                                  \
+                outIndices,                                                    \
+                stream);                                                       \
+    } while (0);
 
-#define IVFINT_CODECS(METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q)   \
-  do {                                                                  \
-    if (!scalarQ) {                                                     \
-      using CodecT = CodecFloat;                                        \
-      CodecT codec(dim * sizeof(float));                                \
-      IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-    } else {                                                            \
-      switch (scalarQ->qtype) {                                         \
-        case ScalarQuantizer::QuantizerType::QT_8bit:                   \
-        {                                                               \
-          using CodecT = Codec<ScalarQuantizer::QuantizerType::QT_8bit, 1>; \
-          CodecT                                                        \
-            codec(scalarQ->code_size,                                   \
-                  scalarQ->gpuTrained.data(),                           \
-                  scalarQ->gpuTrained.data() + dim);                    \
-          IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-        }                                                               \
-        break;                                                          \
-        case ScalarQuantizer::QuantizerType::QT_8bit_uniform:           \
-        {                                                               \
-          using CodecT = Codec<ScalarQuantizer::QuantizerType::QT_8bit_uniform, 1>; \
-          CodecT                                                        \
-            codec(scalarQ->code_size, scalarQ->trained[0], scalarQ->trained[1]); \
-          IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-        }                                                               \
-        break;                                                          \
-        case ScalarQuantizer::QuantizerType::QT_fp16:                   \
-        {                                                               \
-          using CodecT = Codec<ScalarQuantizer::QuantizerType::QT_fp16, 1>; \
-          CodecT                                                        \
-            codec(scalarQ->code_size);                                  \
-          IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-        }                                                               \
-        break;                                                          \
-        case ScalarQuantizer::QuantizerType::QT_8bit_direct:            \
-        {                                                               \
-          using CodecT = Codec<ScalarQuantizer::QuantizerType::QT_8bit_direct, 1>; \
-          Codec<ScalarQuantizer::QuantizerType::QT_8bit_direct, 1>      \
-            codec(scalarQ->code_size);                                  \
-          IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-        }                                                               \
-        break;                                                          \
-        case ScalarQuantizer::QuantizerType::QT_6bit:                   \
-        {                                                               \
-          using CodecT = Codec<ScalarQuantizer::QuantizerType::QT_6bit, 1>; \
-          Codec<ScalarQuantizer::QuantizerType::QT_6bit, 1>             \
-            codec(scalarQ->code_size,                                   \
-                  scalarQ->gpuTrained.data(),                           \
-                  scalarQ->gpuTrained.data() + dim);                    \
-          IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-        }                                                               \
-        break;                                                          \
-        case ScalarQuantizer::QuantizerType::QT_4bit:                   \
-        {                                                               \
-          using CodecT = Codec<ScalarQuantizer::QuantizerType::QT_4bit, 1>; \
-          Codec<ScalarQuantizer::QuantizerType::QT_4bit, 1>             \
-            codec(scalarQ->code_size,                                   \
-                  scalarQ->gpuTrained.data(),                           \
-                  scalarQ->gpuTrained.data() + dim);                    \
-          IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-        }                                                               \
-        break;                                                          \
-        case ScalarQuantizer::QuantizerType::QT_4bit_uniform:           \
-        {                                                               \
-          using CodecT = Codec<ScalarQuantizer::QuantizerType::QT_4bit_uniform, 1>; \
-          Codec<ScalarQuantizer::QuantizerType::QT_4bit_uniform, 1>     \
-            codec(scalarQ->code_size, scalarQ->trained[0], scalarQ->trained[1]); \
-          IVFINT_RUN(CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
-        }                                                               \
-        break;                                                          \
-        default:                                                        \
-          FAISS_ASSERT(false);                                          \
-      }                                                                 \
-    }                                                                   \
-  } while (0)
+#define IVFINT_CODECS(METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q)          \
+    do {                                                                       \
+        if (!scalarQ) {                                                        \
+            using CodecT = CodecFloat;                                         \
+            CodecT codec(dim * sizeof(float));                                 \
+            IVFINT_RUN(                                                        \
+                    CodecT, METRIC_TYPE, THREADS, NUM_WARP_Q, NUM_THREAD_Q);   \
+        } else {                                                               \
+            switch (scalarQ->qtype) {                                          \
+                case ScalarQuantizer::QuantizerType::QT_8bit: {                \
+                    using CodecT =                                             \
+                            Codec<ScalarQuantizer::QuantizerType::QT_8bit, 1>; \
+                    CodecT codec(                                              \
+                            scalarQ->code_size,                                \
+                            scalarQ->gpuTrained.data(),                        \
+                            scalarQ->gpuTrained.data() + dim);                 \
+                    IVFINT_RUN(                                                \
+                            CodecT,                                            \
+                            METRIC_TYPE,                                       \
+                            THREADS,                                           \
+                            NUM_WARP_Q,                                        \
+                            NUM_THREAD_Q);                                     \
+                } break;                                                       \
+                case ScalarQuantizer::QuantizerType::QT_8bit_uniform: {        \
+                    using CodecT = Codec<                                      \
+                            ScalarQuantizer::QuantizerType::QT_8bit_uniform,   \
+                            1>;                                                \
+                    CodecT codec(                                              \
+                            scalarQ->code_size,                                \
+                            scalarQ->trained[0],                               \
+                            scalarQ->trained[1]);                              \
+                    IVFINT_RUN(                                                \
+                            CodecT,                                            \
+                            METRIC_TYPE,                                       \
+                            THREADS,                                           \
+                            NUM_WARP_Q,                                        \
+                            NUM_THREAD_Q);                                     \
+                } break;                                                       \
+                case ScalarQuantizer::QuantizerType::QT_fp16: {                \
+                    using CodecT =                                             \
+                            Codec<ScalarQuantizer::QuantizerType::QT_fp16, 1>; \
+                    CodecT codec(scalarQ->code_size);                          \
+                    IVFINT_RUN(                                                \
+                            CodecT,                                            \
+                            METRIC_TYPE,                                       \
+                            THREADS,                                           \
+                            NUM_WARP_Q,                                        \
+                            NUM_THREAD_Q);                                     \
+                } break;                                                       \
+                case ScalarQuantizer::QuantizerType::QT_8bit_direct: {         \
+                    using CodecT = Codec<                                      \
+                            ScalarQuantizer::QuantizerType::QT_8bit_direct,    \
+                            1>;                                                \
+                    Codec<ScalarQuantizer::QuantizerType::QT_8bit_direct, 1>   \
+                            codec(scalarQ->code_size);                         \
+                    IVFINT_RUN(                                                \
+                            CodecT,                                            \
+                            METRIC_TYPE,                                       \
+                            THREADS,                                           \
+                            NUM_WARP_Q,                                        \
+                            NUM_THREAD_Q);                                     \
+                } break;                                                       \
+                case ScalarQuantizer::QuantizerType::QT_6bit: {                \
+                    using CodecT =                                             \
+                            Codec<ScalarQuantizer::QuantizerType::QT_6bit, 1>; \
+                    Codec<ScalarQuantizer::QuantizerType::QT_6bit, 1> codec(   \
+                            scalarQ->code_size,                                \
+                            scalarQ->gpuTrained.data(),                        \
+                            scalarQ->gpuTrained.data() + dim);                 \
+                    IVFINT_RUN(                                                \
+                            CodecT,                                            \
+                            METRIC_TYPE,                                       \
+                            THREADS,                                           \
+                            NUM_WARP_Q,                                        \
+                            NUM_THREAD_Q);                                     \
+                } break;                                                       \
+                case ScalarQuantizer::QuantizerType::QT_4bit: {                \
+                    using CodecT =                                             \
+                            Codec<ScalarQuantizer::QuantizerType::QT_4bit, 1>; \
+                    Codec<ScalarQuantizer::QuantizerType::QT_4bit, 1> codec(   \
+                            scalarQ->code_size,                                \
+                            scalarQ->gpuTrained.data(),                        \
+                            scalarQ->gpuTrained.data() + dim);                 \
+                    IVFINT_RUN(                                                \
+                            CodecT,                                            \
+                            METRIC_TYPE,                                       \
+                            THREADS,                                           \
+                            NUM_WARP_Q,                                        \
+                            NUM_THREAD_Q);                                     \
+                } break;                                                       \
+                case ScalarQuantizer::QuantizerType::QT_4bit_uniform: {        \
+                    using CodecT = Codec<                                      \
+                            ScalarQuantizer::QuantizerType::QT_4bit_uniform,   \
+                            1>;                                                \
+                    Codec<ScalarQuantizer::QuantizerType::QT_4bit_uniform, 1>  \
+                            codec(scalarQ->code_size,                          \
+                                  scalarQ->trained[0],                         \
+                                  scalarQ->trained[1]);                        \
+                    IVFINT_RUN(                                                \
+                            CodecT,                                            \
+                            METRIC_TYPE,                                       \
+                            THREADS,                                           \
+                            NUM_WARP_Q,                                        \
+                            NUM_THREAD_Q);                                     \
+                } break;                                                       \
+                default:                                                       \
+                    FAISS_ASSERT(false);                                       \
+            }                                                                  \
+        }                                                                      \
+    } while (0)
 
-#define IVFINT_METRICS(THREADS, NUM_WARP_Q, NUM_THREAD_Q)               \
-  do {                                                                  \
-    auto stream = res->getDefaultStreamCurrentDevice();                 \
-    auto nq = queries.getSize(0);                                       \
-    auto dim = queries.getSize(1);                                      \
-    auto nprobe = listIds.getSize(1);                                   \
-                                                                        \
-    DeviceTensor<float, 3, true> distanceTemp(                          \
-      res, makeTempAlloc(AllocType::Other, stream),                     \
-      {queries.getSize(0), listIds.getSize(1), k});                     \
-    DeviceTensor<int, 3, true> indicesTemp(                             \
-      res, makeTempAlloc(AllocType::Other, stream),                     \
-      {queries.getSize(0), listIds.getSize(1), k});                     \
-                                                                        \
-    if (metric == MetricType::METRIC_L2) {                              \
-      L2Distance metric;                                                \
-      IVFINT_CODECS(L2Distance, THREADS, NUM_WARP_Q, NUM_THREAD_Q);     \
-    } else if (metric == MetricType::METRIC_INNER_PRODUCT) {            \
-      IPDistance metric;                                                \
-      IVFINT_CODECS(IPDistance, THREADS, NUM_WARP_Q, NUM_THREAD_Q);     \
-    } else {                                                            \
-      FAISS_ASSERT(false);                                              \
-    }                                                                   \
-  } while (0)
+#define IVFINT_METRICS(THREADS, NUM_WARP_Q, NUM_THREAD_Q)                 \
+    do {                                                                  \
+        auto stream = res->getDefaultStreamCurrentDevice();               \
+        auto nq = queries.getSize(0);                                     \
+        auto dim = queries.getSize(1);                                    \
+        auto nprobe = listIds.getSize(1);                                 \
+                                                                          \
+        DeviceTensor<float, 3, true> distanceTemp(                        \
+                res,                                                      \
+                makeTempAlloc(AllocType::Other, stream),                  \
+                {queries.getSize(0), listIds.getSize(1), k});             \
+        DeviceTensor<int, 3, true> indicesTemp(                           \
+                res,                                                      \
+                makeTempAlloc(AllocType::Other, stream),                  \
+                {queries.getSize(0), listIds.getSize(1), k});             \
+                                                                          \
+        if (metric == MetricType::METRIC_L2) {                            \
+            L2Distance metric;                                            \
+            IVFINT_CODECS(L2Distance, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
+        } else if (metric == MetricType::METRIC_INNER_PRODUCT) {          \
+            IPDistance metric;                                            \
+            IVFINT_CODECS(IPDistance, THREADS, NUM_WARP_Q, NUM_THREAD_Q); \
+        } else {                                                          \
+            FAISS_ASSERT(false);                                          \
+        }                                                                 \
+    } while (0)
 
 // Top-level IVF scan function for the interleaved by 32 layout
 // with all implementations
-void runIVFInterleavedScan(Tensor<float, 2, true>& queries,
-                           Tensor<int, 2, true>& listIds,
-                           thrust::device_vector<void*>& listData,
-                           thrust::device_vector<void*>& listIndices,
-                           IndicesOptions indicesOptions,
-                           thrust::device_vector<int>& listLengths,
-                           int k,
-                           faiss::MetricType metric,
-                           bool useResidual,
-                           Tensor<float, 3, true>& residualBase,
-                           GpuScalarQuantizer* scalarQ,
-                           // output
-                           Tensor<float, 2, true>& outDistances,
-                           // output
-                           Tensor<Index::idx_t, 2, true>& outIndices,
-                           GpuResources* res);
+void runIVFInterleavedScan(
+        Tensor<float, 2, true>& queries,
+        Tensor<int, 2, true>& listIds,
+        thrust::device_vector<void*>& listData,
+        thrust::device_vector<void*>& listIndices,
+        IndicesOptions indicesOptions,
+        thrust::device_vector<int>& listLengths,
+        int k,
+        faiss::MetricType metric,
+        bool useResidual,
+        Tensor<float, 3, true>& residualBase,
+        GpuScalarQuantizer* scalarQ,
+        // output
+        Tensor<float, 2, true>& outDistances,
+        // output
+        Tensor<Index::idx_t, 2, true>& outIndices,
+        GpuResources* res);
 
 // Second pass of IVF list scanning to perform final k-selection and look up the
 // user indices
-void runIVFInterleavedScan2(Tensor<float, 3, true>& distanceIn,
-                            Tensor<int, 3, true>& indicesIn,
-                            Tensor<int, 2, true>& listIds,
-                            int k,
-                            thrust::device_vector<void*>& listIndices,
-                            IndicesOptions indicesOptions,
-                            bool dir,
-                            Tensor<float, 2, true>& distanceOut,
-                            Tensor<Index::idx_t, 2, true>& indicesOut,
-                            cudaStream_t stream);
+void runIVFInterleavedScan2(
+        Tensor<float, 3, true>& distanceIn,
+        Tensor<int, 3, true>& indicesIn,
+        Tensor<int, 2, true>& listIds,
+        int k,
+        thrust::device_vector<void*>& listIndices,
+        IndicesOptions indicesOptions,
+        bool dir,
+        Tensor<float, 2, true>& distanceOut,
+        Tensor<Index::idx_t, 2, true>& indicesOut,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/IVFPQ.cuh

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/Index.h>
@@ -13,151 +12,160 @@
 #include <faiss/gpu/impl/IVFBase.cuh>
 #include <faiss/gpu/utils/Float16.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// Implementing class for IVFPQ on the GPU
 class IVFPQ : public IVFBase {
- public:
-  IVFPQ(GpuResources* resources,
-        faiss::MetricType metric,
-        float metricArg,
-        /// We do not own this reference
-        FlatIndex* quantizer,
-        int numSubQuantizers,
-        int bitsPerSubQuantizer,
-        bool useFloat16LookupTables,
-        bool useMMCodeDistance,
-        bool interleavedLayout,
-        float* pqCentroidData,
-        IndicesOptions indicesOptions,
-        MemorySpace space);
+   public:
+    IVFPQ(GpuResources* resources,
+          faiss::MetricType metric,
+          float metricArg,
+          /// We do not own this reference
+          FlatIndex* quantizer,
+          int numSubQuantizers,
+          int bitsPerSubQuantizer,
+          bool useFloat16LookupTables,
+          bool useMMCodeDistance,
+          bool interleavedLayout,
+          float* pqCentroidData,
+          IndicesOptions indicesOptions,
+          MemorySpace space);
 
-  /// Returns true if we support PQ in this size
-  static bool isSupportedPQCodeLength(int size);
+    /// Returns true if we support PQ in this size
+    static bool isSupportedPQCodeLength(int size);
 
-  ~IVFPQ() override;
+    ~IVFPQ() override;
 
-  /// Enable or disable pre-computed codes
-  void setPrecomputedCodes(bool enable);
+    /// Enable or disable pre-computed codes
+    void setPrecomputedCodes(bool enable);
 
-  /// Find the approximate k nearest neigbors for `queries` against
-  /// our database
-  void query(Tensor<float, 2, true>& queries,
-             int nprobe,
-             int k,
-             Tensor<float, 2, true>& outDistances,
-             Tensor<Index::idx_t, 2, true>& outIndices);
+    /// Find the approximate k nearest neigbors for `queries` against
+    /// our database
+    void query(
+            Tensor<float, 2, true>& queries,
+            int nprobe,
+            int k,
+            Tensor<float, 2, true>& outDistances,
+            Tensor<Index::idx_t, 2, true>& outIndices);
 
-  /// Returns our set of sub-quantizers of the form
-  /// (sub q)(code id)(sub dim)
-  Tensor<float, 3, true> getPQCentroids();
+    /// Returns our set of sub-quantizers of the form
+    /// (sub q)(code id)(sub dim)
+    Tensor<float, 3, true> getPQCentroids();
 
- protected:
-  /// Returns the encoding size for a PQ-encoded IVF list
-  size_t getGpuVectorsEncodingSize_(int numVecs) const override;
-  size_t getCpuVectorsEncodingSize_(int numVecs) const override;
+   protected:
+    /// Returns the encoding size for a PQ-encoded IVF list
+    size_t getGpuVectorsEncodingSize_(int numVecs) const override;
+    size_t getCpuVectorsEncodingSize_(int numVecs) const override;
 
-  /// Translate to our preferred GPU encoding
-  std::vector<uint8_t> translateCodesToGpu_(std::vector<uint8_t> codes,
-                                            size_t numVecs) const override;
+    /// Translate to our preferred GPU encoding
+    std::vector<uint8_t> translateCodesToGpu_(
+            std::vector<uint8_t> codes,
+            size_t numVecs) const override;
 
-  /// Translate from our preferred GPU encoding
-  std::vector<uint8_t> translateCodesFromGpu_(std::vector<uint8_t> codes,
-                                              size_t numVecs) const override;
+    /// Translate from our preferred GPU encoding
+    std::vector<uint8_t> translateCodesFromGpu_(
+            std::vector<uint8_t> codes,
+            size_t numVecs) const override;
 
-  /// Encode the vectors that we're adding and append to our IVF lists
-  void appendVectors_(Tensor<float, 2, true>& vecs,
-                      Tensor<Index::idx_t, 1, true>& indices,
-                      Tensor<int, 1, true>& uniqueLists,
-                      Tensor<int, 1, true>& vectorsByUniqueList,
-                      Tensor<int, 1, true>& uniqueListVectorStart,
-                      Tensor<int, 1, true>& uniqueListStartOffset,
-                      Tensor<int, 1, true>& listIds,
-                      Tensor<int, 1, true>& listOffset,
-                      cudaStream_t stream) override;
+    /// Encode the vectors that we're adding and append to our IVF lists
+    void appendVectors_(
+            Tensor<float, 2, true>& vecs,
+            Tensor<Index::idx_t, 1, true>& indices,
+            Tensor<int, 1, true>& uniqueLists,
+            Tensor<int, 1, true>& vectorsByUniqueList,
+            Tensor<int, 1, true>& uniqueListVectorStart,
+            Tensor<int, 1, true>& uniqueListStartOffset,
+            Tensor<int, 1, true>& listIds,
+            Tensor<int, 1, true>& listOffset,
+            cudaStream_t stream) override;
 
-  /// Sets the current product quantizer centroids; the data can be
-  /// resident on either the host or the device. It will be transposed
-  /// into our preferred data layout
-  /// Data must be a row-major, 3-d array of size
-  /// (numSubQuantizers, numSubQuantizerCodes, dim / numSubQuantizers)
-  void setPQCentroids_(float* data);
+    /// Sets the current product quantizer centroids; the data can be
+    /// resident on either the host or the device. It will be transposed
+    /// into our preferred data layout
+    /// Data must be a row-major, 3-d array of size
+    /// (numSubQuantizers, numSubQuantizerCodes, dim / numSubQuantizers)
+    void setPQCentroids_(float* data);
 
-  /// Calculate precomputed residual distance information
-  void precomputeCodes_();
+    /// Calculate precomputed residual distance information
+    void precomputeCodes_();
 
-  /// Calculate precomputed residual distance information (for different coarse
-  /// centroid type)
-  template <typename CentroidT>
-  void precomputeCodesT_();
+    /// Calculate precomputed residual distance information (for different
+    /// coarse centroid type)
+    template <typename CentroidT>
+    void precomputeCodesT_();
 
-  /// Runs kernels for scanning inverted lists with precomputed codes
-  void runPQPrecomputedCodes_(Tensor<float, 2, true>& queries,
-                              DeviceTensor<float, 2, true>& coarseDistances,
-                              DeviceTensor<int, 2, true>& coarseIndices,
-                              int k,
-                              Tensor<float, 2, true>& outDistances,
-                              Tensor<Index::idx_t, 2, true>& outIndices);
+    /// Runs kernels for scanning inverted lists with precomputed codes
+    void runPQPrecomputedCodes_(
+            Tensor<float, 2, true>& queries,
+            DeviceTensor<float, 2, true>& coarseDistances,
+            DeviceTensor<int, 2, true>& coarseIndices,
+            int k,
+            Tensor<float, 2, true>& outDistances,
+            Tensor<Index::idx_t, 2, true>& outIndices);
 
-  /// Runs kernels for scanning inverted lists without precomputed codes
-  void runPQNoPrecomputedCodes_(Tensor<float, 2, true>& queries,
-                                DeviceTensor<float, 2, true>& coarseDistances,
-                                DeviceTensor<int, 2, true>& coarseIndices,
-                                int k,
-                                Tensor<float, 2, true>& outDistances,
-                                Tensor<Index::idx_t, 2, true>& outIndices);
+    /// Runs kernels for scanning inverted lists without precomputed codes
+    void runPQNoPrecomputedCodes_(
+            Tensor<float, 2, true>& queries,
+            DeviceTensor<float, 2, true>& coarseDistances,
+            DeviceTensor<int, 2, true>& coarseIndices,
+            int k,
+            Tensor<float, 2, true>& outDistances,
+            Tensor<Index::idx_t, 2, true>& outIndices);
 
-  /// Runs kernels for scanning inverted lists without precomputed codes (for
-  /// different coarse centroid type)
-  template <typename CentroidT>
-  void runPQNoPrecomputedCodesT_(Tensor<float, 2, true>& queries,
-                                 DeviceTensor<float, 2, true>& coarseDistances,
-                                 DeviceTensor<int, 2, true>& coarseIndices,
-                                 int k,
-                                 Tensor<float, 2, true>& outDistances,
-                                 Tensor<Index::idx_t, 2, true>& outIndices);
+    /// Runs kernels for scanning inverted lists without precomputed codes (for
+    /// different coarse centroid type)
+    template <typename CentroidT>
+    void runPQNoPrecomputedCodesT_(
+            Tensor<float, 2, true>& queries,
+            DeviceTensor<float, 2, true>& coarseDistances,
+            DeviceTensor<int, 2, true>& coarseIndices,
+            int k,
+            Tensor<float, 2, true>& outDistances,
+            Tensor<Index::idx_t, 2, true>& outIndices);
 
- private:
-  /// Number of sub-quantizers per vector
-  const int numSubQuantizers_;
+   private:
+    /// Number of sub-quantizers per vector
+    const int numSubQuantizers_;
 
-  /// Number of bits per sub-quantizer
-  const int bitsPerSubQuantizer_;
+    /// Number of bits per sub-quantizer
+    const int bitsPerSubQuantizer_;
 
-  /// Number of per sub-quantizer codes (2^bits)
-  const int numSubQuantizerCodes_;
+    /// Number of per sub-quantizer codes (2^bits)
+    const int numSubQuantizerCodes_;
 
-  /// Number of dimensions per each sub-quantizer
-  const int dimPerSubQuantizer_;
+    /// Number of dimensions per each sub-quantizer
+    const int dimPerSubQuantizer_;
 
-  /// Do we maintain precomputed terms and lookup tables in float16
-  /// form?
-  const bool useFloat16LookupTables_;
+    /// Do we maintain precomputed terms and lookup tables in float16
+    /// form?
+    const bool useFloat16LookupTables_;
 
-  /// For usage without precomputed codes, do we force usage of the
-  /// general-purpose MM code distance computation? This is for testing
-  /// purposes.
-  const bool useMMCodeDistance_;
+    /// For usage without precomputed codes, do we force usage of the
+    /// general-purpose MM code distance computation? This is for testing
+    /// purposes.
+    const bool useMMCodeDistance_;
 
-  /// On the GPU, we prefer different PQ centroid data layouts for
-  /// different purposes.
-  ///
-  /// (sub q)(sub dim)(code id)
-  DeviceTensor<float, 3, true> pqCentroidsInnermostCode_;
+    /// On the GPU, we prefer different PQ centroid data layouts for
+    /// different purposes.
+    ///
+    /// (sub q)(sub dim)(code id)
+    DeviceTensor<float, 3, true> pqCentroidsInnermostCode_;
 
-  /// (sub q)(code id)(sub dim)
-  DeviceTensor<float, 3, true> pqCentroidsMiddleCode_;
+    /// (sub q)(code id)(sub dim)
+    DeviceTensor<float, 3, true> pqCentroidsMiddleCode_;
 
-  /// Are precomputed codes enabled? (additional factoring and
-  /// precomputation of the residual distance, to reduce query-time work)
-  bool precomputedCodes_;
+    /// Are precomputed codes enabled? (additional factoring and
+    /// precomputation of the residual distance, to reduce query-time work)
+    bool precomputedCodes_;
 
-  /// Precomputed term 2 in float form
-  /// (centroid id)(sub q)(code id)
-  DeviceTensor<float, 3, true> precomputedCode_;
+    /// Precomputed term 2 in float form
+    /// (centroid id)(sub q)(code id)
+    DeviceTensor<float, 3, true> precomputedCode_;
 
-  /// Precomputed term 2 in half form
-  DeviceTensor<half, 3, true> precomputedCodeHalf_;
+    /// Precomputed term 2 in half form
+    DeviceTensor<half, 3, true> precomputedCodeHalf_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/IVFUtils.cuh

@@ -5,51 +5,55 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/Index.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
-#include <faiss/gpu/utils/Tensor.cuh>
 #include <thrust/device_vector.h>
+#include <faiss/gpu/utils/Tensor.cuh>
 
 // A collection of utility functions for IVFPQ and IVFFlat, for
 // post-processing and k-selecting the results
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
 /// Function for multi-pass scanning that collects the length of
 /// intermediate results for all (query, probe) pair
-void runCalcListOffsets(GpuResources* res,
-                        Tensor<int, 2, true>& topQueryToCentroid,
-                        thrust::device_vector<int>& listLengths,
-                        Tensor<int, 2, true>& prefixSumOffsets,
-                        Tensor<char, 1, true>& thrustMem,
-                        cudaStream_t stream);
+void runCalcListOffsets(
+        GpuResources* res,
+        Tensor<int, 2, true>& topQueryToCentroid,
+        thrust::device_vector<int>& listLengths,
+        Tensor<int, 2, true>& prefixSumOffsets,
+        Tensor<char, 1, true>& thrustMem,
+        cudaStream_t stream);
 
 /// Performs a first pass of k-selection on the results
-void runPass1SelectLists(Tensor<int, 2, true>& prefixSumOffsets,
-                         Tensor<float, 1, true>& distance,
-                         int nprobe,
-                         int k,
-                         bool chooseLargest,
-                         Tensor<float, 3, true>& heapDistances,
-                         Tensor<int, 3, true>& heapIndices,
-                         cudaStream_t stream);
+void runPass1SelectLists(
+        Tensor<int, 2, true>& prefixSumOffsets,
+        Tensor<float, 1, true>& distance,
+        int nprobe,
+        int k,
+        bool chooseLargest,
+        Tensor<float, 3, true>& heapDistances,
+        Tensor<int, 3, true>& heapIndices,
+        cudaStream_t stream);
 
 /// Performs a final pass of k-selection on the results, producing the
 /// final indices
-void runPass2SelectLists(Tensor<float, 2, true>& heapDistances,
-                         Tensor<int, 2, true>& heapIndices,
-                         thrust::device_vector<void*>& listIndices,
-                         IndicesOptions indicesOptions,
-                         Tensor<int, 2, true>& prefixSumOffsets,
-                         Tensor<int, 2, true>& topQueryToCentroid,
-                         int k,
-                         bool chooseLargest,
-                         Tensor<float, 2, true>& outDistances,
-                         Tensor<Index::idx_t, 2, true>& outIndices,
-                         cudaStream_t stream);
+void runPass2SelectLists(
+        Tensor<float, 2, true>& heapDistances,
+        Tensor<int, 2, true>& heapIndices,
+        thrust::device_vector<void*>& listIndices,
+        IndicesOptions indicesOptions,
+        Tensor<int, 2, true>& prefixSumOffsets,
+        Tensor<int, 2, true>& topQueryToCentroid,
+        int k,
+        bool chooseLargest,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<Index::idx_t, 2, true>& outIndices,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/L2Norm.cuh

@@ -5,23 +5,26 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
-void runL2Norm(Tensor<float, 2, true>& input,
-               bool inputRowMajor,
-               Tensor<float, 1, true>& output,
-               bool normSquared,
-               cudaStream_t stream);
+void runL2Norm(
+        Tensor<float, 2, true>& input,
+        bool inputRowMajor,
+        Tensor<float, 1, true>& output,
+        bool normSquared,
+        cudaStream_t stream);
 
-void runL2Norm(Tensor<half, 2, true>& input,
-               bool inputRowMajor,
-               Tensor<float, 1, true>& output,
-               bool normSquared,
-               cudaStream_t stream);
+void runL2Norm(
+        Tensor<half, 2, true>& input,
+        bool inputRowMajor,
+        Tensor<float, 1, true>& output,
+        bool normSquared,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/L2Select.cuh

@@ -5,18 +5,20 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
-void runL2SelectMin(Tensor<float, 2, true>& productDistances,
-                    Tensor<float, 1, true>& centroidDistances,
-                    Tensor<float, 2, true>& outDistances,
-                    Tensor<int, 2, true>& outIndices,
-                    int k,
-                    cudaStream_t stream);
+void runL2SelectMin(
+        Tensor<float, 2, true>& productDistances,
+        Tensor<float, 1, true>& centroidDistances,
+        Tensor<float, 2, true>& outDistances,
+        Tensor<int, 2, true>& outIndices,
+        int k,
+        cudaStream_t stream);
 
-} } // namespace
+}
+} // namespace faiss

faiss/gpu/impl/PQCodeDistances.cuh

@@ -5,14 +5,14 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/gpu/utils/Tensor.cuh>
-#include <faiss/gpu/utils/NoTypeTensor.cuh>
 #include <cublas_v2.h>
+#include <faiss/gpu/utils/NoTypeTensor.cuh>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class DeviceMemory;
 
@@ -21,17 +21,19 @@ class DeviceMemory;
 /// each sub-code vector, for the given list of query results in
 /// coarseIndices
 template <typename CentroidT>
-void runPQCodeDistances(GpuResources* res,
-                        Tensor<float, 3, true>& pqCentroids,
-                        Tensor<float, 2, true>& queries,
-                        Tensor<CentroidT, 2, true>& coarseCentroids,
-                        Tensor<float, 2, true>& coarseDistances,
-                        Tensor<int, 2, true>& coarseIndices,
-                        NoTypeTensor<4, true>& outCodeDistances,
-                        bool useMMImplementation,
-                        bool l2Distance,
-                        bool useFloat16Lookup);
+void runPQCodeDistances(
+        GpuResources* res,
+        Tensor<float, 3, true>& pqCentroids,
+        Tensor<float, 2, true>& queries,
+        Tensor<CentroidT, 2, true>& coarseCentroids,
+        Tensor<float, 2, true>& coarseDistances,
+        Tensor<int, 2, true>& coarseIndices,
+        NoTypeTensor<4, true>& outCodeDistances,
+        bool useMMImplementation,
+        bool l2Distance,
+        bool useFloat16Lookup);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss
 
 #include <faiss/gpu/impl/PQCodeDistances-inl.cuh>

faiss/gpu/impl/PQCodeLoad.cuh

@@ -5,12 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/PtxUtils.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 #if __CUDA_ARCH__ >= 350
 // Use the CC 3.5+ read-only texture cache (nc)
@@ -31,327 +31,352 @@ namespace faiss { namespace gpu {
 
 // Type-specific wrappers around the PTX bfe.* instruction, for
 // quantization code extraction
-inline __device__ unsigned int getByte(unsigned char v,
-                                       int pos,
-                                       int width) {
-  return v;
+inline __device__ unsigned int getByte(unsigned char v, int pos, int width) {
+    return v;
 }
 
-inline __device__ unsigned int getByte(unsigned short v,
-                                       int pos,
-                                       int width) {
-  return getBitfield((unsigned int) v, pos, width);
+inline __device__ unsigned int getByte(unsigned short v, int pos, int width) {
+    return getBitfield((unsigned int)v, pos, width);
 }
 
-inline __device__ unsigned int getByte(unsigned int v,
-                                       int pos,
-                                       int width) {
-  return getBitfield(v, pos, width);
+inline __device__ unsigned int getByte(unsigned int v, int pos, int width) {
+    return getBitfield(v, pos, width);
 }
 
-inline __device__ unsigned int getByte(uint64_t v,
-                                       int pos,
-                                       int width) {
-  return getBitfield(v, pos, width);
+inline __device__ unsigned int getByte(uint64_t v, int pos, int width) {
+    return getBitfield(v, pos, width);
 }
 
 template <int NumSubQuantizers>
 struct LoadCode32 {};
 
-template<>
+template <>
 struct LoadCode32<1> {
-  static inline __device__ void load(unsigned int code32[1],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 1;
-    asm("ld.global.cs.u8 {%0}, [%1];" :
-        "=r"(code32[0]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[1],
+            uint8_t* p,
+            int offset) {
+        p += offset * 1;
+        asm("ld.global.cs.u8 {%0}, [%1];" : "=r"(code32[0]) : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<2> {
-  static inline __device__ void load(unsigned int code32[1],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 2;
-    asm("ld.global.cs.u16 {%0}, [%1];" :
-        "=r"(code32[0]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[1],
+            uint8_t* p,
+            int offset) {
+        p += offset * 2;
+        asm("ld.global.cs.u16 {%0}, [%1];" : "=r"(code32[0]) : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<3> {
-  static inline __device__ void load(unsigned int code32[1],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 3;
-    unsigned int a;
-    unsigned int b;
-    unsigned int c;
+    static inline __device__ void load(
+            unsigned int code32[1],
+            uint8_t* p,
+            int offset) {
+        p += offset * 3;
+        unsigned int a;
+        unsigned int b;
+        unsigned int c;
 
-    // FIXME: this is a non-coalesced, unaligned, non-vectorized load
-    // unfortunately need to reorganize memory layout by warp
-    asm("ld.global.cs.u8 {%0}, [%1 + 0];" :
-        "=r"(a) : "l"(p));
-    asm("ld.global.cs.u8 {%0}, [%1 + 1];" :
-        "=r"(b) : "l"(p));
-    asm("ld.global.cs.u8 {%0}, [%1 + 2];" :
-        "=r"(c) : "l"(p));
+        // FIXME: this is a non-coalesced, unaligned, non-vectorized load
+        // unfortunately need to reorganize memory layout by warp
+        asm("ld.global.cs.u8 {%0}, [%1 + 0];" : "=r"(a) : "l"(p));
+        asm("ld.global.cs.u8 {%0}, [%1 + 1];" : "=r"(b) : "l"(p));
+        asm("ld.global.cs.u8 {%0}, [%1 + 2];" : "=r"(c) : "l"(p));
 
-    // FIXME: this is also slow, since we have to recover the
-    // individual bytes loaded
-    code32[0] = (c << 16) | (b << 8) | a;
-  }
+        // FIXME: this is also slow, since we have to recover the
+        // individual bytes loaded
+        code32[0] = (c << 16) | (b << 8) | a;
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<4> {
-  static inline __device__ void load(unsigned int code32[1],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 4;
-    asm("ld.global.cs.u32 {%0}, [%1];" :
-        "=r"(code32[0]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[1],
+            uint8_t* p,
+            int offset) {
+        p += offset * 4;
+        asm("ld.global.cs.u32 {%0}, [%1];" : "=r"(code32[0]) : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<8> {
-  static inline __device__ void load(unsigned int code32[2],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 8;
-    asm("ld.global.cs.v2.u32 {%0, %1}, [%2];" :
-        "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[2],
+            uint8_t* p,
+            int offset) {
+        p += offset * 8;
+        asm("ld.global.cs.v2.u32 {%0, %1}, [%2];"
+            : "=r"(code32[0]), "=r"(code32[1])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<12> {
-  static inline __device__ void load(unsigned int code32[3],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 12;
-    // FIXME: this is a non-coalesced, unaligned, non-vectorized load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V1 " {%0}, [%1 + 0];" :
-        "=r"(code32[0]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 4];" :
-        "=r"(code32[1]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 8];" :
-        "=r"(code32[2]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[3],
+            uint8_t* p,
+            int offset) {
+        p += offset * 12;
+        // FIXME: this is a non-coalesced, unaligned, non-vectorized load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V1 " {%0}, [%1 + 0];" : "=r"(code32[0]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 4];" : "=r"(code32[1]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 8];" : "=r"(code32[2]) : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<16> {
-  static inline __device__ void load(unsigned int code32[4],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 16;
-    asm("ld.global.cs.v4.u32 {%0, %1, %2, %3}, [%4];" :
-        "=r"(code32[0]), "=r"(code32[1]),
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[4],
+            uint8_t* p,
+            int offset) {
+        p += offset * 16;
+        asm("ld.global.cs.v4.u32 {%0, %1, %2, %3}, [%4];"
+            : "=r"(code32[0]), "=r"(code32[1]), "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<20> {
-  static inline __device__ void load(unsigned int code32[5],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 20;
-    // FIXME: this is a non-coalesced, unaligned, non-vectorized load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V1 " {%0}, [%1 + 0];" :
-        "=r"(code32[0]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 4];" :
-        "=r"(code32[1]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 8];" :
-        "=r"(code32[2]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 12];" :
-        "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 16];" :
-        "=r"(code32[4]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[5],
+            uint8_t* p,
+            int offset) {
+        p += offset * 20;
+        // FIXME: this is a non-coalesced, unaligned, non-vectorized load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V1 " {%0}, [%1 + 0];" : "=r"(code32[0]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 4];" : "=r"(code32[1]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 8];" : "=r"(code32[2]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 12];" : "=r"(code32[3]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 16];" : "=r"(code32[4]) : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<24> {
-  static inline __device__ void load(unsigned int code32[6],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 24;
-    // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 0];" :
-        "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 8];" :
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 16];" :
-        "=r"(code32[4]), "=r"(code32[5]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[6],
+            uint8_t* p,
+            int offset) {
+        p += offset * 24;
+        // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 0];"
+            : "=r"(code32[0]), "=r"(code32[1])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 8];"
+            : "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 16];"
+            : "=r"(code32[4]), "=r"(code32[5])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<28> {
-  static inline __device__ void load(unsigned int code32[7],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 28;
-    // FIXME: this is a non-coalesced, unaligned, non-vectorized load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V1 " {%0}, [%1 + 0];" :
-        "=r"(code32[0]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 4];" :
-        "=r"(code32[1]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 8];" :
-        "=r"(code32[2]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 12];" :
-        "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 16];" :
-        "=r"(code32[4]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 20];" :
-        "=r"(code32[5]) : "l"(p));
-    asm(LD_NC_V1 " {%0}, [%1 + 24];" :
-        "=r"(code32[6]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[7],
+            uint8_t* p,
+            int offset) {
+        p += offset * 28;
+        // FIXME: this is a non-coalesced, unaligned, non-vectorized load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V1 " {%0}, [%1 + 0];" : "=r"(code32[0]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 4];" : "=r"(code32[1]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 8];" : "=r"(code32[2]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 12];" : "=r"(code32[3]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 16];" : "=r"(code32[4]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 20];" : "=r"(code32[5]) : "l"(p));
+        asm(LD_NC_V1 " {%0}, [%1 + 24];" : "=r"(code32[6]) : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<32> {
-  static inline __device__ void load(unsigned int code32[8],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 32;
-    // FIXME: this is a non-coalesced load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];" :
-        "=r"(code32[0]), "=r"(code32[1]),
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];" :
-        "=r"(code32[4]), "=r"(code32[5]),
-        "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[8],
+            uint8_t* p,
+            int offset) {
+        p += offset * 32;
+        // FIXME: this is a non-coalesced load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];"
+            : "=r"(code32[0]), "=r"(code32[1]), "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];"
+            : "=r"(code32[4]), "=r"(code32[5]), "=r"(code32[6]), "=r"(code32[7])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<40> {
-  static inline __device__ void load(unsigned int code32[10],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 40;
-    // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 0];" :
-        "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 8];" :
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 16];" :
-        "=r"(code32[4]), "=r"(code32[5]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 24];" :
-        "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 32];" :
-        "=r"(code32[8]), "=r"(code32[9]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[10],
+            uint8_t* p,
+            int offset) {
+        p += offset * 40;
+        // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 0];"
+            : "=r"(code32[0]), "=r"(code32[1])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 8];"
+            : "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 16];"
+            : "=r"(code32[4]), "=r"(code32[5])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 24];"
+            : "=r"(code32[6]), "=r"(code32[7])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 32];"
+            : "=r"(code32[8]), "=r"(code32[9])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<48> {
-  static inline __device__ void load(unsigned int code32[12],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 48;
-    // FIXME: this is a non-coalesced load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];" :
-        "=r"(code32[0]), "=r"(code32[1]),
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];" :
-        "=r"(code32[4]), "=r"(code32[5]),
-        "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 32];" :
-        "=r"(code32[8]), "=r"(code32[9]),
-        "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[12],
+            uint8_t* p,
+            int offset) {
+        p += offset * 48;
+        // FIXME: this is a non-coalesced load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];"
+            : "=r"(code32[0]), "=r"(code32[1]), "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];"
+            : "=r"(code32[4]), "=r"(code32[5]), "=r"(code32[6]), "=r"(code32[7])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 32];"
+            : "=r"(code32[8]),
+              "=r"(code32[9]),
+              "=r"(code32[10]),
+              "=r"(code32[11])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<56> {
-  static inline __device__ void load(unsigned int code32[14],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 56;
-    // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 0];" :
-        "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 8];" :
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 16];" :
-        "=r"(code32[4]), "=r"(code32[5]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 24];" :
-        "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 32];" :
-        "=r"(code32[8]), "=r"(code32[9]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 40];" :
-        "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
-    asm(LD_NC_V2 " {%0, %1}, [%2 + 48];" :
-        "=r"(code32[12]), "=r"(code32[13]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[14],
+            uint8_t* p,
+            int offset) {
+        p += offset * 56;
+        // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 0];"
+            : "=r"(code32[0]), "=r"(code32[1])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 8];"
+            : "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 16];"
+            : "=r"(code32[4]), "=r"(code32[5])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 24];"
+            : "=r"(code32[6]), "=r"(code32[7])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 32];"
+            : "=r"(code32[8]), "=r"(code32[9])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 40];"
+            : "=r"(code32[10]), "=r"(code32[11])
+            : "l"(p));
+        asm(LD_NC_V2 " {%0, %1}, [%2 + 48];"
+            : "=r"(code32[12]), "=r"(code32[13])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<64> {
-  static inline __device__ void load(unsigned int code32[16],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 64;
-    // FIXME: this is a non-coalesced load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];" :
-        "=r"(code32[0]), "=r"(code32[1]),
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];" :
-        "=r"(code32[4]), "=r"(code32[5]),
-        "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 32];" :
-        "=r"(code32[8]), "=r"(code32[9]),
-        "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 48];" :
-        "=r"(code32[12]), "=r"(code32[13]),
-        "=r"(code32[14]), "=r"(code32[15]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[16],
+            uint8_t* p,
+            int offset) {
+        p += offset * 64;
+        // FIXME: this is a non-coalesced load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];"
+            : "=r"(code32[0]), "=r"(code32[1]), "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];"
+            : "=r"(code32[4]), "=r"(code32[5]), "=r"(code32[6]), "=r"(code32[7])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 32];"
+            : "=r"(code32[8]),
+              "=r"(code32[9]),
+              "=r"(code32[10]),
+              "=r"(code32[11])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 48];"
+            : "=r"(code32[12]),
+              "=r"(code32[13]),
+              "=r"(code32[14]),
+              "=r"(code32[15])
+            : "l"(p));
+    }
 };
 
-template<>
+template <>
 struct LoadCode32<96> {
-  static inline __device__ void load(unsigned int code32[24],
-                                     uint8_t* p,
-                                     int offset) {
-    p += offset * 96;
-    // FIXME: this is a non-coalesced load
-    // unfortunately need to reorganize memory layout by warp
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];" :
-        "=r"(code32[0]), "=r"(code32[1]),
-        "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];" :
-        "=r"(code32[4]), "=r"(code32[5]),
-        "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 32];" :
-        "=r"(code32[8]), "=r"(code32[9]),
-        "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 48];" :
-        "=r"(code32[12]), "=r"(code32[13]),
-        "=r"(code32[14]), "=r"(code32[15]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 64];" :
-        "=r"(code32[16]), "=r"(code32[17]),
-        "=r"(code32[18]), "=r"(code32[19]) : "l"(p));
-    asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 80];" :
-        "=r"(code32[20]), "=r"(code32[21]),
-        "=r"(code32[22]), "=r"(code32[23]) : "l"(p));
-  }
+    static inline __device__ void load(
+            unsigned int code32[24],
+            uint8_t* p,
+            int offset) {
+        p += offset * 96;
+        // FIXME: this is a non-coalesced load
+        // unfortunately need to reorganize memory layout by warp
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4];"
+            : "=r"(code32[0]), "=r"(code32[1]), "=r"(code32[2]), "=r"(code32[3])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 16];"
+            : "=r"(code32[4]), "=r"(code32[5]), "=r"(code32[6]), "=r"(code32[7])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 32];"
+            : "=r"(code32[8]),
+              "=r"(code32[9]),
+              "=r"(code32[10]),
+              "=r"(code32[11])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 48];"
+            : "=r"(code32[12]),
+              "=r"(code32[13]),
+              "=r"(code32[14]),
+              "=r"(code32[15])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 64];"
+            : "=r"(code32[16]),
+              "=r"(code32[17]),
+              "=r"(code32[18]),
+              "=r"(code32[19])
+            : "l"(p));
+        asm(LD_NC_V4 " {%0, %1, %2, %3}, [%4 + 80];"
+            : "=r"(code32[20]),
+              "=r"(code32[21]),
+              "=r"(code32[22]),
+              "=r"(code32[23])
+            : "l"(p));
+    }
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/PQScanMultiPassNoPrecomputed.cuh

@@ -5,44 +5,46 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/Index.h>
 #include <faiss/MetricType.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
-#include <faiss/gpu/utils/Tensor.cuh>
 #include <thrust/device_vector.h>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
 template <typename CentroidT>
-void runPQScanMultiPassNoPrecomputed(Tensor<float, 2, true>& queries,
-                                     Tensor<CentroidT, 2, true>& centroids,
-                                     Tensor<float, 3, true>& pqCentroidsInnermostCode,
-                                     Tensor<float, 2, true>& coarseDistances,
-                                     Tensor<int, 2, true>& coarseIndices,
-                                     bool useFloat16Lookup,
-                                     bool useMMCodeDistance,
-                                     bool interleavedCodeLayout,
-                                     int bitsPerSubQuantizer,
-                                     int numSubQuantizers,
-                                     int numSubQuantizerCodes,
-                                     thrust::device_vector<void*>& listCodes,
-                                     thrust::device_vector<void*>& listIndices,
-                                     IndicesOptions indicesOptions,
-                                     thrust::device_vector<int>& listLengths,
-                                     int maxListLength,
-                                     int k,
-                                     faiss::MetricType metric,
-                                     // output
-                                     Tensor<float, 2, true>& outDistances,
-                                     // output
-                                     Tensor<Index::idx_t, 2, true>& outIndices,
-                                     GpuResources* res);
+void runPQScanMultiPassNoPrecomputed(
+        Tensor<float, 2, true>& queries,
+        Tensor<CentroidT, 2, true>& centroids,
+        Tensor<float, 3, true>& pqCentroidsInnermostCode,
+        Tensor<float, 2, true>& coarseDistances,
+        Tensor<int, 2, true>& coarseIndices,
+        bool useFloat16Lookup,
+        bool useMMCodeDistance,
+        bool interleavedCodeLayout,
+        int bitsPerSubQuantizer,
+        int numSubQuantizers,
+        int numSubQuantizerCodes,
+        thrust::device_vector<void*>& listCodes,
+        thrust::device_vector<void*>& listIndices,
+        IndicesOptions indicesOptions,
+        thrust::device_vector<int>& listLengths,
+        int maxListLength,
+        int k,
+        faiss::MetricType metric,
+        // output
+        Tensor<float, 2, true>& outDistances,
+        // output
+        Tensor<Index::idx_t, 2, true>& outIndices,
+        GpuResources* res);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss
 
 #include <faiss/gpu/impl/PQScanMultiPassNoPrecomputed-inl.cuh>

faiss/gpu/impl/PQScanMultiPassPrecomputed.cuh

@@ -5,39 +5,41 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/Index.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
-#include <faiss/gpu/utils/Tensor.cuh>
-#include <faiss/gpu/utils/NoTypeTensor.cuh>
 #include <thrust/device_vector.h>
+#include <faiss/gpu/utils/NoTypeTensor.cuh>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
-void runPQScanMultiPassPrecomputed(Tensor<float, 2, true>& queries,
-                                   Tensor<float, 2, true>& precompTerm1,
-                                   NoTypeTensor<3, true>& precompTerm2,
-                                   NoTypeTensor<3, true>& precompTerm3,
-                                   Tensor<int, 2, true>& topQueryToCentroid,
-                                   bool useFloat16Lookup,
-                                   bool interleavedCodeLayout,
-                                   int bitsPerSubQuantizer,
-                                   int numSubQuantizers,
-                                   int numSubQuantizerCodes,
-                                   thrust::device_vector<void*>& listCodes,
-                                   thrust::device_vector<void*>& listIndices,
-                                   IndicesOptions indicesOptions,
-                                   thrust::device_vector<int>& listLengths,
-                                   int maxListLength,
-                                   int k,
-                                   // output
-                                   Tensor<float, 2, true>& outDistances,
-                                   // output
-                                   Tensor<Index::idx_t, 2, true>& outIndices,
-                                   GpuResources* res);
+void runPQScanMultiPassPrecomputed(
+        Tensor<float, 2, true>& queries,
+        Tensor<float, 2, true>& precompTerm1,
+        NoTypeTensor<3, true>& precompTerm2,
+        NoTypeTensor<3, true>& precompTerm3,
+        Tensor<int, 2, true>& topQueryToCentroid,
+        bool useFloat16Lookup,
+        bool interleavedCodeLayout,
+        int bitsPerSubQuantizer,
+        int numSubQuantizers,
+        int numSubQuantizerCodes,
+        thrust::device_vector<void*>& listCodes,
+        thrust::device_vector<void*>& listIndices,
+        IndicesOptions indicesOptions,
+        thrust::device_vector<int>& listLengths,
+        int maxListLength,
+        int k,
+        // output
+        Tensor<float, 2, true>& outDistances,
+        // output
+        Tensor<Index::idx_t, 2, true>& outIndices,
+        GpuResources* res);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/VectorResidual.cuh

@@ -5,35 +5,40 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // Calculates residual v_i - c_j for all v_i in vecs where j = vecToCentroid[i]
-void runCalcResidual(Tensor<float, 2, true>& vecs,
-                     Tensor<float, 2, true>& centroids,
-                     Tensor<int, 1, true>& vecToCentroid,
-                     Tensor<float, 2, true>& residuals,
-                     cudaStream_t stream);
+void runCalcResidual(
+        Tensor<float, 2, true>& vecs,
+        Tensor<float, 2, true>& centroids,
+        Tensor<int, 1, true>& vecToCentroid,
+        Tensor<float, 2, true>& residuals,
+        cudaStream_t stream);
 
-void runCalcResidual(Tensor<float, 2, true>& vecs,
-                     Tensor<half, 2, true>& centroids,
-                     Tensor<int, 1, true>& vecToCentroid,
-                     Tensor<float, 2, true>& residuals,
-                     cudaStream_t stream);
+void runCalcResidual(
+        Tensor<float, 2, true>& vecs,
+        Tensor<half, 2, true>& centroids,
+        Tensor<int, 1, true>& vecToCentroid,
+        Tensor<float, 2, true>& residuals,
+        cudaStream_t stream);
 
 // Gather vectors
-void runReconstruct(Tensor<int, 1, true>& listIds,
-                    Tensor<float, 2, true>& vecs,
-                    Tensor<float, 2, true>& out,
-                    cudaStream_t stream);
+void runReconstruct(
+        Tensor<int, 1, true>& listIds,
+        Tensor<float, 2, true>& vecs,
+        Tensor<float, 2, true>& out,
+        cudaStream_t stream);
 
-void runReconstruct(Tensor<int, 1, true>& listIds,
-                    Tensor<half, 2, true>& vecs,
-                    Tensor<float, 2, true>& out,
-                    cudaStream_t stream);
+void runReconstruct(
+        Tensor<int, 1, true>& listIds,
+        Tensor<half, 2, true>& vecs,
+        Tensor<float, 2, true>& out,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/impl/scan/IVFInterleavedImpl.cuh

@@ -10,66 +10,67 @@
 #include <faiss/gpu/impl/IVFInterleaved.cuh>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
 
-#define IVF_INTERLEAVED_IMPL(THREADS, WARP_Q, THREAD_Q)                 \
-                                                                        \
-void ivfInterleavedScanImpl_ ## WARP_Q ## _(                            \
-  Tensor<float, 2, true>& queries,                                      \
-  Tensor<int, 2, true>& listIds,                                        \
-  thrust::device_vector<void*>& listData,                               \
-  thrust::device_vector<void*>& listIndices,                            \
-  IndicesOptions indicesOptions,                                        \
-  thrust::device_vector<int>& listLengths,                              \
-  int k,                                                                \
-  faiss::MetricType metric,                                             \
-  bool useResidual,                                                     \
-  Tensor<float, 3, true>& residualBase,                                 \
-  GpuScalarQuantizer* scalarQ,                                          \
-  Tensor<float, 2, true>& outDistances,                                 \
-  Tensor<Index::idx_t, 2, true>& outIndices,                            \
-  GpuResources* res) {                                                  \
-  FAISS_ASSERT(k <= WARP_Q);                                            \
-                                                                        \
-  IVFINT_METRICS(THREADS, WARP_Q, THREAD_Q);                            \
-                                                                        \
-  CUDA_TEST_ERROR();                                                    \
-}
+#define IVF_INTERLEAVED_IMPL(THREADS, WARP_Q, THREAD_Q) \
+                                                        \
+    void ivfInterleavedScanImpl_##WARP_Q##_(            \
+            Tensor<float, 2, true>& queries,            \
+            Tensor<int, 2, true>& listIds,              \
+            thrust::device_vector<void*>& listData,     \
+            thrust::device_vector<void*>& listIndices,  \
+            IndicesOptions indicesOptions,              \
+            thrust::device_vector<int>& listLengths,    \
+            int k,                                      \
+            faiss::MetricType metric,                   \
+            bool useResidual,                           \
+            Tensor<float, 3, true>& residualBase,       \
+            GpuScalarQuantizer* scalarQ,                \
+            Tensor<float, 2, true>& outDistances,       \
+            Tensor<Index::idx_t, 2, true>& outIndices,  \
+            GpuResources* res) {                        \
+        FAISS_ASSERT(k <= WARP_Q);                      \
+                                                        \
+        IVFINT_METRICS(THREADS, WARP_Q, THREAD_Q);      \
+                                                        \
+        CUDA_TEST_ERROR();                              \
+    }
 
-#define IVF_INTERLEAVED_DECL(WARP_Q)                                    \
-                                                                        \
-void ivfInterleavedScanImpl_ ## WARP_Q ## _(                            \
-  Tensor<float, 2, true>& queries,                                      \
-  Tensor<int, 2, true>& listIds,                                        \
-  thrust::device_vector<void*>& listData,                               \
-  thrust::device_vector<void*>& listIndices,                            \
-  IndicesOptions indicesOptions,                                        \
-  thrust::device_vector<int>& listLengths,                              \
-  int k,                                                                \
-  faiss::MetricType metric,                                             \
-  bool useResidual,                                                     \
-  Tensor<float, 3, true>& residualBase,                                 \
-  GpuScalarQuantizer* scalarQ,                                          \
-  Tensor<float, 2, true>& outDistances,                                 \
-  Tensor<Index::idx_t, 2, true>& outIndices,                            \
-  GpuResources* res)
+#define IVF_INTERLEAVED_DECL(WARP_Q)                   \
+                                                       \
+    void ivfInterleavedScanImpl_##WARP_Q##_(           \
+            Tensor<float, 2, true>& queries,           \
+            Tensor<int, 2, true>& listIds,             \
+            thrust::device_vector<void*>& listData,    \
+            thrust::device_vector<void*>& listIndices, \
+            IndicesOptions indicesOptions,             \
+            thrust::device_vector<int>& listLengths,   \
+            int k,                                     \
+            faiss::MetricType metric,                  \
+            bool useResidual,                          \
+            Tensor<float, 3, true>& residualBase,      \
+            GpuScalarQuantizer* scalarQ,               \
+            Tensor<float, 2, true>& outDistances,      \
+            Tensor<Index::idx_t, 2, true>& outIndices, \
+            GpuResources* res)
 
-#define IVF_INTERLEAVED_CALL(WARP_Q)                    \
-  ivfInterleavedScanImpl_ ## WARP_Q ## _(               \
-  queries,                                              \
-  listIds,                                              \
-  listData,                                             \
-  listIndices,                                          \
-  indicesOptions,                                       \
-  listLengths,                                          \
-  k,                                                    \
-  metric,                                               \
-  useResidual,                                          \
-  residualBase,                                         \
-  scalarQ,                                              \
-  outDistances,                                         \
-  outIndices,                                           \
-  res)
+#define IVF_INTERLEAVED_CALL(WARP_Q)    \
+    ivfInterleavedScanImpl_##WARP_Q##_( \
+            queries,                    \
+            listIds,                    \
+            listData,                   \
+            listIndices,                \
+            indicesOptions,             \
+            listLengths,                \
+            k,                          \
+            metric,                     \
+            useResidual,                \
+            residualBase,               \
+            scalarQ,                    \
+            outDistances,               \
+            outIndices,                 \
+            res)
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 IVF_INTERLEAVED_DECL(1);
 IVF_INTERLEAVED_DECL(32);
@@ -83,4 +84,5 @@ IVF_INTERLEAVED_DECL(1024);
 IVF_INTERLEAVED_DECL(2048);
 #endif
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/BlockSelectKernel.cuh

@@ -9,127 +9,157 @@
 
 #include <faiss/gpu/utils/Select.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
-template <typename K,
-          typename IndexType,
-          bool Dir,
-          int NumWarpQ,
-          int NumThreadQ,
-          int ThreadsPerBlock>
-__global__ void blockSelect(Tensor<K, 2, true> in,
-                            Tensor<K, 2, true> outK,
-                            Tensor<IndexType, 2, true> outV,
-                            K initK,
-                            IndexType initV,
-                            int k) {
-  constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
+template <
+        typename K,
+        typename IndexType,
+        bool Dir,
+        int NumWarpQ,
+        int NumThreadQ,
+        int ThreadsPerBlock>
+__global__ void blockSelect(
+        Tensor<K, 2, true> in,
+        Tensor<K, 2, true> outK,
+        Tensor<IndexType, 2, true> outV,
+        K initK,
+        IndexType initV,
+        int k) {
+    constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
 
-  __shared__ K smemK[kNumWarps * NumWarpQ];
-  __shared__ IndexType smemV[kNumWarps * NumWarpQ];
+    __shared__ K smemK[kNumWarps * NumWarpQ];
+    __shared__ IndexType smemV[kNumWarps * NumWarpQ];
 
-  BlockSelect<K, IndexType, Dir, Comparator<K>,
-              NumWarpQ, NumThreadQ, ThreadsPerBlock>
-    heap(initK, initV, smemK, smemV, k);
+    BlockSelect<
+            K,
+            IndexType,
+            Dir,
+            Comparator<K>,
+            NumWarpQ,
+            NumThreadQ,
+            ThreadsPerBlock>
+            heap(initK, initV, smemK, smemV, k);
 
-  // Grid is exactly sized to rows available
-  int row = blockIdx.x;
+    // Grid is exactly sized to rows available
+    int row = blockIdx.x;
 
-  int i = threadIdx.x;
-  K* inStart = in[row][i].data();
+    int i = threadIdx.x;
+    K* inStart = in[row][i].data();
 
-  // Whole warps must participate in the selection
-  int limit = utils::roundDown(in.getSize(1), kWarpSize);
+    // Whole warps must participate in the selection
+    int limit = utils::roundDown(in.getSize(1), kWarpSize);
 
-  for (; i < limit; i += ThreadsPerBlock) {
-    heap.add(*inStart, (IndexType) i);
-    inStart += ThreadsPerBlock;
-  }
+    for (; i < limit; i += ThreadsPerBlock) {
+        heap.add(*inStart, (IndexType)i);
+        inStart += ThreadsPerBlock;
+    }
 
-  // Handle last remainder fraction of a warp of elements
-  if (i < in.getSize(1)) {
-    heap.addThreadQ(*inStart, (IndexType) i);
-  }
+    // Handle last remainder fraction of a warp of elements
+    if (i < in.getSize(1)) {
+        heap.addThreadQ(*inStart, (IndexType)i);
+    }
 
-  heap.reduce();
+    heap.reduce();
 
-  for (int i = threadIdx.x; i < k; i += ThreadsPerBlock) {
-    outK[row][i] = smemK[i];
-    outV[row][i] = smemV[i];
-  }
+    for (int i = threadIdx.x; i < k; i += ThreadsPerBlock) {
+        outK[row][i] = smemK[i];
+        outV[row][i] = smemV[i];
+    }
 }
 
-template <typename K,
-          typename IndexType,
-          bool Dir,
-          int NumWarpQ,
-          int NumThreadQ,
-          int ThreadsPerBlock>
-__global__ void blockSelectPair(Tensor<K, 2, true> inK,
-                                Tensor<IndexType, 2, true> inV,
-                                Tensor<K, 2, true> outK,
-                                Tensor<IndexType, 2, true> outV,
-                                K initK,
-                                IndexType initV,
-                                int k) {
-  constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
+template <
+        typename K,
+        typename IndexType,
+        bool Dir,
+        int NumWarpQ,
+        int NumThreadQ,
+        int ThreadsPerBlock>
+__global__ void blockSelectPair(
+        Tensor<K, 2, true> inK,
+        Tensor<IndexType, 2, true> inV,
+        Tensor<K, 2, true> outK,
+        Tensor<IndexType, 2, true> outV,
+        K initK,
+        IndexType initV,
+        int k) {
+    constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
 
-  __shared__ K smemK[kNumWarps * NumWarpQ];
-  __shared__ IndexType smemV[kNumWarps * NumWarpQ];
+    __shared__ K smemK[kNumWarps * NumWarpQ];
+    __shared__ IndexType smemV[kNumWarps * NumWarpQ];
 
-  BlockSelect<K, IndexType, Dir, Comparator<K>,
-              NumWarpQ, NumThreadQ, ThreadsPerBlock>
-    heap(initK, initV, smemK, smemV, k);
+    BlockSelect<
+            K,
+            IndexType,
+            Dir,
+            Comparator<K>,
+            NumWarpQ,
+            NumThreadQ,
+            ThreadsPerBlock>
+            heap(initK, initV, smemK, smemV, k);
 
-  // Grid is exactly sized to rows available
-  int row = blockIdx.x;
+    // Grid is exactly sized to rows available
+    int row = blockIdx.x;
 
-  int i = threadIdx.x;
-  K* inKStart = inK[row][i].data();
-  IndexType* inVStart = inV[row][i].data();
+    int i = threadIdx.x;
+    K* inKStart = inK[row][i].data();
+    IndexType* inVStart = inV[row][i].data();
 
-  // Whole warps must participate in the selection
-  int limit = utils::roundDown(inK.getSize(1), kWarpSize);
+    // Whole warps must participate in the selection
+    int limit = utils::roundDown(inK.getSize(1), kWarpSize);
 
-  for (; i < limit; i += ThreadsPerBlock) {
-    heap.add(*inKStart, *inVStart);
-    inKStart += ThreadsPerBlock;
-    inVStart += ThreadsPerBlock;
-  }
+    for (; i < limit; i += ThreadsPerBlock) {
+        heap.add(*inKStart, *inVStart);
+        inKStart += ThreadsPerBlock;
+        inVStart += ThreadsPerBlock;
+    }
 
-  // Handle last remainder fraction of a warp of elements
-  if (i < inK.getSize(1)) {
-    heap.addThreadQ(*inKStart, *inVStart);
-  }
+    // Handle last remainder fraction of a warp of elements
+    if (i < inK.getSize(1)) {
+        heap.addThreadQ(*inKStart, *inVStart);
+    }
 
-  heap.reduce();
+    heap.reduce();
 
-  for (int i = threadIdx.x; i < k; i += ThreadsPerBlock) {
-    outK[row][i] = smemK[i];
-    outV[row][i] = smemV[i];
-  }
+    for (int i = threadIdx.x; i < k; i += ThreadsPerBlock) {
+        outK[row][i] = smemK[i];
+        outV[row][i] = smemV[i];
+    }
 }
 
-void runBlockSelect(Tensor<float, 2, true>& in,
-                    Tensor<float, 2, true>& outKeys,
-                    Tensor<int, 2, true>& outIndices,
-                    bool dir, int k, cudaStream_t stream);
+void runBlockSelect(
+        Tensor<float, 2, true>& in,
+        Tensor<float, 2, true>& outKeys,
+        Tensor<int, 2, true>& outIndices,
+        bool dir,
+        int k,
+        cudaStream_t stream);
 
-void runBlockSelectPair(Tensor<float, 2, true>& inKeys,
-                        Tensor<int, 2, true>& inIndices,
-                        Tensor<float, 2, true>& outKeys,
-                        Tensor<int, 2, true>& outIndices,
-                        bool dir, int k, cudaStream_t stream);
+void runBlockSelectPair(
+        Tensor<float, 2, true>& inKeys,
+        Tensor<int, 2, true>& inIndices,
+        Tensor<float, 2, true>& outKeys,
+        Tensor<int, 2, true>& outIndices,
+        bool dir,
+        int k,
+        cudaStream_t stream);
 
-void runBlockSelect(Tensor<half, 2, true>& in,
-                    Tensor<half, 2, true>& outKeys,
-                    Tensor<int, 2, true>& outIndices,
-                    bool dir, int k, cudaStream_t stream);
+void runBlockSelect(
+        Tensor<half, 2, true>& in,
+        Tensor<half, 2, true>& outKeys,
+        Tensor<int, 2, true>& outIndices,
+        bool dir,
+        int k,
+        cudaStream_t stream);
 
-void runBlockSelectPair(Tensor<half, 2, true>& inKeys,
-                        Tensor<int, 2, true>& inIndices,
-                        Tensor<half, 2, true>& outKeys,
-                        Tensor<int, 2, true>& outIndices,
-                        bool dir, int k, cudaStream_t stream);
+void runBlockSelectPair(
+        Tensor<half, 2, true>& inKeys,
+        Tensor<int, 2, true>& inIndices,
+        Tensor<half, 2, true>& outKeys,
+        Tensor<int, 2, true>& outIndices,
+        bool dir,
+        int k,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/Comparators.cuh

@@ -5,42 +5,43 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cuda.h>
 #include <faiss/gpu/utils/Float16.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T>
 struct Comparator {
-  __device__ static inline bool lt(T a, T b) {
-    return a < b;
-  }
+    __device__ static inline bool lt(T a, T b) {
+        return a < b;
+    }
 
-  __device__ static inline bool gt(T a, T b) {
-    return a > b;
-  }
+    __device__ static inline bool gt(T a, T b) {
+        return a > b;
+    }
 };
 
 template <>
 struct Comparator<half> {
-  __device__ static inline bool lt(half a, half b) {
+    __device__ static inline bool lt(half a, half b) {
 #if FAISS_USE_FULL_FLOAT16
-    return __hlt(a, b);
+        return __hlt(a, b);
 #else
-    return __half2float(a) < __half2float(b);
+        return __half2float(a) < __half2float(b);
 #endif // FAISS_USE_FULL_FLOAT16
-  }
+    }
 
-  __device__ static inline bool gt(half a, half b) {
+    __device__ static inline bool gt(half a, half b) {
 #if FAISS_USE_FULL_FLOAT16
-    return __hgt(a, b);
+        return __hgt(a, b);
 #else
-    return __half2float(a) > __half2float(b);
+        return __half2float(a) > __half2float(b);
 #endif // FAISS_USE_FULL_FLOAT16
-  }
+    }
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/ConversionOperators.cuh

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/MetricType.h>
@@ -16,7 +15,8 @@
 #include <thrust/execution_policy.h>
 #include <thrust/transform.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 //
 // Conversion utilities
@@ -24,111 +24,134 @@ namespace faiss { namespace gpu {
 
 template <typename From, typename To>
 struct Convert {
-  inline __device__ To operator()(From v) const {
-    return (To) v;
-  }
+    inline __device__ To operator()(From v) const {
+        return (To)v;
+    }
 };
 
 template <>
 struct Convert<float, half> {
-  inline __device__ half operator()(float v) const {
-    return __float2half(v);
-  }
+    inline __device__ half operator()(float v) const {
+        return __float2half(v);
+    }
 };
 
 template <>
 struct Convert<half, float> {
-  inline __device__ float operator()(half v) const {
-    return __half2float(v);
-  }
+    inline __device__ float operator()(half v) const {
+        return __half2float(v);
+    }
 };
 
 template <typename T>
-struct ConvertTo {
-};
+struct ConvertTo {};
 
 template <>
 struct ConvertTo<float> {
-  static inline __device__ float to(float v) { return v; }
-  static inline __device__ float to(half v) { return __half2float(v); }
+    static inline __device__ float to(float v) {
+        return v;
+    }
+    static inline __device__ float to(half v) {
+        return __half2float(v);
+    }
 };
 
 template <>
 struct ConvertTo<float2> {
-  static inline __device__ float2 to(float2 v) { return v; }
-  static inline __device__ float2 to(half2 v) { return __half22float2(v); }
+    static inline __device__ float2 to(float2 v) {
+        return v;
+    }
+    static inline __device__ float2 to(half2 v) {
+        return __half22float2(v);
+    }
 };
 
 template <>
 struct ConvertTo<float4> {
-  static inline __device__ float4 to(float4 v) { return v; }
-  static inline __device__ float4 to(Half4 v) { return half4ToFloat4(v); }
+    static inline __device__ float4 to(float4 v) {
+        return v;
+    }
+    static inline __device__ float4 to(Half4 v) {
+        return half4ToFloat4(v);
+    }
 };
 
 template <>
 struct ConvertTo<half> {
-  static inline __device__ half to(float v) { return __float2half(v); }
-  static inline __device__ half to(half v) { return v; }
+    static inline __device__ half to(float v) {
+        return __float2half(v);
+    }
+    static inline __device__ half to(half v) {
+        return v;
+    }
 };
 
 template <>
 struct ConvertTo<half2> {
-  static inline __device__ half2 to(float2 v) { return __float22half2_rn(v); }
-  static inline __device__ half2 to(half2 v) { return v; }
+    static inline __device__ half2 to(float2 v) {
+        return __float22half2_rn(v);
+    }
+    static inline __device__ half2 to(half2 v) {
+        return v;
+    }
 };
 
 template <>
 struct ConvertTo<Half4> {
-  static inline __device__ Half4 to(float4 v) { return float4ToHalf4(v); }
-  static inline __device__ Half4 to(Half4 v) { return v; }
+    static inline __device__ Half4 to(float4 v) {
+        return float4ToHalf4(v);
+    }
+    static inline __device__ Half4 to(Half4 v) {
+        return v;
+    }
 };
 
 // Tensor conversion
 template <typename From, typename To>
-void runConvert(const From* in,
-                To* out,
-                size_t num,
-                cudaStream_t stream) {
-  thrust::transform(thrust::cuda::par.on(stream),
-                    in, in + num, out, Convert<From, To>());
+void runConvert(const From* in, To* out, size_t num, cudaStream_t stream) {
+    thrust::transform(
+            thrust::cuda::par.on(stream),
+            in,
+            in + num,
+            out,
+            Convert<From, To>());
 }
 
 template <typename From, typename To, int Dim>
-void convertTensor(cudaStream_t stream,
-                   Tensor<From, Dim, true>& in,
-                   Tensor<To, Dim, true>& out) {
-  FAISS_ASSERT(in.numElements() == out.numElements());
+void convertTensor(
+        cudaStream_t stream,
+        Tensor<From, Dim, true>& in,
+        Tensor<To, Dim, true>& out) {
+    FAISS_ASSERT(in.numElements() == out.numElements());
 
-  runConvert<From, To>(in.data(), out.data(), in.numElements(), stream);
+    runConvert<From, To>(in.data(), out.data(), in.numElements(), stream);
 }
 
 template <typename From, typename To, int Dim>
-DeviceTensor<To, Dim, true>
-convertTensorTemporary(GpuResources* res,
-                       cudaStream_t stream,
-                       Tensor<From, Dim, true>& in) {
-  FAISS_ASSERT(res);
-  DeviceTensor<To, Dim, true> out(
-    res, makeTempAlloc(AllocType::Other, stream),
-    in.sizes());
+DeviceTensor<To, Dim, true> convertTensorTemporary(
+        GpuResources* res,
+        cudaStream_t stream,
+        Tensor<From, Dim, true>& in) {
+    FAISS_ASSERT(res);
+    DeviceTensor<To, Dim, true> out(
+            res, makeTempAlloc(AllocType::Other, stream), in.sizes());
 
-  convertTensor(stream, in, out);
-  return out;
+    convertTensor(stream, in, out);
+    return out;
 }
 
 template <typename From, typename To, int Dim>
-DeviceTensor<To, Dim, true>
-convertTensorNonTemporary(GpuResources* res,
-                          cudaStream_t stream,
-                          Tensor<From, Dim, true>& in) {
-  FAISS_ASSERT(res);
-  DeviceTensor<To, Dim, true> out(
-    res, makeDevAlloc(AllocType::Other, stream),
-    in.sizes());
+DeviceTensor<To, Dim, true> convertTensorNonTemporary(
+        GpuResources* res,
+        cudaStream_t stream,
+        Tensor<From, Dim, true>& in) {
+    FAISS_ASSERT(res);
+    DeviceTensor<To, Dim, true> out(
+            res, makeDevAlloc(AllocType::Other, stream), in.sizes());
 
-  convertTensor(stream, in, out);
-  return out;
+    convertTensor(stream, in, out);
+    return out;
 }
 
-
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/CopyUtils.cuh

@@ -5,138 +5,134 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 #include <faiss/gpu/utils/HostTensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// Ensure the memory at `p` is either on the given device, or copy it
 /// to the device in a new temporary allocation.
 template <typename T, int Dim>
 DeviceTensor<T, Dim, true> toDeviceTemporary(
-  GpuResources* resources,
-  int dstDevice,
-  T* src,
-  cudaStream_t stream,
-  std::initializer_list<int> sizes) {
-  int dev = getDeviceForAddress(src);
-  DeviceTensor<T, Dim, true> oldT(src, sizes);
+        GpuResources* resources,
+        int dstDevice,
+        T* src,
+        cudaStream_t stream,
+        std::initializer_list<int> sizes) {
+    int dev = getDeviceForAddress(src);
+    DeviceTensor<T, Dim, true> oldT(src, sizes);
 
-  if (dev == dstDevice) {
-    // On device we expect
-    return oldT;
-  } else {
-    // On different device or on host
-    DeviceScope scope(dstDevice);
+    if (dev == dstDevice) {
+        // On device we expect
+        return oldT;
+    } else {
+        // On different device or on host
+        DeviceScope scope(dstDevice);
 
-    DeviceTensor<T, Dim, true> newT(
-      resources, makeTempAlloc(AllocType::Other, stream), sizes);
+        DeviceTensor<T, Dim, true> newT(
+                resources, makeTempAlloc(AllocType::Other, stream), sizes);
 
-    newT.copyFrom(oldT, stream);
-    return newT;
-  }
+        newT.copyFrom(oldT, stream);
+        return newT;
+    }
 }
 
 template <typename T, int Dim>
 DeviceTensor<T, Dim, true> toDeviceNonTemporary(
-  GpuResources* resources,
-  int dstDevice,
-  T* src,
-  cudaStream_t stream,
-  std::initializer_list<int> sizes) {
-  int dev = getDeviceForAddress(src);
-  DeviceTensor<T, Dim, true> oldT(src, sizes);
+        GpuResources* resources,
+        int dstDevice,
+        T* src,
+        cudaStream_t stream,
+        std::initializer_list<int> sizes) {
+    int dev = getDeviceForAddress(src);
+    DeviceTensor<T, Dim, true> oldT(src, sizes);
 
-  if (dev == dstDevice) {
-    // On device we expect
-    return oldT;
-  } else {
-    // On different device or on host
-    DeviceScope scope(dstDevice);
+    if (dev == dstDevice) {
+        // On device we expect
+        return oldT;
+    } else {
+        // On different device or on host
+        DeviceScope scope(dstDevice);
 
-    DeviceTensor<T, Dim, true> newT(
-      resources, makeDevAlloc(AllocType::Other, stream), sizes);
+        DeviceTensor<T, Dim, true> newT(
+                resources, makeDevAlloc(AllocType::Other, stream), sizes);
 
-    newT.copyFrom(oldT, stream);
-    return newT;
-  }
+        newT.copyFrom(oldT, stream);
+        return newT;
+    }
 }
 
 template <typename T>
 DeviceTensor<T, 1, true> toDeviceTemporary(
-  GpuResources* resources,
-  const std::vector<T>& src,
-  cudaStream_t stream,
-  int device = -1) {
-  // Uses the current device if device == -1
-  DeviceScope scope(device);
+        GpuResources* resources,
+        const std::vector<T>& src,
+        cudaStream_t stream,
+        int device = -1) {
+    // Uses the current device if device == -1
+    DeviceScope scope(device);
 
-  FAISS_ASSERT(src.size() <
-               (size_t) std::numeric_limits<int>::max());
+    FAISS_ASSERT(src.size() < (size_t)std::numeric_limits<int>::max());
 
-  DeviceTensor<T, 1, true> out(
-    resources, makeTempAlloc(AllocType::Other, stream),
-    {(int) src.size()});
+    DeviceTensor<T, 1, true> out(
+            resources,
+            makeTempAlloc(AllocType::Other, stream),
+            {(int)src.size()});
 
-  out.copyFrom(src, stream);
+    out.copyFrom(src, stream);
 
-  return out;
+    return out;
 }
 
 /// Copies data to the CPU, if it is not already on the CPU
 template <typename T, int Dim>
-HostTensor<T, Dim, true> toHost(T* src,
-                                cudaStream_t stream,
-                                std::initializer_list<int> sizes) {
-  int dev = getDeviceForAddress(src);
+HostTensor<T, Dim, true> toHost(
+        T* src,
+        cudaStream_t stream,
+        std::initializer_list<int> sizes) {
+    int dev = getDeviceForAddress(src);
 
-  if (dev == -1) {
-    // Already on the CPU, just wrap in a HostTensor that doesn't own this
-    // memory
-    return HostTensor<T, Dim, true>(src, sizes);
-  } else {
-    HostTensor<T, Dim, true> out(sizes);
-    Tensor<T, Dim, true> devData(src, sizes);
-    out.copyFrom(devData, stream);
+    if (dev == -1) {
+        // Already on the CPU, just wrap in a HostTensor that doesn't own this
+        // memory
+        return HostTensor<T, Dim, true>(src, sizes);
+    } else {
+        HostTensor<T, Dim, true> out(sizes);
+        Tensor<T, Dim, true> devData(src, sizes);
+        out.copyFrom(devData, stream);
 
-    return out;
-  }
+        return out;
+    }
 }
 
 /// Copies a device array's allocation to an address, if necessary
 template <typename T>
 inline void fromDevice(T* src, T* dst, size_t num, cudaStream_t stream) {
-  // It is possible that the array already represents memory at `p`,
-  // in which case no copy is needed
-  if (src == dst) {
-    return;
-  }
+    // It is possible that the array already represents memory at `p`,
+    // in which case no copy is needed
+    if (src == dst) {
+        return;
+    }
 
-  int dev = getDeviceForAddress(dst);
+    int dev = getDeviceForAddress(dst);
 
-  if (dev == -1) {
-    CUDA_VERIFY(cudaMemcpyAsync(dst,
-                                src,
-                                num * sizeof(T),
-                                cudaMemcpyDeviceToHost,
-                                stream));
-  } else {
-    CUDA_VERIFY(cudaMemcpyAsync(dst,
-                                src,
-                                num * sizeof(T),
-                                cudaMemcpyDeviceToDevice,
-                                stream));
-  }
+    if (dev == -1) {
+        CUDA_VERIFY(cudaMemcpyAsync(
+                dst, src, num * sizeof(T), cudaMemcpyDeviceToHost, stream));
+    } else {
+        CUDA_VERIFY(cudaMemcpyAsync(
+                dst, src, num * sizeof(T), cudaMemcpyDeviceToDevice, stream));
+    }
 }
 
 /// Copies a device array's allocation to an address, if necessary
 template <typename T, int Dim>
 void fromDevice(Tensor<T, Dim, true>& src, T* dst, cudaStream_t stream) {
-  FAISS_ASSERT(src.isContiguous());
-  fromDevice(src.data(), dst, src.numElements(), stream);
+    FAISS_ASSERT(src.isContiguous());
+    fromDevice(src.data(), dst, src.numElements(), stream);
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/DeviceDefs.cuh

@@ -5,12 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cuda.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // We require at least CUDA 8.0 for compilation
 #if CUDA_VERSION < 8000
@@ -22,12 +22,11 @@ constexpr int kWarpSize = 32;
 
 // This is a memory barrier for intra-warp writes to shared memory.
 __forceinline__ __device__ void warpFence() {
-
 #if CUDA_VERSION >= 9000
-  __syncwarp();
+    __syncwarp();
 #else
-  // For the time being, assume synchronicity.
-  //  __threadfence_block();
+    // For the time being, assume synchronicity.
+    //  __threadfence_block();
 #endif
 }
 
@@ -40,4 +39,5 @@ __forceinline__ __device__ void warpFence() {
 #define GPU_MAX_SELECTION_K 1024
 #endif
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/DeviceTensor-inl.cuh

@@ -5,140 +5,188 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #include <utility> // std::move
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor() :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>() {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor()
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>() {}
+
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
+        DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>() {
+    this->operator=(std::move(t));
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
-  DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>() {
-  this->operator=(std::move(t));
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>& DeviceTensor<
+        T,
+        Dim,
+        InnerContig,
+        IndexT,
+        PtrTraits>::
+operator=(DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) {
+    this->Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
+            std::move(t));
+
+    this->reservation_ = std::move(t.reservation_);
+    return *this;
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
-  DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) {
-  this->Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
-    std::move(t));
-
-  this->reservation_ = std::move(t.reservation_);
-  return *this;
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::~DeviceTensor() {
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::~DeviceTensor() {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
+        GpuResources* res,
+        const AllocInfo& info,
+        const IndexT sizes[Dim])
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes) {
+    this->reservation_ = std::move(
+            res->allocMemoryHandle(AllocRequest(info, this->getSizeInBytes())));
+    this->data_ = (T*)reservation_.get();
+
+    FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
-  GpuResources* res,
-  const AllocInfo& info,
-  const IndexT sizes[Dim]) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
+        GpuResources* res,
+        const AllocInfo& info,
+        std::initializer_list<IndexT> sizes)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes) {
+    this->reservation_ = std::move(
+            res->allocMemoryHandle(AllocRequest(info, this->getSizeInBytes())));
+    this->data_ = (T*)reservation_.get();
 
-  this->reservation_ =
-    std::move(
-      res->allocMemoryHandle(AllocRequest(info, this->getSizeInBytes())));
-  this->data_ = (T*) reservation_.get();
-
-  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
+    FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
-  GpuResources* res,
-  const AllocInfo& info,
-  std::initializer_list<IndexT> sizes) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
+        DataPtrType data,
+        const IndexT sizes[Dim])
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes) {}
 
-  this->reservation_ =
-    std::move(res->allocMemoryHandle(AllocRequest(info, this->getSizeInBytes())));
-  this->data_ = (T*) reservation_.get();
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
+        DataPtrType data,
+        std::initializer_list<IndexT> sizes)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes) {}
 
-  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
+        DataPtrType data,
+        const IndexT sizes[Dim],
+        const IndexT strides[Dim])
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes, strides) {
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
-  DataPtrType data,
-  const IndexT sizes[Dim]) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
+        GpuResources* res,
+        const AllocInfo& info,
+        Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(
+                  nullptr,
+                  t.sizes(),
+                  t.strides()) {
+    this->reservation_ = std::move(
+            res->allocMemoryHandle(AllocRequest(info, this->getSizeInBytes())));
+    this->data_ = (T*)reservation_.get();
+
+    FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
+
+    this->copyFrom(t, info.stream);
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
-  DataPtrType data,
-  std::initializer_list<IndexT> sizes) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>& DeviceTensor<
+        T,
+        Dim,
+        InnerContig,
+        IndexT,
+        PtrTraits>::zero(cudaStream_t stream) {
+    if (this->data_) {
+        // Region must be contiguous
+        FAISS_ASSERT(this->isContiguous());
+
+        CUDA_VERIFY(cudaMemsetAsync(
+                this->data_, 0, this->getSizeInBytes(), stream));
+    }
+
+    return *this;
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
-  DataPtrType data,
-  const IndexT sizes[Dim],
-  const IndexT strides[Dim]) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes, strides) {
-}
-
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
-  GpuResources* res,
-  const AllocInfo& info,
-  Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(
-      nullptr, t.sizes(), t.strides()) {
-
-  this->reservation_ =
-    std::move(res->allocMemoryHandle(AllocRequest(info, this->getSizeInBytes())));
-  this->data_ = (T*) reservation_.get();
-
-  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
-
-  this->copyFrom(t, info.stream);
-}
-
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
-DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::zero(
-  cudaStream_t stream) {
-  if (this->data_) {
-    // Region must be contiguous
-    FAISS_ASSERT(this->isContiguous());
-
-    CUDA_VERIFY(cudaMemsetAsync(
-                  this->data_, 0, this->getSizeInBytes(), stream));
-  }
-
-  return *this;
-}
-
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/DeviceTensor.cuh

@@ -5,84 +5,91 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/gpu/utils/Tensor.cuh>
 #include <faiss/gpu/GpuResources.h>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
-template <typename T,
-          int Dim,
-          bool InnerContig = false,
-          typename IndexT = int,
-          template <typename U> class PtrTraits = traits::DefaultPtrTraits>
+template <
+        typename T,
+        int Dim,
+        bool InnerContig = false,
+        typename IndexT = int,
+        template <typename U> class PtrTraits = traits::DefaultPtrTraits>
 class DeviceTensor : public Tensor<T, Dim, InnerContig, IndexT, PtrTraits> {
- public:
-  typedef IndexT IndexType;
-  typedef typename PtrTraits<T>::PtrType DataPtrType;
+   public:
+    typedef IndexT IndexType;
+    typedef typename PtrTraits<T>::PtrType DataPtrType;
 
-  /// Default constructor
-  __host__ DeviceTensor();
+    /// Default constructor
+    __host__ DeviceTensor();
 
-  /// Destructor
-  __host__ ~DeviceTensor();
+    /// Destructor
+    __host__ ~DeviceTensor();
 
-  /// Move constructor
-  __host__ DeviceTensor(DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
+    /// Move constructor
+    __host__ DeviceTensor(
+            DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
 
-  /// Move assignment
-  __host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
-  operator=(DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
+    /// Move assignment
+    __host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>& operator=(
+            DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
 
-  /// Constructs a tensor of the given size, allocating memory for it
-  /// via temporary or other allocation.
-  /// `stream` specifies the stream on which the memory will be used
-  __host__ DeviceTensor(GpuResources* res,
-                        const AllocInfo& info,
-                        const IndexT sizes[Dim]);
+    /// Constructs a tensor of the given size, allocating memory for it
+    /// via temporary or other allocation.
+    /// `stream` specifies the stream on which the memory will be used
+    __host__ DeviceTensor(
+            GpuResources* res,
+            const AllocInfo& info,
+            const IndexT sizes[Dim]);
 
-  /// Constructs a tensor of the given size, allocating memory for it
-  /// via temporary or other allocation.
-  /// `stream` specifies the stream on which the memory will be used
-  __host__ DeviceTensor(GpuResources* res,
-                        const AllocInfo& info,
-                        std::initializer_list<IndexT> sizes);
+    /// Constructs a tensor of the given size, allocating memory for it
+    /// via temporary or other allocation.
+    /// `stream` specifies the stream on which the memory will be used
+    __host__ DeviceTensor(
+            GpuResources* res,
+            const AllocInfo& info,
+            std::initializer_list<IndexT> sizes);
 
-  /// Constructs a tensor of the given size and stride, referencing a
-  /// memory region we do not own
-  __host__ DeviceTensor(DataPtrType data,
-                        const IndexT sizes[Dim]);
+    /// Constructs a tensor of the given size and stride, referencing a
+    /// memory region we do not own
+    __host__ DeviceTensor(DataPtrType data, const IndexT sizes[Dim]);
 
-  /// Constructs a tensor of the given size and stride, referencing a
-  /// memory region we do not own
-  __host__ DeviceTensor(DataPtrType data,
-                        std::initializer_list<IndexT> sizes);
+    /// Constructs a tensor of the given size and stride, referencing a
+    /// memory region we do not own
+    __host__ DeviceTensor(
+            DataPtrType data,
+            std::initializer_list<IndexT> sizes);
 
-  /// Constructs a tensor of the given size and stride, referencing a
-  /// memory region we do not own
-  __host__ DeviceTensor(DataPtrType data,
-                        const IndexT sizes[Dim],
-                        const IndexT strides[Dim]);
+    /// Constructs a tensor of the given size and stride, referencing a
+    /// memory region we do not own
+    __host__ DeviceTensor(
+            DataPtrType data,
+            const IndexT sizes[Dim],
+            const IndexT strides[Dim]);
 
-  /// Copies a tensor into ourselves, allocating memory for it.
-  /// `stream` specifies the stream of the copy and thus the stream on which the
-  /// memory will initially be used.
-  __host__ DeviceTensor(GpuResources* res,
-                        const AllocInfo& info,
-                        Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t);
+    /// Copies a tensor into ourselves, allocating memory for it.
+    /// `stream` specifies the stream of the copy and thus the stream on which
+    /// the memory will initially be used.
+    __host__ DeviceTensor(
+            GpuResources* res,
+            const AllocInfo& info,
+            Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t);
 
-  /// Call to zero out memory
-  __host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
-  zero(cudaStream_t stream);
+    /// Call to zero out memory
+    __host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>& zero(
+            cudaStream_t stream);
 
- private:
-  /// If we own the memory (temporary or non-temporary memory reservation), this
-  /// holds the memory and will release it when we are destroyed
-  GpuMemoryReservation reservation_;
+   private:
+    /// If we own the memory (temporary or non-temporary memory reservation),
+    /// this holds the memory and will release it when we are destroyed
+    GpuMemoryReservation reservation_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss
 
 #include <faiss/gpu/utils/DeviceTensor-inl.cuh>

faiss/gpu/utils/DeviceVector.cuh

@@ -5,18 +5,18 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/gpu/utils/DeviceUtils.h>
-#include <faiss/gpu/GpuResources.h>
-#include <faiss/gpu/utils/StaticUtils.h>
-#include <algorithm>
 #include <cuda.h>
+#include <faiss/gpu/GpuResources.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/gpu/utils/StaticUtils.h>
+#include <faiss/impl/FaissAssert.h>
+#include <algorithm>
 #include <vector>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// A simple version of thrust::device_vector<T>, but has more control
 /// over whether resize() initializes new space with T() (which we
@@ -27,181 +27,203 @@ namespace faiss { namespace gpu {
 /// the user.
 template <typename T>
 class DeviceVector {
- public:
-  DeviceVector(GpuResources* res, AllocInfo allocInfo)
-      : num_(0),
-        capacity_(0),
-        res_(res),
-        allocInfo_(allocInfo) {
-    FAISS_ASSERT(res_);
-  }
-
-  ~DeviceVector() {
-    clear();
-  }
-
-  // Clear all allocated memory; reset to zero size
-  void clear() {
-    alloc_.release();
-    num_ = 0;
-    capacity_ = 0;
-  }
-
-  size_t size() const { return num_; }
-  size_t capacity() const { return capacity_; }
-  T* data() { return (T*) alloc_.data; }
-  const T* data() const { return (const T*) alloc_.data; }
-
-  template <typename OutT>
-  std::vector<OutT> copyToHost(cudaStream_t stream) const {
-    FAISS_ASSERT(num_ * sizeof(T) % sizeof(OutT) == 0);
-
-    std::vector<OutT> out((num_ * sizeof(T)) / sizeof(OutT));
-
-    if (num_ > 0) {
-      FAISS_ASSERT(data());
-      CUDA_VERIFY(cudaMemcpyAsync(out.data(), data(), num_ * sizeof(T),
-                                  cudaMemcpyDeviceToHost, stream));
+   public:
+    DeviceVector(GpuResources* res, AllocInfo allocInfo)
+            : num_(0), capacity_(0), res_(res), allocInfo_(allocInfo) {
+        FAISS_ASSERT(res_);
     }
 
-    return out;
-  }
-
-  // Returns true if we actually reallocated memory
-  // If `reserveExact` is true, then we reserve only the memory that
-  // we need for what we're appending
-  bool append(const T* d,
-              size_t n,
-              cudaStream_t stream,
-              bool reserveExact = false) {
-    bool mem = false;
-
-    if (n > 0) {
-      size_t reserveSize = num_ + n;
-      if (!reserveExact) {
-        reserveSize = getNewCapacity_(reserveSize);
-      }
-
-      mem = reserve(reserveSize, stream);
-
-      int dev = getDeviceForAddress(d);
-      if (dev == -1) {
-        CUDA_VERIFY(cudaMemcpyAsync(data() + num_, d, n * sizeof(T),
-                                    cudaMemcpyHostToDevice, stream));
-      } else {
-        CUDA_VERIFY(cudaMemcpyAsync(data() + num_, d, n * sizeof(T),
-                                    cudaMemcpyDeviceToDevice, stream));
-      }
-      num_ += n;
+    ~DeviceVector() {
+        clear();
     }
 
-    return mem;
-  }
-
-  // Returns true if we actually reallocated memory
-  bool resize(size_t newSize, cudaStream_t stream) {
-    bool mem = false;
-
-    if (num_ < newSize) {
-      mem = reserve(getNewCapacity_(newSize), stream);
+    // Clear all allocated memory; reset to zero size
+    void clear() {
+        alloc_.release();
+        num_ = 0;
+        capacity_ = 0;
     }
 
-    // Don't bother zero initializing the newly accessible memory
-    // (unlike thrust::device_vector)
-    num_ = newSize;
-
-    return mem;
-  }
-
-  // Clean up after oversized allocations, while leaving some space to
-  // remain for subsequent allocations (if `exact` false) or to
-  // exactly the space we need (if `exact` true); returns space
-  // reclaimed in bytes
-  size_t reclaim(bool exact, cudaStream_t stream) {
-    size_t free = capacity_ - num_;
-
-    if (exact) {
-      realloc_(num_, stream);
-      return free * sizeof(T);
+    size_t size() const {
+        return num_;
+    }
+    size_t capacity() const {
+        return capacity_;
+    }
+    T* data() {
+        return (T*)alloc_.data;
+    }
+    const T* data() const {
+        return (const T*)alloc_.data;
     }
 
-    // If more than 1/4th of the space is free, then we want to
-    // truncate to only having 1/8th of the space free; this still
-    // preserves some space for new elements, but won't force us to
-    // double our size right away
-    if (free > (capacity_ / 4)) {
-      size_t newFree = capacity_ / 8;
-      size_t newCapacity = num_ + newFree;
+    template <typename OutT>
+    std::vector<OutT> copyToHost(cudaStream_t stream) const {
+        FAISS_ASSERT(num_ * sizeof(T) % sizeof(OutT) == 0);
 
-      size_t oldCapacity = capacity_;
-      FAISS_ASSERT(newCapacity < oldCapacity);
+        std::vector<OutT> out((num_ * sizeof(T)) / sizeof(OutT));
 
-      realloc_(newCapacity, stream);
+        if (num_ > 0) {
+            FAISS_ASSERT(data());
+            CUDA_VERIFY(cudaMemcpyAsync(
+                    out.data(),
+                    data(),
+                    num_ * sizeof(T),
+                    cudaMemcpyDeviceToHost,
+                    stream));
+        }
 
-      return (oldCapacity - newCapacity) * sizeof(T);
+        return out;
     }
 
-    return 0;
-  }
+    // Returns true if we actually reallocated memory
+    // If `reserveExact` is true, then we reserve only the memory that
+    // we need for what we're appending
+    bool append(
+            const T* d,
+            size_t n,
+            cudaStream_t stream,
+            bool reserveExact = false) {
+        bool mem = false;
 
-  // Returns true if we actually reallocated memory
-  bool reserve(size_t newCapacity, cudaStream_t stream) {
-    if (newCapacity <= capacity_) {
-      return false;
+        if (n > 0) {
+            size_t reserveSize = num_ + n;
+            if (!reserveExact) {
+                reserveSize = getNewCapacity_(reserveSize);
+            }
+
+            mem = reserve(reserveSize, stream);
+
+            int dev = getDeviceForAddress(d);
+            if (dev == -1) {
+                CUDA_VERIFY(cudaMemcpyAsync(
+                        data() + num_,
+                        d,
+                        n * sizeof(T),
+                        cudaMemcpyHostToDevice,
+                        stream));
+            } else {
+                CUDA_VERIFY(cudaMemcpyAsync(
+                        data() + num_,
+                        d,
+                        n * sizeof(T),
+                        cudaMemcpyDeviceToDevice,
+                        stream));
+            }
+            num_ += n;
+        }
+
+        return mem;
     }
 
-    // Otherwise, we need new space.
-    realloc_(newCapacity, stream);
-    return true;
-  }
+    // Returns true if we actually reallocated memory
+    bool resize(size_t newSize, cudaStream_t stream) {
+        bool mem = false;
 
- private:
-  void realloc_(size_t newCapacity, cudaStream_t stream) {
-    FAISS_ASSERT(num_ <= newCapacity);
+        if (num_ < newSize) {
+            mem = reserve(getNewCapacity_(newSize), stream);
+        }
 
-    size_t newSizeInBytes = newCapacity * sizeof(T);
-    size_t oldSizeInBytes = num_ * sizeof(T);
+        // Don't bother zero initializing the newly accessible memory
+        // (unlike thrust::device_vector)
+        num_ = newSize;
 
-    // The new allocation will occur on this stream
-    allocInfo_.stream = stream;
+        return mem;
+    }
 
-    auto newAlloc =
-      res_->allocMemoryHandle(AllocRequest(allocInfo_, newSizeInBytes));
+    // Clean up after oversized allocations, while leaving some space to
+    // remain for subsequent allocations (if `exact` false) or to
+    // exactly the space we need (if `exact` true); returns space
+    // reclaimed in bytes
+    size_t reclaim(bool exact, cudaStream_t stream) {
+        size_t free = capacity_ - num_;
 
-    // Copy over any old data
-    CUDA_VERIFY(cudaMemcpyAsync(newAlloc.data,
-                                data(),
-                                oldSizeInBytes,
-                                cudaMemcpyDeviceToDevice, stream));
+        if (exact) {
+            realloc_(num_, stream);
+            return free * sizeof(T);
+        }
 
-    // Zero out the new space past the data we just copied
-    CUDA_VERIFY(cudaMemsetAsync((uint8_t*) newAlloc.data + oldSizeInBytes,
-                                0,
-                                newSizeInBytes - oldSizeInBytes,
-                                stream));
+        // If more than 1/4th of the space is free, then we want to
+        // truncate to only having 1/8th of the space free; this still
+        // preserves some space for new elements, but won't force us to
+        // double our size right away
+        if (free > (capacity_ / 4)) {
+            size_t newFree = capacity_ / 8;
+            size_t newCapacity = num_ + newFree;
 
-    alloc_ = std::move(newAlloc);
-    capacity_ = newCapacity;
-  }
+            size_t oldCapacity = capacity_;
+            FAISS_ASSERT(newCapacity < oldCapacity);
 
-  size_t getNewCapacity_(size_t preferredSize) {
-    return utils::nextHighestPowerOf2(preferredSize);
-  }
+            realloc_(newCapacity, stream);
 
-  /// Our current memory allocation, if any
-  GpuMemoryReservation alloc_;
+            return (oldCapacity - newCapacity) * sizeof(T);
+        }
 
-  /// current valid number of T present
-  size_t num_;
+        return 0;
+    }
 
-  /// current space of T present (bytes is sizeof(T) * capacity_)
-  size_t capacity_;
+    // Returns true if we actually reallocated memory
+    bool reserve(size_t newCapacity, cudaStream_t stream) {
+        if (newCapacity <= capacity_) {
+            return false;
+        }
 
-  /// Where we should allocate memory
-  GpuResources* res_;
+        // Otherwise, we need new space.
+        realloc_(newCapacity, stream);
+        return true;
+    }
 
-  /// How we should allocate memory
-  AllocInfo allocInfo_;
+   private:
+    void realloc_(size_t newCapacity, cudaStream_t stream) {
+        FAISS_ASSERT(num_ <= newCapacity);
+
+        size_t newSizeInBytes = newCapacity * sizeof(T);
+        size_t oldSizeInBytes = num_ * sizeof(T);
+
+        // The new allocation will occur on this stream
+        allocInfo_.stream = stream;
+
+        auto newAlloc = res_->allocMemoryHandle(
+                AllocRequest(allocInfo_, newSizeInBytes));
+
+        // Copy over any old data
+        CUDA_VERIFY(cudaMemcpyAsync(
+                newAlloc.data,
+                data(),
+                oldSizeInBytes,
+                cudaMemcpyDeviceToDevice,
+                stream));
+
+        // Zero out the new space past the data we just copied
+        CUDA_VERIFY(cudaMemsetAsync(
+                (uint8_t*)newAlloc.data + oldSizeInBytes,
+                0,
+                newSizeInBytes - oldSizeInBytes,
+                stream));
+
+        alloc_ = std::move(newAlloc);
+        capacity_ = newCapacity;
+    }
+
+    size_t getNewCapacity_(size_t preferredSize) {
+        return utils::nextHighestPowerOf2(preferredSize);
+    }
+
+    /// Our current memory allocation, if any
+    GpuMemoryReservation alloc_;
+
+    /// current valid number of T present
+    size_t num_;
+
+    /// current space of T present (bytes is sizeof(T) * capacity_)
+    size_t capacity_;
+
+    /// Where we should allocate memory
+    GpuResources* res_;
+
+    /// How we should allocate memory
+    AllocInfo allocInfo_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/Float16.cuh

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cuda.h>
@@ -19,55 +18,56 @@
 
 #include <cuda_fp16.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // 64 bytes containing 4 half (float16) values
 struct Half4 {
-  half2 a;
-  half2 b;
+    half2 a;
+    half2 b;
 };
 
 inline __device__ float4 half4ToFloat4(Half4 v) {
-  float2 a = __half22float2(v.a);
-  float2 b = __half22float2(v.b);
+    float2 a = __half22float2(v.a);
+    float2 b = __half22float2(v.b);
 
-  float4 out;
-  out.x = a.x;
-  out.y = a.y;
-  out.z = b.x;
-  out.w = b.y;
+    float4 out;
+    out.x = a.x;
+    out.y = a.y;
+    out.z = b.x;
+    out.w = b.y;
 
-  return out;
+    return out;
 }
 
 inline __device__ Half4 float4ToHalf4(float4 v) {
-  float2 a;
-  a.x = v.x;
-  a.y = v.y;
+    float2 a;
+    a.x = v.x;
+    a.y = v.y;
 
-  float2 b;
-  b.x = v.z;
-  b.y = v.w;
+    float2 b;
+    b.x = v.z;
+    b.y = v.w;
 
-  Half4 out;
-  out.a = __float22half2_rn(a);
-  out.b = __float22half2_rn(b);
+    Half4 out;
+    out.a = __float22half2_rn(a);
+    out.b = __float22half2_rn(b);
 
-  return out;
+    return out;
 }
 
 // 128 bytes containing 8 half (float16) values
 struct Half8 {
-  Half4 a;
-  Half4 b;
+    Half4 a;
+    Half4 b;
 };
 
 /// Returns true if the given device supports native float16 math
 inline bool getDeviceSupportsFloat16Math(int device) {
-  const auto& prop = getDeviceProperties(device);
+    const auto& prop = getDeviceProperties(device);
 
-  return (prop.major >= 6 ||
-          (prop.major == 5 && prop.minor >= 3));
+    return (prop.major >= 6 || (prop.major == 5 && prop.minor >= 3));
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/HostTensor-inl.cuh

@@ -5,176 +5,233 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+namespace faiss {
+namespace gpu {
 
-namespace faiss { namespace gpu {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor()
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(),
+          state_(AllocState::NotOwner) {}
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor() :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(),
-    state_(AllocState::NotOwner) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::~HostTensor() {
+    if (state_ == AllocState::Owner) {
+        FAISS_ASSERT(this->data_ != nullptr);
+        delete[] this->data_;
+        this->data_ = nullptr;
+    }
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::~HostTensor() {
-  if (state_ == AllocState::Owner) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
+        HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(),
+          state_(AllocState::NotOwner) {
+    this->operator=(std::move(t));
+}
+
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& HostTensor<
+        T,
+        Dim,
+        InnerContig,
+        IndexT,
+        PtrTraits>::
+operator=(HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) {
+    if (this->state_ == AllocState::Owner) {
+        FAISS_ASSERT(this->data_ != nullptr);
+        delete[] this->data_;
+        this->data_ = nullptr;
+    }
+
+    this->Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
+            std::move(t));
+
+    this->state_ = t.state_;
+    t.state_ = AllocState::NotOwner;
+
+    return *this;
+}
+
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
+        const IndexT sizes[Dim])
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
+          state_(AllocState::Owner) {
+    this->data_ = new T[this->numElements()];
     FAISS_ASSERT(this->data_ != nullptr);
-    delete[] this->data_;
-    this->data_ = nullptr;
-  }
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
-  HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(),
-    state_(AllocState::NotOwner) {
-  this->operator=(std::move(t));
-}
-
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
-  HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) {
-  if (this->state_ == AllocState::Owner) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
+        std::initializer_list<IndexT> sizes)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
+          state_(AllocState::Owner) {
+    this->data_ = new T[this->numElements()];
     FAISS_ASSERT(this->data_ != nullptr);
-    delete[] this->data_;
-    this->data_ = nullptr;
-  }
-
-  this->Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
-    std::move(t));
-
-  this->state_ = t.state_; t.state_ = AllocState::NotOwner;
-
-  return *this;
 }
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
-  const IndexT sizes[Dim]) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
-    state_(AllocState::Owner) {
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
+        DataPtrType data,
+        const IndexT sizes[Dim])
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes),
+          state_(AllocState::NotOwner) {}
 
-  this->data_ = new T[this->numElements()];
-  FAISS_ASSERT(this->data_ != nullptr);
-}
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
+        DataPtrType data,
+        std::initializer_list<IndexT> sizes)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes),
+          state_(AllocState::NotOwner) {}
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
-  std::initializer_list<IndexT> sizes) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
-    state_(AllocState::Owner) {
-  this->data_ = new T[this->numElements()];
-  FAISS_ASSERT(this->data_ != nullptr);
-}
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
+        DataPtrType data,
+        const IndexT sizes[Dim],
+        const IndexT strides[Dim])
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes, strides),
+          state_(AllocState::NotOwner) {}
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
-  DataPtrType data,
-  const IndexT sizes[Dim]) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes),
-    state_(AllocState::NotOwner) {
-}
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
+        const Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
+        cudaStream_t stream)
+        : Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(
+                  nullptr,
+                  t.sizes(),
+                  t.strides()),
+          state_(AllocState::Owner) {
+    // Only contiguous arrays handled for now
+    FAISS_ASSERT(t.isContiguous());
 
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
-  DataPtrType data,
-  std::initializer_list<IndexT> sizes) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes),
-    state_(AllocState::NotOwner) {
-}
-
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
-  DataPtrType data,
-  const IndexT sizes[Dim],
-  const IndexT strides[Dim]) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes, strides),
-    state_(AllocState::NotOwner) {
-}
-
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::HostTensor(
-  const Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
-  cudaStream_t stream) :
-    Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, t.sizes(), t.strides()),
-    state_(AllocState::Owner) {
-  // Only contiguous arrays handled for now
-  FAISS_ASSERT(t.isContiguous());
-
-  this->data_ = new T[t.numElements()];
-  this->copyFrom(t, stream);
+    this->data_ = new T[t.numElements()];
+    this->copyFrom(t, stream);
 }
 
 /// Call to zero out memory
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::zero() {
-  // Region must be contiguous
-  FAISS_ASSERT(this->isContiguous());
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& HostTensor<
+        T,
+        Dim,
+        InnerContig,
+        IndexT,
+        PtrTraits>::zero() {
+    // Region must be contiguous
+    FAISS_ASSERT(this->isContiguous());
 
-  if (this->data_ != nullptr) {
-    memset(this->data_, 0, this->getSizeInBytes());
-  }
-
-  return *this;
-}
-
-template <typename T, int Dim, bool InnerContig,
-          typename IndexT, template <typename U> class PtrTraits>
-__host__ T
-HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::maxDiff(
-  const HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) const {
-  auto size = this->numElements();
-
-  FAISS_ASSERT(size == t.numElements());
-  FAISS_ASSERT(size > 0);
-
-  if (InnerContig) {
-    auto a = this->data();
-    auto b = t.data();
-
-    T maxDiff = a[0] - b[0];
-    // FIXME: type-specific abs()
-    maxDiff = maxDiff < 0 ? maxDiff * (T) -1 : maxDiff;
-
-    for (IndexT i = 1; i < size; ++i) {
-      auto diff = a[i] - b[i];
-      // FIXME: type-specific abs
-      diff = diff < 0 ? diff * (T) -1 : diff;
-      if (diff > maxDiff) {
-        maxDiff = diff;
-      }
+    if (this->data_ != nullptr) {
+        memset(this->data_, 0, this->getSizeInBytes());
     }
 
-    return maxDiff;
-  } else {
-    // non-contiguous
-    // FIXME
-    FAISS_ASSERT(false);
-    return (T) 0;
-  }
+    return *this;
 }
 
-} } // namespace
+template <
+        typename T,
+        int Dim,
+        bool InnerContig,
+        typename IndexT,
+        template <typename U>
+        class PtrTraits>
+__host__ T HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>::maxDiff(
+        const HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) const {
+    auto size = this->numElements();
+
+    FAISS_ASSERT(size == t.numElements());
+    FAISS_ASSERT(size > 0);
+
+    if (InnerContig) {
+        auto a = this->data();
+        auto b = t.data();
+
+        T maxDiff = a[0] - b[0];
+        // FIXME: type-specific abs()
+        maxDiff = maxDiff < 0 ? maxDiff * (T)-1 : maxDiff;
+
+        for (IndexT i = 1; i < size; ++i) {
+            auto diff = a[i] - b[i];
+            // FIXME: type-specific abs
+            diff = diff < 0 ? diff * (T)-1 : diff;
+            if (diff > maxDiff) {
+                maxDiff = diff;
+            }
+        }
+
+        return maxDiff;
+    } else {
+        // non-contiguous
+        // FIXME
+        FAISS_ASSERT(false);
+        return (T)0;
+    }
+}
+
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/HostTensor.cuh

@@ -5,87 +5,89 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
-template <typename T,
-          int Dim,
-          bool InnerContig = false,
-          typename IndexT = int,
-          template <typename U> class PtrTraits = traits::DefaultPtrTraits>
+template <
+        typename T,
+        int Dim,
+        bool InnerContig = false,
+        typename IndexT = int,
+        template <typename U> class PtrTraits = traits::DefaultPtrTraits>
 class HostTensor : public Tensor<T, Dim, InnerContig, IndexT, PtrTraits> {
- public:
-  typedef IndexT IndexType;
-  typedef typename PtrTraits<T>::PtrType DataPtrType;
+   public:
+    typedef IndexT IndexType;
+    typedef typename PtrTraits<T>::PtrType DataPtrType;
 
-  /// Default constructor
-  __host__ HostTensor();
+    /// Default constructor
+    __host__ HostTensor();
 
-  /// Destructor
-  __host__ ~HostTensor();
+    /// Destructor
+    __host__ ~HostTensor();
 
-  /// Move constructor
-  __host__ HostTensor(HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
+    /// Move constructor
+    __host__ HostTensor(HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
 
-  /// Move assignment
-  __host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
-  operator=(HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
+    /// Move assignment
+    __host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& operator=(
+            HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
 
-  /// Constructs a tensor of the given size, allocating memory for it
-  /// locally
-  __host__ HostTensor(const IndexT sizes[Dim]);
-  __host__ HostTensor(std::initializer_list<IndexT> sizes);
+    /// Constructs a tensor of the given size, allocating memory for it
+    /// locally
+    __host__ HostTensor(const IndexT sizes[Dim]);
+    __host__ HostTensor(std::initializer_list<IndexT> sizes);
 
-  /// Constructs a tensor of the given size and stride, referencing a
-  /// memory region we do not own
-  __host__ HostTensor(DataPtrType data,
-                      const IndexT sizes[Dim]);
-  __host__ HostTensor(DataPtrType data,
-                      std::initializer_list<IndexT> sizes);
+    /// Constructs a tensor of the given size and stride, referencing a
+    /// memory region we do not own
+    __host__ HostTensor(DataPtrType data, const IndexT sizes[Dim]);
+    __host__ HostTensor(DataPtrType data, std::initializer_list<IndexT> sizes);
 
-  /// Constructs a tensor of the given size and stride, referencing a
-  /// memory region we do not own
-  __host__ HostTensor(DataPtrType data,
-                      const IndexT sizes[Dim],
-                      const IndexT strides[Dim]);
+    /// Constructs a tensor of the given size and stride, referencing a
+    /// memory region we do not own
+    __host__ HostTensor(
+            DataPtrType data,
+            const IndexT sizes[Dim],
+            const IndexT strides[Dim]);
 
-  /// Copies a tensor into ourselves, allocating memory for it
-  /// locally. If the tensor is on the GPU, then we will copy it to
-  /// ourselves wrt the given stream.
-  __host__ HostTensor(const Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
-                      cudaStream_t stream);
+    /// Copies a tensor into ourselves, allocating memory for it
+    /// locally. If the tensor is on the GPU, then we will copy it to
+    /// ourselves wrt the given stream.
+    __host__ HostTensor(
+            const Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
+            cudaStream_t stream);
 
-  /// Call to zero out memory
-  __host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& zero();
+    /// Call to zero out memory
+    __host__ HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& zero();
 
-  /// Returns the maximum difference seen between two tensors
-  __host__ T
-  maxDiff(const HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) const;
+    /// Returns the maximum difference seen between two tensors
+    __host__ T
+    maxDiff(const HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) const;
 
-  /// Are the two tensors exactly equal?
-  __host__ bool
-  equal(const HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) const {
-    return (maxDiff(t) == (T) 0);
-  }
+    /// Are the two tensors exactly equal?
+    __host__ bool equal(
+            const HostTensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) const {
+        return (maxDiff(t) == (T)0);
+    }
 
- private:
-  enum AllocState {
-    /// This tensor itself owns the memory, which must be freed via
-    /// cudaFree
-    Owner,
+   private:
+    enum AllocState {
+        /// This tensor itself owns the memory, which must be freed via
+        /// cudaFree
+        Owner,
 
-    /// This tensor itself is not an owner of the memory; there is
-    /// nothing to free
-    NotOwner,
-  };
+        /// This tensor itself is not an owner of the memory; there is
+        /// nothing to free
+        NotOwner,
+    };
 
-  AllocState state_;
+    AllocState state_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss
 
 #include <faiss/gpu/utils/HostTensor-inl.cuh>

faiss/gpu/utils/Limits.cuh

@@ -5,17 +5,16 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Pair.cuh>
 #include <limits>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T>
-struct Limits {
-};
+struct Limits {};
 
 // Unfortunately we can't use constexpr because there is no
 // constexpr constructor for half
@@ -25,34 +24,34 @@ constexpr float kFloatMin = std::numeric_limits<float>::lowest();
 
 template <>
 struct Limits<float> {
-  static __device__ __host__ inline float getMin() {
-    return kFloatMin;
-  }
-  static __device__ __host__ inline float getMax() {
-    return kFloatMax;
-  }
+    static __device__ __host__ inline float getMin() {
+        return kFloatMin;
+    }
+    static __device__ __host__ inline float getMax() {
+        return kFloatMax;
+    }
 };
 
 inline __device__ __host__ half kGetHalf(unsigned short v) {
 #if CUDA_VERSION >= 9000
-  __half_raw h;
-  h.x = v;
-  return __half(h);
+    __half_raw h;
+    h.x = v;
+    return __half(h);
 #else
-  half h;
-  h.x = v;
-  return h;
+    half h;
+    h.x = v;
+    return h;
 #endif
 }
 
 template <>
 struct Limits<half> {
-  static __device__ __host__ inline half getMin() {
-    return kGetHalf(0xfbffU);
-  }
-  static __device__ __host__ inline half getMax() {
-    return kGetHalf(0x7bffU);
-  }
+    static __device__ __host__ inline half getMin() {
+        return kGetHalf(0xfbffU);
+    }
+    static __device__ __host__ inline half getMax() {
+        return kGetHalf(0x7bffU);
+    }
 };
 
 constexpr int kIntMax = std::numeric_limits<int>::max();
@@ -60,23 +59,24 @@ constexpr int kIntMin = std::numeric_limits<int>::lowest();
 
 template <>
 struct Limits<int> {
-  static __device__ __host__ inline int getMin() {
-    return kIntMin;
-  }
-  static __device__ __host__ inline int getMax() {
-    return kIntMax;
-  }
+    static __device__ __host__ inline int getMin() {
+        return kIntMin;
+    }
+    static __device__ __host__ inline int getMax() {
+        return kIntMax;
+    }
 };
 
-template<typename K, typename V>
+template <typename K, typename V>
 struct Limits<Pair<K, V>> {
-  static __device__ __host__ inline Pair<K, V> getMin() {
-    return Pair<K, V>(Limits<K>::getMin(), Limits<V>::getMin());
-  }
+    static __device__ __host__ inline Pair<K, V> getMin() {
+        return Pair<K, V>(Limits<K>::getMin(), Limits<V>::getMin());
+    }
 
-  static __device__ __host__ inline Pair<K, V> getMax() {
-    return Pair<K, V>(Limits<K>::getMax(), Limits<V>::getMax());
-  }
+    static __device__ __host__ inline Pair<K, V> getMax() {
+        return Pair<K, V>(Limits<K>::getMax(), Limits<V>::getMax());
+    }
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/LoadStoreOperators.cuh

@@ -5,86 +5,97 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/Float16.cuh>
 
 #ifndef __HALF2_TO_UI
 // cuda_fp16.hpp doesn't export this
-#define __HALF2_TO_UI(var) *(reinterpret_cast<unsigned int *>(&(var)))
+#define __HALF2_TO_UI(var) *(reinterpret_cast<unsigned int*>(&(var)))
 #endif
 
-
 //
 // Templated wrappers to express load/store for different scalar and vector
 // types, so kernels can have the same written form but can operate
 // over half and float, and on vector types transparently
 //
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T>
 struct LoadStore {
-  static inline __device__ T load(void* p) {
-    return *((T*) p);
-  }
+    static inline __device__ T load(void* p) {
+        return *((T*)p);
+    }
 
-  static inline __device__ void store(void* p, const T& v) {
-    *((T*) p) = v;
-  }
+    static inline __device__ void store(void* p, const T& v) {
+        *((T*)p) = v;
+    }
 };
 
 template <>
 struct LoadStore<Half4> {
-  static inline __device__ Half4 load(void* p) {
-    Half4 out;
+    static inline __device__ Half4 load(void* p) {
+        Half4 out;
 #if CUDA_VERSION >= 9000
-    asm("ld.global.v2.u32 {%0, %1}, [%2];" :
-        "=r"(__HALF2_TO_UI(out.a)), "=r"(__HALF2_TO_UI(out.b)) : "l"(p));
+        asm("ld.global.v2.u32 {%0, %1}, [%2];"
+            : "=r"(__HALF2_TO_UI(out.a)), "=r"(__HALF2_TO_UI(out.b))
+            : "l"(p));
 #else
-    asm("ld.global.v2.u32 {%0, %1}, [%2];" :
-        "=r"(out.a.x), "=r"(out.b.x) : "l"(p));
+        asm("ld.global.v2.u32 {%0, %1}, [%2];"
+            : "=r"(out.a.x), "=r"(out.b.x)
+            : "l"(p));
 #endif
-    return out;
-  }
+        return out;
+    }
 
-  static inline __device__ void store(void* p, Half4& v) {
+    static inline __device__ void store(void* p, Half4& v) {
 #if CUDA_VERSION >= 9000
-    asm("st.v2.u32 [%0], {%1, %2};" : : "l"(p),
-        "r"(__HALF2_TO_UI(v.a)), "r"(__HALF2_TO_UI(v.b)));
+        asm("st.v2.u32 [%0], {%1, %2};"
+            :
+            : "l"(p), "r"(__HALF2_TO_UI(v.a)), "r"(__HALF2_TO_UI(v.b)));
 #else
-    asm("st.v2.u32 [%0], {%1, %2};" : : "l"(p), "r"(v.a.x), "r"(v.b.x));
+        asm("st.v2.u32 [%0], {%1, %2};" : : "l"(p), "r"(v.a.x), "r"(v.b.x));
 #endif
-  }
+    }
 };
 
 template <>
 struct LoadStore<Half8> {
-  static inline __device__ Half8 load(void* p) {
-    Half8 out;
+    static inline __device__ Half8 load(void* p) {
+        Half8 out;
 #if CUDA_VERSION >= 9000
-    asm("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];" :
-        "=r"(__HALF2_TO_UI(out.a.a)), "=r"(__HALF2_TO_UI(out.a.b)),
-        "=r"(__HALF2_TO_UI(out.b.a)), "=r"(__HALF2_TO_UI(out.b.b)) : "l"(p));
+        asm("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+            : "=r"(__HALF2_TO_UI(out.a.a)),
+              "=r"(__HALF2_TO_UI(out.a.b)),
+              "=r"(__HALF2_TO_UI(out.b.a)),
+              "=r"(__HALF2_TO_UI(out.b.b))
+            : "l"(p));
 #else
-    asm("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];" :
-        "=r"(out.a.a.x), "=r"(out.a.b.x),
-        "=r"(out.b.a.x), "=r"(out.b.b.x) : "l"(p));
+        asm("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+            : "=r"(out.a.a.x), "=r"(out.a.b.x), "=r"(out.b.a.x), "=r"(out.b.b.x)
+            : "l"(p));
 #endif
-    return out;
-  }
+        return out;
+    }
 
-  static inline __device__ void store(void* p, Half8& v) {
+    static inline __device__ void store(void* p, Half8& v) {
 #if CUDA_VERSION >= 9000
-    asm("st.v4.u32 [%0], {%1, %2, %3, %4};"
-        : : "l"(p), "r"(__HALF2_TO_UI(v.a.a)), "r"(__HALF2_TO_UI(v.a.b)),
-          "r"(__HALF2_TO_UI(v.b.a)), "r"(__HALF2_TO_UI(v.b.b)));
+        asm("st.v4.u32 [%0], {%1, %2, %3, %4};"
+            :
+            : "l"(p),
+              "r"(__HALF2_TO_UI(v.a.a)),
+              "r"(__HALF2_TO_UI(v.a.b)),
+              "r"(__HALF2_TO_UI(v.b.a)),
+              "r"(__HALF2_TO_UI(v.b.b)));
 #else
-    asm("st.v4.u32 [%0], {%1, %2, %3, %4};"
-        : : "l"(p), "r"(v.a.a.x), "r"(v.a.b.x), "r"(v.b.a.x), "r"(v.b.b.x));
+        asm("st.v4.u32 [%0], {%1, %2, %3, %4};"
+            :
+            : "l"(p), "r"(v.a.a.x), "r"(v.a.b.x), "r"(v.b.a.x), "r"(v.b.b.x));
 #endif
-  }
+    }
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/MathOperators.cuh

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/ConversionOperators.cuh>
@@ -17,543 +16,545 @@
 // over half and float, and on vector types transparently
 //
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T>
 struct Math {
-  typedef T ScalarType;
+    typedef T ScalarType;
 
-  static inline __device__ T add(T a, T b) {
-    return a + b;
-  }
+    static inline __device__ T add(T a, T b) {
+        return a + b;
+    }
 
-  static inline __device__ T sub(T a, T b) {
-    return a - b;
-  }
+    static inline __device__ T sub(T a, T b) {
+        return a - b;
+    }
 
-  static inline __device__ T mul(T a, T b) {
-    return a * b;
-  }
+    static inline __device__ T mul(T a, T b) {
+        return a * b;
+    }
 
-  static inline __device__ T neg(T v) {
-    return -v;
-  }
+    static inline __device__ T neg(T v) {
+        return -v;
+    }
 
-  /// For a vector type, this is a horizontal add, returning sum(v_i)
-  static inline __device__ float reduceAdd(T v) {
-    return ConvertTo<float>::to(v);
-  }
+    /// For a vector type, this is a horizontal add, returning sum(v_i)
+    static inline __device__ float reduceAdd(T v) {
+        return ConvertTo<float>::to(v);
+    }
 
-  static inline __device__ bool lt(T a, T b) {
-    return a < b;
-  }
+    static inline __device__ bool lt(T a, T b) {
+        return a < b;
+    }
 
-  static inline __device__ bool gt(T a, T b) {
-    return a > b;
-  }
+    static inline __device__ bool gt(T a, T b) {
+        return a > b;
+    }
 
-  static inline __device__ bool eq(T a, T b) {
-    return a == b;
-  }
+    static inline __device__ bool eq(T a, T b) {
+        return a == b;
+    }
 
-  static inline __device__ T zero() {
-    return (T) 0;
-  }
+    static inline __device__ T zero() {
+        return (T)0;
+    }
 };
 
 template <>
 struct Math<float2> {
-  typedef float ScalarType;
+    typedef float ScalarType;
 
-  static inline __device__ float2 add(float2 a, float2 b) {
-    float2 v;
-    v.x = a.x + b.x;
-    v.y = a.y + b.y;
-    return v;
-  }
+    static inline __device__ float2 add(float2 a, float2 b) {
+        float2 v;
+        v.x = a.x + b.x;
+        v.y = a.y + b.y;
+        return v;
+    }
 
-  static inline __device__ float2 sub(float2 a, float2 b) {
-    float2 v;
-    v.x = a.x - b.x;
-    v.y = a.y - b.y;
-    return v;
-  }
+    static inline __device__ float2 sub(float2 a, float2 b) {
+        float2 v;
+        v.x = a.x - b.x;
+        v.y = a.y - b.y;
+        return v;
+    }
 
-  static inline __device__ float2 add(float2 a, float b) {
-    float2 v;
-    v.x = a.x + b;
-    v.y = a.y + b;
-    return v;
-  }
+    static inline __device__ float2 add(float2 a, float b) {
+        float2 v;
+        v.x = a.x + b;
+        v.y = a.y + b;
+        return v;
+    }
 
-  static inline __device__ float2 sub(float2 a, float b) {
-    float2 v;
-    v.x = a.x - b;
-    v.y = a.y - b;
-    return v;
-  }
+    static inline __device__ float2 sub(float2 a, float b) {
+        float2 v;
+        v.x = a.x - b;
+        v.y = a.y - b;
+        return v;
+    }
 
-  static inline __device__ float2 mul(float2 a, float2 b) {
-    float2 v;
-    v.x = a.x * b.x;
-    v.y = a.y * b.y;
-    return v;
-  }
+    static inline __device__ float2 mul(float2 a, float2 b) {
+        float2 v;
+        v.x = a.x * b.x;
+        v.y = a.y * b.y;
+        return v;
+    }
 
-  static inline __device__ float2 mul(float2 a, float b) {
-    float2 v;
-    v.x = a.x * b;
-    v.y = a.y * b;
-    return v;
-  }
+    static inline __device__ float2 mul(float2 a, float b) {
+        float2 v;
+        v.x = a.x * b;
+        v.y = a.y * b;
+        return v;
+    }
 
-  static inline __device__ float2 neg(float2 v) {
-    v.x = -v.x;
-    v.y = -v.y;
-    return v;
-  }
+    static inline __device__ float2 neg(float2 v) {
+        v.x = -v.x;
+        v.y = -v.y;
+        return v;
+    }
 
-  /// For a vector type, this is a horizontal add, returning sum(v_i)
-  static inline __device__ float reduceAdd(float2 v) {
-    return v.x + v.y;
-  }
+    /// For a vector type, this is a horizontal add, returning sum(v_i)
+    static inline __device__ float reduceAdd(float2 v) {
+        return v.x + v.y;
+    }
 
-  // not implemented for vector types
-  // static inline __device__ bool lt(float2 a, float2 b);
-  // static inline __device__ bool gt(float2 a, float2 b);
-  // static inline __device__ bool eq(float2 a, float2 b);
+    // not implemented for vector types
+    // static inline __device__ bool lt(float2 a, float2 b);
+    // static inline __device__ bool gt(float2 a, float2 b);
+    // static inline __device__ bool eq(float2 a, float2 b);
 
-  static inline __device__ float2 zero() {
-    float2 v;
-    v.x = 0.0f;
-    v.y = 0.0f;
-    return v;
-  }
+    static inline __device__ float2 zero() {
+        float2 v;
+        v.x = 0.0f;
+        v.y = 0.0f;
+        return v;
+    }
 };
 
 template <>
 struct Math<float4> {
-  typedef float ScalarType;
+    typedef float ScalarType;
 
-  static inline __device__ float4 add(float4 a, float4 b) {
-    float4 v;
-    v.x = a.x + b.x;
-    v.y = a.y + b.y;
-    v.z = a.z + b.z;
-    v.w = a.w + b.w;
-    return v;
-  }
+    static inline __device__ float4 add(float4 a, float4 b) {
+        float4 v;
+        v.x = a.x + b.x;
+        v.y = a.y + b.y;
+        v.z = a.z + b.z;
+        v.w = a.w + b.w;
+        return v;
+    }
 
-  static inline __device__ float4 sub(float4 a, float4 b) {
-    float4 v;
-    v.x = a.x - b.x;
-    v.y = a.y - b.y;
-    v.z = a.z - b.z;
-    v.w = a.w - b.w;
-    return v;
-  }
+    static inline __device__ float4 sub(float4 a, float4 b) {
+        float4 v;
+        v.x = a.x - b.x;
+        v.y = a.y - b.y;
+        v.z = a.z - b.z;
+        v.w = a.w - b.w;
+        return v;
+    }
 
-  static inline __device__ float4 add(float4 a, float b) {
-    float4 v;
-    v.x = a.x + b;
-    v.y = a.y + b;
-    v.z = a.z + b;
-    v.w = a.w + b;
-    return v;
-  }
+    static inline __device__ float4 add(float4 a, float b) {
+        float4 v;
+        v.x = a.x + b;
+        v.y = a.y + b;
+        v.z = a.z + b;
+        v.w = a.w + b;
+        return v;
+    }
 
-  static inline __device__ float4 sub(float4 a, float b) {
-    float4 v;
-    v.x = a.x - b;
-    v.y = a.y - b;
-    v.z = a.z - b;
-    v.w = a.w - b;
-    return v;
-  }
+    static inline __device__ float4 sub(float4 a, float b) {
+        float4 v;
+        v.x = a.x - b;
+        v.y = a.y - b;
+        v.z = a.z - b;
+        v.w = a.w - b;
+        return v;
+    }
 
-  static inline __device__ float4 mul(float4 a, float4 b) {
-    float4 v;
-    v.x = a.x * b.x;
-    v.y = a.y * b.y;
-    v.z = a.z * b.z;
-    v.w = a.w * b.w;
-    return v;
-  }
+    static inline __device__ float4 mul(float4 a, float4 b) {
+        float4 v;
+        v.x = a.x * b.x;
+        v.y = a.y * b.y;
+        v.z = a.z * b.z;
+        v.w = a.w * b.w;
+        return v;
+    }
 
-  static inline __device__ float4 mul(float4 a, float b) {
-    float4 v;
-    v.x = a.x * b;
-    v.y = a.y * b;
-    v.z = a.z * b;
-    v.w = a.w * b;
-    return v;
-  }
+    static inline __device__ float4 mul(float4 a, float b) {
+        float4 v;
+        v.x = a.x * b;
+        v.y = a.y * b;
+        v.z = a.z * b;
+        v.w = a.w * b;
+        return v;
+    }
 
-  static inline __device__ float4 neg(float4 v) {
-    v.x = -v.x;
-    v.y = -v.y;
-    v.z = -v.z;
-    v.w = -v.w;
-    return v;
-  }
+    static inline __device__ float4 neg(float4 v) {
+        v.x = -v.x;
+        v.y = -v.y;
+        v.z = -v.z;
+        v.w = -v.w;
+        return v;
+    }
 
-  /// For a vector type, this is a horizontal add, returning sum(v_i)
-  static inline __device__ float reduceAdd(float4 v) {
-    return v.x + v.y + v.z + v.w;
-  }
+    /// For a vector type, this is a horizontal add, returning sum(v_i)
+    static inline __device__ float reduceAdd(float4 v) {
+        return v.x + v.y + v.z + v.w;
+    }
 
-  // not implemented for vector types
-  // static inline __device__ bool lt(float4 a, float4 b);
-  // static inline __device__ bool gt(float4 a, float4 b);
-  // static inline __device__ bool eq(float4 a, float4 b);
+    // not implemented for vector types
+    // static inline __device__ bool lt(float4 a, float4 b);
+    // static inline __device__ bool gt(float4 a, float4 b);
+    // static inline __device__ bool eq(float4 a, float4 b);
 
-  static inline __device__ float4 zero() {
-    float4 v;
-    v.x = 0.0f;
-    v.y = 0.0f;
-    v.z = 0.0f;
-    v.w = 0.0f;
-    return v;
-  }
+    static inline __device__ float4 zero() {
+        float4 v;
+        v.x = 0.0f;
+        v.y = 0.0f;
+        v.z = 0.0f;
+        v.w = 0.0f;
+        return v;
+    }
 };
 
 template <>
 struct Math<half> {
-  typedef half ScalarType;
+    typedef half ScalarType;
 
-  static inline __device__ half add(half a, half b) {
+    static inline __device__ half add(half a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hadd(a, b);
+        return __hadd(a, b);
 #else
-    return __float2half(__half2float(a) + __half2float(b));
+        return __float2half(__half2float(a) + __half2float(b));
 #endif
-  }
+    }
 
-  static inline __device__ half sub(half a, half b) {
+    static inline __device__ half sub(half a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hsub(a, b);
+        return __hsub(a, b);
 #else
-    return __float2half(__half2float(a) - __half2float(b));
+        return __float2half(__half2float(a) - __half2float(b));
 #endif
-  }
+    }
 
-  static inline __device__ half mul(half a, half b) {
+    static inline __device__ half mul(half a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hmul(a, b);
+        return __hmul(a, b);
 #else
-    return __float2half(__half2float(a) * __half2float(b));
+        return __float2half(__half2float(a) * __half2float(b));
 #endif
-  }
+    }
 
-  static inline __device__ half neg(half v) {
+    static inline __device__ half neg(half v) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hneg(v);
+        return __hneg(v);
 #else
-    return __float2half(-__half2float(v));
+        return __float2half(-__half2float(v));
 #endif
-  }
+    }
 
-  static inline __device__ float reduceAdd(half v) {
-    return ConvertTo<float>::to(v);
-  }
+    static inline __device__ float reduceAdd(half v) {
+        return ConvertTo<float>::to(v);
+    }
 
-  static inline __device__ bool lt(half a, half b) {
+    static inline __device__ bool lt(half a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hlt(a, b);
+        return __hlt(a, b);
 #else
-    return __half2float(a) < __half2float(b);
+        return __half2float(a) < __half2float(b);
 #endif
-  }
+    }
 
-  static inline __device__ bool gt(half a, half b) {
+    static inline __device__ bool gt(half a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hgt(a, b);
+        return __hgt(a, b);
 #else
-    return __half2float(a) > __half2float(b);
+        return __half2float(a) > __half2float(b);
 #endif
-  }
+    }
 
-  static inline __device__ bool eq(half a, half b) {
+    static inline __device__ bool eq(half a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __heq(a, b);
+        return __heq(a, b);
 #else
-    return __half2float(a) == __half2float(b);
+        return __half2float(a) == __half2float(b);
 #endif
-  }
+    }
 
-  static inline __device__ half zero() {
+    static inline __device__ half zero() {
 #if CUDA_VERSION >= 9000
-    return 0;
+        return 0;
 #else
-    half h;
-    h.x = 0;
-    return h;
+        half h;
+        h.x = 0;
+        return h;
 #endif
-  }
+    }
 };
 
 template <>
 struct Math<half2> {
-  typedef half ScalarType;
+    typedef half ScalarType;
 
-  static inline __device__ half2 add(half2 a, half2 b) {
+    static inline __device__ half2 add(half2 a, half2 b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hadd2(a, b);
+        return __hadd2(a, b);
 #else
-  float2 af = __half22float2(a);
-  float2 bf = __half22float2(b);
+        float2 af = __half22float2(a);
+        float2 bf = __half22float2(b);
 
-  af.x += bf.x;
-  af.y += bf.y;
+        af.x += bf.x;
+        af.y += bf.y;
 
-  return __float22half2_rn(af);
+        return __float22half2_rn(af);
 #endif
-  }
+    }
 
-  static inline __device__ half2 sub(half2 a, half2 b) {
+    static inline __device__ half2 sub(half2 a, half2 b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hsub2(a, b);
+        return __hsub2(a, b);
 #else
-  float2 af = __half22float2(a);
-  float2 bf = __half22float2(b);
+        float2 af = __half22float2(a);
+        float2 bf = __half22float2(b);
 
-  af.x -= bf.x;
-  af.y -= bf.y;
+        af.x -= bf.x;
+        af.y -= bf.y;
 
-  return __float22half2_rn(af);
+        return __float22half2_rn(af);
 #endif
-  }
+    }
 
-  static inline __device__ half2 add(half2 a, half b) {
+    static inline __device__ half2 add(half2 a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    half2 b2 = __half2half2(b);
-    return __hadd2(a, b2);
+        half2 b2 = __half2half2(b);
+        return __hadd2(a, b2);
 #else
-  float2 af = __half22float2(a);
-  float bf = __half2float(b);
+        float2 af = __half22float2(a);
+        float bf = __half2float(b);
 
-  af.x += bf;
-  af.y += bf;
+        af.x += bf;
+        af.y += bf;
 
-  return __float22half2_rn(af);
+        return __float22half2_rn(af);
 #endif
-  }
+    }
 
-  static inline __device__ half2 sub(half2 a, half b) {
+    static inline __device__ half2 sub(half2 a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    half2 b2 = __half2half2(b);
-    return __hsub2(a, b2);
+        half2 b2 = __half2half2(b);
+        return __hsub2(a, b2);
 #else
-  float2 af = __half22float2(a);
-  float bf = __half2float(b);
+        float2 af = __half22float2(a);
+        float bf = __half2float(b);
 
-  af.x -= bf;
-  af.y -= bf;
+        af.x -= bf;
+        af.y -= bf;
 
-  return __float22half2_rn(af);
+        return __float22half2_rn(af);
 #endif
-  }
+    }
 
-  static inline __device__ half2 mul(half2 a, half2 b) {
+    static inline __device__ half2 mul(half2 a, half2 b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hmul2(a, b);
+        return __hmul2(a, b);
 #else
-  float2 af = __half22float2(a);
-  float2 bf = __half22float2(b);
+        float2 af = __half22float2(a);
+        float2 bf = __half22float2(b);
 
-  af.x *= bf.x;
-  af.y *= bf.y;
+        af.x *= bf.x;
+        af.y *= bf.y;
 
-  return __float22half2_rn(af);
+        return __float22half2_rn(af);
 #endif
-  }
+    }
 
-  static inline __device__ half2 mul(half2 a, half b) {
+    static inline __device__ half2 mul(half2 a, half b) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    half2 b2 = __half2half2(b);
-    return __hmul2(a, b2);
+        half2 b2 = __half2half2(b);
+        return __hmul2(a, b2);
 #else
-  float2 af = __half22float2(a);
-  float bf = __half2float(b);
+        float2 af = __half22float2(a);
+        float bf = __half2float(b);
 
-  af.x *= bf;
-  af.y *= bf;
+        af.x *= bf;
+        af.y *= bf;
 
-  return __float22half2_rn(af);
+        return __float22half2_rn(af);
 #endif
-  }
+    }
 
-  static inline __device__ half2 neg(half2 v) {
+    static inline __device__ half2 neg(half2 v) {
 #ifdef FAISS_USE_FULL_FLOAT16
-    return __hneg2(v);
+        return __hneg2(v);
 #else
-  float2 vf = __half22float2(v);
-  vf.x = -vf.x;
-  vf.y = -vf.y;
+        float2 vf = __half22float2(v);
+        vf.x = -vf.x;
+        vf.y = -vf.y;
 
-  return __float22half2_rn(vf);
+        return __float22half2_rn(vf);
 #endif
-  }
+    }
 
-  static inline __device__ float reduceAdd(half2 v) {
-    float2 vf = __half22float2(v);
-    vf.x += vf.y;
+    static inline __device__ float reduceAdd(half2 v) {
+        float2 vf = __half22float2(v);
+        vf.x += vf.y;
 
-    return vf.x;
-  }
+        return vf.x;
+    }
 
-  // not implemented for vector types
-  // static inline __device__ bool lt(half2 a, half2 b);
-  // static inline __device__ bool gt(half2 a, half2 b);
-  // static inline __device__ bool eq(half2 a, half2 b);
+    // not implemented for vector types
+    // static inline __device__ bool lt(half2 a, half2 b);
+    // static inline __device__ bool gt(half2 a, half2 b);
+    // static inline __device__ bool eq(half2 a, half2 b);
 
-  static inline __device__ half2 zero() {
-    return __half2half2(Math<half>::zero());
-  }
+    static inline __device__ half2 zero() {
+        return __half2half2(Math<half>::zero());
+    }
 };
 
 template <>
 struct Math<Half4> {
-  typedef half ScalarType;
+    typedef half ScalarType;
 
-  static inline __device__ Half4 add(Half4 a, Half4 b) {
-    Half4 h;
-    h.a = Math<half2>::add(a.a, b.a);
-    h.b = Math<half2>::add(a.b, b.b);
-    return h;
-  }
+    static inline __device__ Half4 add(Half4 a, Half4 b) {
+        Half4 h;
+        h.a = Math<half2>::add(a.a, b.a);
+        h.b = Math<half2>::add(a.b, b.b);
+        return h;
+    }
 
-  static inline __device__ Half4 sub(Half4 a, Half4 b) {
-    Half4 h;
-    h.a = Math<half2>::sub(a.a, b.a);
-    h.b = Math<half2>::sub(a.b, b.b);
-    return h;
-  }
+    static inline __device__ Half4 sub(Half4 a, Half4 b) {
+        Half4 h;
+        h.a = Math<half2>::sub(a.a, b.a);
+        h.b = Math<half2>::sub(a.b, b.b);
+        return h;
+    }
 
-  static inline __device__ Half4 add(Half4 a, half b) {
-    Half4 h;
-    h.a = Math<half2>::add(a.a, b);
-    h.b = Math<half2>::add(a.b, b);
-    return h;
-  }
+    static inline __device__ Half4 add(Half4 a, half b) {
+        Half4 h;
+        h.a = Math<half2>::add(a.a, b);
+        h.b = Math<half2>::add(a.b, b);
+        return h;
+    }
 
-  static inline __device__ Half4 sub(Half4 a, half b) {
-    Half4 h;
-    h.a = Math<half2>::sub(a.a, b);
-    h.b = Math<half2>::sub(a.b, b);
-    return h;
-  }
+    static inline __device__ Half4 sub(Half4 a, half b) {
+        Half4 h;
+        h.a = Math<half2>::sub(a.a, b);
+        h.b = Math<half2>::sub(a.b, b);
+        return h;
+    }
 
-  static inline __device__ Half4 mul(Half4 a, Half4 b) {
-    Half4 h;
-    h.a = Math<half2>::mul(a.a, b.a);
-    h.b = Math<half2>::mul(a.b, b.b);
-    return h;
-  }
+    static inline __device__ Half4 mul(Half4 a, Half4 b) {
+        Half4 h;
+        h.a = Math<half2>::mul(a.a, b.a);
+        h.b = Math<half2>::mul(a.b, b.b);
+        return h;
+    }
 
-  static inline __device__ Half4 mul(Half4 a, half b) {
-    Half4 h;
-    h.a = Math<half2>::mul(a.a, b);
-    h.b = Math<half2>::mul(a.b, b);
-    return h;
-  }
+    static inline __device__ Half4 mul(Half4 a, half b) {
+        Half4 h;
+        h.a = Math<half2>::mul(a.a, b);
+        h.b = Math<half2>::mul(a.b, b);
+        return h;
+    }
 
-  static inline __device__ Half4 neg(Half4 v) {
-    Half4 h;
-    h.a = Math<half2>::neg(v.a);
-    h.b = Math<half2>::neg(v.b);
-    return h;
-  }
+    static inline __device__ Half4 neg(Half4 v) {
+        Half4 h;
+        h.a = Math<half2>::neg(v.a);
+        h.b = Math<half2>::neg(v.b);
+        return h;
+    }
 
-  static inline __device__ float reduceAdd(Half4 v) {
-    float x = Math<half2>::reduceAdd(v.a);
-    float y = Math<half2>::reduceAdd(v.b);
-    return x + y;
-  }
+    static inline __device__ float reduceAdd(Half4 v) {
+        float x = Math<half2>::reduceAdd(v.a);
+        float y = Math<half2>::reduceAdd(v.b);
+        return x + y;
+    }
 
-  // not implemented for vector types
-  // static inline __device__ bool lt(Half4 a, Half4 b);
-  // static inline __device__ bool gt(Half4 a, Half4 b);
-  // static inline __device__ bool eq(Half4 a, Half4 b);
+    // not implemented for vector types
+    // static inline __device__ bool lt(Half4 a, Half4 b);
+    // static inline __device__ bool gt(Half4 a, Half4 b);
+    // static inline __device__ bool eq(Half4 a, Half4 b);
 
-  static inline __device__ Half4 zero() {
-    Half4 h;
-    h.a = Math<half2>::zero();
-    h.b = Math<half2>::zero();
-    return h;
-  }
+    static inline __device__ Half4 zero() {
+        Half4 h;
+        h.a = Math<half2>::zero();
+        h.b = Math<half2>::zero();
+        return h;
+    }
 };
 
 template <>
 struct Math<Half8> {
-  typedef half ScalarType;
+    typedef half ScalarType;
 
-  static inline __device__ Half8 add(Half8 a, Half8 b) {
-    Half8 h;
-    h.a = Math<Half4>::add(a.a, b.a);
-    h.b = Math<Half4>::add(a.b, b.b);
-    return h;
-  }
+    static inline __device__ Half8 add(Half8 a, Half8 b) {
+        Half8 h;
+        h.a = Math<Half4>::add(a.a, b.a);
+        h.b = Math<Half4>::add(a.b, b.b);
+        return h;
+    }
 
-  static inline __device__ Half8 sub(Half8 a, Half8 b) {
-    Half8 h;
-    h.a = Math<Half4>::sub(a.a, b.a);
-    h.b = Math<Half4>::sub(a.b, b.b);
-    return h;
-  }
+    static inline __device__ Half8 sub(Half8 a, Half8 b) {
+        Half8 h;
+        h.a = Math<Half4>::sub(a.a, b.a);
+        h.b = Math<Half4>::sub(a.b, b.b);
+        return h;
+    }
 
-  static inline __device__ Half8 add(Half8 a, half b) {
-    Half8 h;
-    h.a = Math<Half4>::add(a.a, b);
-    h.b = Math<Half4>::add(a.b, b);
-    return h;
-  }
+    static inline __device__ Half8 add(Half8 a, half b) {
+        Half8 h;
+        h.a = Math<Half4>::add(a.a, b);
+        h.b = Math<Half4>::add(a.b, b);
+        return h;
+    }
 
-  static inline __device__ Half8 sub(Half8 a, half b) {
-    Half8 h;
-    h.a = Math<Half4>::sub(a.a, b);
-    h.b = Math<Half4>::sub(a.b, b);
-    return h;
-  }
+    static inline __device__ Half8 sub(Half8 a, half b) {
+        Half8 h;
+        h.a = Math<Half4>::sub(a.a, b);
+        h.b = Math<Half4>::sub(a.b, b);
+        return h;
+    }
 
-  static inline __device__ Half8 mul(Half8 a, Half8 b) {
-    Half8 h;
-    h.a = Math<Half4>::mul(a.a, b.a);
-    h.b = Math<Half4>::mul(a.b, b.b);
-    return h;
-  }
+    static inline __device__ Half8 mul(Half8 a, Half8 b) {
+        Half8 h;
+        h.a = Math<Half4>::mul(a.a, b.a);
+        h.b = Math<Half4>::mul(a.b, b.b);
+        return h;
+    }
 
-  static inline __device__ Half8 mul(Half8 a, half b) {
-    Half8 h;
-    h.a = Math<Half4>::mul(a.a, b);
-    h.b = Math<Half4>::mul(a.b, b);
-    return h;
-  }
+    static inline __device__ Half8 mul(Half8 a, half b) {
+        Half8 h;
+        h.a = Math<Half4>::mul(a.a, b);
+        h.b = Math<Half4>::mul(a.b, b);
+        return h;
+    }
 
-  static inline __device__ Half8 neg(Half8 v) {
-    Half8 h;
-    h.a = Math<Half4>::neg(v.a);
-    h.b = Math<Half4>::neg(v.b);
-    return h;
-  }
+    static inline __device__ Half8 neg(Half8 v) {
+        Half8 h;
+        h.a = Math<Half4>::neg(v.a);
+        h.b = Math<Half4>::neg(v.b);
+        return h;
+    }
 
-  static inline __device__ float reduceAdd(Half8 v) {
-    float x = Math<Half4>::reduceAdd(v.a);
-    float y = Math<Half4>::reduceAdd(v.b);
-    return x + y;
-  }
+    static inline __device__ float reduceAdd(Half8 v) {
+        float x = Math<Half4>::reduceAdd(v.a);
+        float y = Math<Half4>::reduceAdd(v.b);
+        return x + y;
+    }
 
-  // not implemented for vector types
-  // static inline __device__ bool lt(Half8 a, Half8 b);
-  // static inline __device__ bool gt(Half8 a, Half8 b);
-  // static inline __device__ bool eq(Half8 a, Half8 b);
+    // not implemented for vector types
+    // static inline __device__ bool lt(Half8 a, Half8 b);
+    // static inline __device__ bool gt(Half8 a, Half8 b);
+    // static inline __device__ bool eq(Half8 a, Half8 b);
 
-  static inline __device__ Half8 zero() {
-    Half8 h;
-    h.a = Math<Half4>::zero();
-    h.b = Math<Half4>::zero();
-    return h;
-  }
+    static inline __device__ Half8 zero() {
+        Half8 h;
+        h.a = Math<Half4>::zero();
+        h.b = Math<Half4>::zero();
+        return h;
+    }
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/MatrixMult-inl.cuh

@@ -5,258 +5,356 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cublas_v2.h>
-#include <faiss/gpu/utils/Tensor.cuh>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
-#include <faiss/gpu/utils/HostTensor.cuh>
 #include <faiss/gpu/utils/Float16.cuh>
+#include <faiss/gpu/utils/HostTensor.cuh>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T>
 struct GetCudaType;
 
 template <>
 struct GetCudaType<float> {
-  static constexpr cudaDataType_t Type = CUDA_R_32F;
+    static constexpr cudaDataType_t Type = CUDA_R_32F;
 };
 
 template <>
 struct GetCudaType<half> {
-  static constexpr cudaDataType_t Type = CUDA_R_16F;
+    static constexpr cudaDataType_t Type = CUDA_R_16F;
 };
 
 template <typename AT, typename BT>
-cublasStatus_t
-rawGemm(cublasHandle_t handle,
+cublasStatus_t rawGemm(
+        cublasHandle_t handle,
         cublasOperation_t transa,
         cublasOperation_t transb,
         int m,
         int n,
         int k,
         const float fAlpha,
-        const void *A,
+        const void* A,
         int lda,
-        const void *B,
+        const void* B,
         int ldb,
         const float fBeta,
-        float *C,
+        float* C,
         int ldc) {
-  auto cAT = GetCudaType<AT>::Type;
-  auto cBT = GetCudaType<BT>::Type;
+    auto cAT = GetCudaType<AT>::Type;
+    auto cBT = GetCudaType<BT>::Type;
 
-  // FIXME: some weird CUDA 11 bug? where cublasSgemmEx on
-  // f16 (8, 64) x f16 (64, 64)' = f32 (8, 64) returns "not supported".
-  // cublasGemmEx using CUBLAS_COMPUTE_32F also fails, but
-  // CUBLAS_COMPUTE_32F_PEDANTIC does not fail (as seen on a V100).
-  //
-  // Only use the PEDANTIC implementation if the input matrices are f16
-  // and we are on CUDA 11+
+    // FIXME: some weird CUDA 11 bug? where cublasSgemmEx on
+    // f16 (8, 64) x f16 (64, 64)' = f32 (8, 64) returns "not supported".
+    // cublasGemmEx using CUBLAS_COMPUTE_32F also fails, but
+    // CUBLAS_COMPUTE_32F_PEDANTIC does not fail (as seen on a V100).
+    //
+    // Only use the PEDANTIC implementation if the input matrices are f16
+    // and we are on CUDA 11+
 #if CUDA_VERSION >= 11000
-  if (cAT == CUDA_R_16F || cBT == CUDA_R_16F) {
-    return cublasGemmEx(handle, transa, transb, m, n, k,
-                        &fAlpha, A, cAT, lda,
-                        B, cBT, ldb,
-                        &fBeta,
-                        C, CUDA_R_32F, ldc,
-                        CUBLAS_COMPUTE_32F_PEDANTIC, CUBLAS_GEMM_DEFAULT);
-  }
+    if (cAT == CUDA_R_16F || cBT == CUDA_R_16F) {
+        return cublasGemmEx(
+                handle,
+                transa,
+                transb,
+                m,
+                n,
+                k,
+                &fAlpha,
+                A,
+                cAT,
+                lda,
+                B,
+                cBT,
+                ldb,
+                &fBeta,
+                C,
+                CUDA_R_32F,
+                ldc,
+                CUBLAS_COMPUTE_32F_PEDANTIC,
+                CUBLAS_GEMM_DEFAULT);
+    }
 #endif
 
-  // Always accumulate in f32
-  return cublasSgemmEx(handle, transa, transb, m, n, k,
-                       &fAlpha, A, cAT, lda,
-                       B, cBT, ldb,
-                       &fBeta,
-                       C, CUDA_R_32F, ldc);
+    // Always accumulate in f32
+    return cublasSgemmEx(
+            handle,
+            transa,
+            transb,
+            m,
+            n,
+            k,
+            &fAlpha,
+            A,
+            cAT,
+            lda,
+            B,
+            cBT,
+            ldb,
+            &fBeta,
+            C,
+            CUDA_R_32F,
+            ldc);
 }
 
 template <typename AT, typename BT>
-cublasStatus_t
-rawBatchGemm(cublasHandle_t handle,
-             cublasOperation_t transa,
-             cublasOperation_t transb,
-             int m,
-             int n,
-             int k,
-             const float fAlpha,
-             const void* A,
-             int lda,
-             long long int strideA,
-             const void* B,
-             int ldb,
-             long long int strideB,
-             const float fBeta,
-             float* C,
-             int ldc,
-             long long int strideC,
-             int batchCount) {
-  auto cAT = GetCudaType<AT>::Type;
-  auto cBT = GetCudaType<BT>::Type;
+cublasStatus_t rawBatchGemm(
+        cublasHandle_t handle,
+        cublasOperation_t transa,
+        cublasOperation_t transb,
+        int m,
+        int n,
+        int k,
+        const float fAlpha,
+        const void* A,
+        int lda,
+        long long int strideA,
+        const void* B,
+        int ldb,
+        long long int strideB,
+        const float fBeta,
+        float* C,
+        int ldc,
+        long long int strideC,
+        int batchCount) {
+    auto cAT = GetCudaType<AT>::Type;
+    auto cBT = GetCudaType<BT>::Type;
 
-  // Always accumulate in f32
-  return cublasGemmStridedBatchedEx(handle, transa, transb, m, n, k,
-                                    &fAlpha, A, cAT, lda, strideA,
-                                    B, cBT, ldb, strideB, &fBeta,
-                                    C, CUDA_R_32F, ldc, strideC, batchCount,
-                                    CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
+    // Always accumulate in f32
+    return cublasGemmStridedBatchedEx(
+            handle,
+            transa,
+            transb,
+            m,
+            n,
+            k,
+            &fAlpha,
+            A,
+            cAT,
+            lda,
+            strideA,
+            B,
+            cBT,
+            ldb,
+            strideB,
+            &fBeta,
+            C,
+            CUDA_R_32F,
+            ldc,
+            strideC,
+            batchCount,
+            CUDA_R_32F,
+            CUBLAS_GEMM_DEFAULT);
 }
 
 template <typename AT, typename BT>
-void
-runMatrixMult(Tensor<float, 2, true>& c, bool transC,
-              Tensor<AT, 2, true>& a, bool transA,
-              Tensor<BT, 2, true>& b, bool transB,
-              float alpha,
-              float beta,
-              cublasHandle_t handle,
-              cudaStream_t stream) {
-  cublasSetStream(handle, stream);
+void runMatrixMult(
+        Tensor<float, 2, true>& c,
+        bool transC,
+        Tensor<AT, 2, true>& a,
+        bool transA,
+        Tensor<BT, 2, true>& b,
+        bool transB,
+        float alpha,
+        float beta,
+        cublasHandle_t handle,
+        cudaStream_t stream) {
+    cublasSetStream(handle, stream);
 
-  // Check that we have (m x k) * (k x n) = (m x n)
-  // using the input row-major layout
-  int aM = transA ? a.getSize(1) : a.getSize(0);
-  int aK = transA ? a.getSize(0) : a.getSize(1);
+    // Check that we have (m x k) * (k x n) = (m x n)
+    // using the input row-major layout
+    int aM = transA ? a.getSize(1) : a.getSize(0);
+    int aK = transA ? a.getSize(0) : a.getSize(1);
 
-  int bK = transB ? b.getSize(1) : b.getSize(0);
-  int bN = transB ? b.getSize(0) : b.getSize(1);
+    int bK = transB ? b.getSize(1) : b.getSize(0);
+    int bN = transB ? b.getSize(0) : b.getSize(1);
 
-  int cM = transC ? c.getSize(1) : c.getSize(0);
-  int cN = transC ? c.getSize(0) : c.getSize(1);
+    int cM = transC ? c.getSize(1) : c.getSize(0);
+    int cN = transC ? c.getSize(0) : c.getSize(1);
 
-  FAISS_ASSERT(aM == cM);
-  FAISS_ASSERT(aK == bK);
-  FAISS_ASSERT(bN == cN);
+    FAISS_ASSERT(aM == cM);
+    FAISS_ASSERT(aK == bK);
+    FAISS_ASSERT(bN == cN);
 
-  FAISS_ASSERT(a.getStride(1) == 1);
-  FAISS_ASSERT(b.getStride(1) == 1);
-  FAISS_ASSERT(c.getStride(1) == 1);
+    FAISS_ASSERT(a.getStride(1) == 1);
+    FAISS_ASSERT(b.getStride(1) == 1);
+    FAISS_ASSERT(c.getStride(1) == 1);
 
-  // Now, we have to represent the matrix multiplication in
-  // column-major layout
-  float* pC = c.data();
+    // Now, we have to represent the matrix multiplication in
+    // column-major layout
+    float* pC = c.data();
 
-  int m = c.getSize(1); // stride 1 size
-  int n = c.getSize(0); // other size
-  int k = transA ? a.getSize(0) : a.getSize(1);
+    int m = c.getSize(1); // stride 1 size
+    int n = c.getSize(0); // other size
+    int k = transA ? a.getSize(0) : a.getSize(1);
 
-  int lda = transC ? a.getStride(0) : b.getStride(0);
-  int ldb = transC ? b.getStride(0) : a.getStride(0);
-  int ldc = c.getStride(0);
+    int lda = transC ? a.getStride(0) : b.getStride(0);
+    int ldb = transC ? b.getStride(0) : a.getStride(0);
+    int ldc = c.getStride(0);
 
-  auto gemmTrA = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
-  auto gemmTrB = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
+    auto gemmTrA = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
+    auto gemmTrB = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
 
-  if (transC) {
-    gemmTrA = transA ? CUBLAS_OP_N : CUBLAS_OP_T;
-    gemmTrB = transB ? CUBLAS_OP_N : CUBLAS_OP_T;
-  }
+    if (transC) {
+        gemmTrA = transA ? CUBLAS_OP_N : CUBLAS_OP_T;
+        gemmTrB = transB ? CUBLAS_OP_N : CUBLAS_OP_T;
+    }
 
-  cublasStatus_t err;
+    cublasStatus_t err;
 
-  if (transC) {
-    err = rawGemm<AT, BT>(handle,
-                          gemmTrA, gemmTrB,
-                          m, n, k, alpha,
-                          a.data(), lda, b.data(), ldb, beta,
-                          pC, ldc);
-  } else {
-    err = rawGemm<AT, BT>(handle,
-                          gemmTrA, gemmTrB,
-                          m, n, k, alpha,
-                          b.data(), lda, a.data(), ldb, beta,
-                          pC, ldc);
-  }
+    if (transC) {
+        err = rawGemm<AT, BT>(
+                handle,
+                gemmTrA,
+                gemmTrB,
+                m,
+                n,
+                k,
+                alpha,
+                a.data(),
+                lda,
+                b.data(),
+                ldb,
+                beta,
+                pC,
+                ldc);
+    } else {
+        err = rawGemm<AT, BT>(
+                handle,
+                gemmTrA,
+                gemmTrB,
+                m,
+                n,
+                k,
+                alpha,
+                b.data(),
+                lda,
+                a.data(),
+                ldb,
+                beta,
+                pC,
+                ldc);
+    }
 
-  FAISS_ASSERT_FMT(err == CUBLAS_STATUS_SUCCESS,
-                   "cublas failed (%d): "
-                   "(%d, %d)%s x (%d, %d)%s = (%d, %d)%s "
-                   "gemm params m %d n %d k %d trA %s trB %s lda %d ldb %d ldc %d",
-                   (int) err,
-                   a.getSize(0), a.getSize(1), transA ? "'" : "",
-                   b.getSize(0), b.getSize(1), transB ? "'" : "",
-                   c.getSize(0), c.getSize(1), transC ? "'" : "",
-                   m, n, k,
-                   gemmTrA == CUBLAS_OP_T ? "T" : "N",
-                   gemmTrB == CUBLAS_OP_T ? "T" : "N",
-                   lda, ldb, ldc);
-  CUDA_TEST_ERROR();
+    FAISS_ASSERT_FMT(
+            err == CUBLAS_STATUS_SUCCESS,
+            "cublas failed (%d): "
+            "(%d, %d)%s x (%d, %d)%s = (%d, %d)%s "
+            "gemm params m %d n %d k %d trA %s trB %s lda %d ldb %d ldc %d",
+            (int)err,
+            a.getSize(0),
+            a.getSize(1),
+            transA ? "'" : "",
+            b.getSize(0),
+            b.getSize(1),
+            transB ? "'" : "",
+            c.getSize(0),
+            c.getSize(1),
+            transC ? "'" : "",
+            m,
+            n,
+            k,
+            gemmTrA == CUBLAS_OP_T ? "T" : "N",
+            gemmTrB == CUBLAS_OP_T ? "T" : "N",
+            lda,
+            ldb,
+            ldc);
+    CUDA_TEST_ERROR();
 }
 
 template <typename AT, typename BT>
-void
-runBatchMatrixMult(Tensor<float, 3, true>& c, bool transC,
-                   Tensor<AT, 3, true>& a, bool transA,
-                   Tensor<BT, 3, true>& b, bool transB,
-                   float alpha,
-                   float beta,
-                   cublasHandle_t handle,
-                   cudaStream_t stream) {
-  FAISS_ASSERT(c.getSize(0) == a.getSize(0));
-  FAISS_ASSERT(a.getSize(0) == b.getSize(0));
+void runBatchMatrixMult(
+        Tensor<float, 3, true>& c,
+        bool transC,
+        Tensor<AT, 3, true>& a,
+        bool transA,
+        Tensor<BT, 3, true>& b,
+        bool transB,
+        float alpha,
+        float beta,
+        cublasHandle_t handle,
+        cudaStream_t stream) {
+    FAISS_ASSERT(c.getSize(0) == a.getSize(0));
+    FAISS_ASSERT(a.getSize(0) == b.getSize(0));
 
-  // This uses the strided batch MM, which assumes a uniform stride
-  FAISS_ASSERT(a.getStride(0) == a.getSize(1) * a.getSize(2));
-  FAISS_ASSERT(b.getStride(0) == b.getSize(1) * b.getSize(2));
-  FAISS_ASSERT(c.getStride(0) == c.getSize(1) * c.getSize(2));
+    // This uses the strided batch MM, which assumes a uniform stride
+    FAISS_ASSERT(a.getStride(0) == a.getSize(1) * a.getSize(2));
+    FAISS_ASSERT(b.getStride(0) == b.getSize(1) * b.getSize(2));
+    FAISS_ASSERT(c.getStride(0) == c.getSize(1) * c.getSize(2));
 
-  cublasSetStream(handle, stream);
+    cublasSetStream(handle, stream);
 
-  // Check that we have (m x k) * (k x n) = (m x n)
-  // using the input row-major layout
-  int aM = transA ? a.getSize(2) : a.getSize(1);
-  int aK = transA ? a.getSize(1) : a.getSize(2);
+    // Check that we have (m x k) * (k x n) = (m x n)
+    // using the input row-major layout
+    int aM = transA ? a.getSize(2) : a.getSize(1);
+    int aK = transA ? a.getSize(1) : a.getSize(2);
 
-  int bK = transB ? b.getSize(2) : b.getSize(1);
-  int bN = transB ? b.getSize(1) : b.getSize(2);
+    int bK = transB ? b.getSize(2) : b.getSize(1);
+    int bN = transB ? b.getSize(1) : b.getSize(2);
 
-  int cM = transC ? c.getSize(2) : c.getSize(1);
-  int cN = transC ? c.getSize(1) : c.getSize(2);
+    int cM = transC ? c.getSize(2) : c.getSize(1);
+    int cN = transC ? c.getSize(1) : c.getSize(2);
 
-  FAISS_ASSERT(aM == cM);
-  FAISS_ASSERT(aK == bK);
-  FAISS_ASSERT(bN == cN);
+    FAISS_ASSERT(aM == cM);
+    FAISS_ASSERT(aK == bK);
+    FAISS_ASSERT(bN == cN);
 
-  // Now, we have to represent the matrix multiplication in
-  // column-major layout
-  void* pA = transC ? (void*) a.data() : (void*) b.data();
-  void* pB = transC ? (void*) b.data() : (void*) a.data();
-  float* pC = c.data();
+    // Now, we have to represent the matrix multiplication in
+    // column-major layout
+    void* pA = transC ? (void*)a.data() : (void*)b.data();
+    void* pB = transC ? (void*)b.data() : (void*)a.data();
+    float* pC = c.data();
 
-  int m = c.getSize(2); // stride 1 size
-  int n = c.getSize(1); // other size
-  int k = transA ? a.getSize(1) : a.getSize(2);
+    int m = c.getSize(2); // stride 1 size
+    int n = c.getSize(1); // other size
+    int k = transA ? a.getSize(1) : a.getSize(2);
 
-  int lda = transC ? a.getStride(1) : b.getStride(1);
-  int ldb = transC ? b.getStride(1) : a.getStride(1);
-  int ldc = c.getStride(1);
+    int lda = transC ? a.getStride(1) : b.getStride(1);
+    int ldb = transC ? b.getStride(1) : a.getStride(1);
+    int ldc = c.getStride(1);
 
-  auto gemmTrA = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
-  auto gemmTrB = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
+    auto gemmTrA = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
+    auto gemmTrB = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
 
-  if (transC) {
-    gemmTrA = transA ? CUBLAS_OP_N : CUBLAS_OP_T;
-    gemmTrB = transB ? CUBLAS_OP_N : CUBLAS_OP_T;
-  }
+    if (transC) {
+        gemmTrA = transA ? CUBLAS_OP_N : CUBLAS_OP_T;
+        gemmTrB = transB ? CUBLAS_OP_N : CUBLAS_OP_T;
+    }
 
-  long long int gemmStrideA = transC ? a.getStride(0) : b.getStride(0);
-  long long int gemmStrideB = transC ? b.getStride(0) : a.getStride(0);
-  long long int gemmStrideC = c.getStride(0);
+    long long int gemmStrideA = transC ? a.getStride(0) : b.getStride(0);
+    long long int gemmStrideB = transC ? b.getStride(0) : a.getStride(0);
+    long long int gemmStrideC = c.getStride(0);
 
-  auto err =
-    rawBatchGemm<AT, BT>(handle,
-                         gemmTrA, gemmTrB,
-                         m, n, k, alpha,
-                         pA, lda, gemmStrideA,
-                         pB, ldb, gemmStrideB, beta,
-                         pC, ldc, gemmStrideC, a.getSize(0));
+    auto err = rawBatchGemm<AT, BT>(
+            handle,
+            gemmTrA,
+            gemmTrB,
+            m,
+            n,
+            k,
+            alpha,
+            pA,
+            lda,
+            gemmStrideA,
+            pB,
+            ldb,
+            gemmStrideB,
+            beta,
+            pC,
+            ldc,
+            gemmStrideC,
+            a.getSize(0));
 
-  FAISS_ASSERT_FMT(err == CUBLAS_STATUS_SUCCESS,
-                   "cublasGemmStridedBatchedEx failed (%d)", (int) err);
-  CUDA_TEST_ERROR();
+    FAISS_ASSERT_FMT(
+            err == CUBLAS_STATUS_SUCCESS,
+            "cublasGemmStridedBatchedEx failed (%d)",
+            (int)err);
+    CUDA_TEST_ERROR();
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/MatrixMult.cuh

@@ -5,43 +5,51 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cublas_v2.h>
-#include <faiss/gpu/utils/Tensor.cuh>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
-#include <faiss/gpu/utils/HostTensor.cuh>
 #include <faiss/gpu/utils/Float16.cuh>
+#include <faiss/gpu/utils/HostTensor.cuh>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class GpuResources;
 
 /// C = alpha * A * B + beta * C
 /// Expects row major layout, not fortran/blas column major!
 template <typename AT, typename BT>
-void
-runMatrixMult(Tensor<float, 2, true>& c, bool transC,
-              Tensor<AT, 2, true>& a, bool transA,
-              Tensor<BT, 2, true>& b, bool transB,
-              float alpha,
-              float beta,
-              cublasHandle_t handle,
-              cudaStream_t stream);
+void runMatrixMult(
+        Tensor<float, 2, true>& c,
+        bool transC,
+        Tensor<AT, 2, true>& a,
+        bool transA,
+        Tensor<BT, 2, true>& b,
+        bool transB,
+        float alpha,
+        float beta,
+        cublasHandle_t handle,
+        cudaStream_t stream);
 
 /// C_i = alpha * A_i * B_i + beta * C_i
 /// where `i` is the outermost dimension, via iterated gemm
 /// Expects row major layout, not fortran/blas column major!
 template <typename AT, typename BT>
-void runIteratedMatrixMult(Tensor<float, 3, true>& c, bool transC,
-                           Tensor<AT, 3, true>& a, bool transA,
-                           Tensor<BT, 3, true>& b, bool transB,
-                           float alpha,
-                           float beta,
-                           cublasHandle_t handle,
-                           cudaStream_t stream);
+void runIteratedMatrixMult(
+        Tensor<float, 3, true>& c,
+        bool transC,
+        Tensor<AT, 3, true>& a,
+        bool transA,
+        Tensor<BT, 3, true>& b,
+        bool transB,
+        float alpha,
+        float beta,
+        cublasHandle_t handle,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss
 
 #include <faiss/gpu/utils/MatrixMult-inl.cuh>

faiss/gpu/utils/MergeNetworkBlock.cuh

@@ -7,283 +7,315 @@
 
 #pragma once
 
+#include <cuda.h>
+#include <faiss/gpu/utils/StaticUtils.h>
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
 #include <faiss/gpu/utils/MergeNetworkUtils.cuh>
 #include <faiss/gpu/utils/PtxUtils.cuh>
-#include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/gpu/utils/WarpShuffles.cuh>
-#include <faiss/impl/FaissAssert.h>
-#include <cuda.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // Merge pairs of lists smaller than blockDim.x (NumThreads)
-template <int NumThreads,
-          typename K,
-          typename V,
-          int N,
-          int L,
-          bool AllThreads,
-          bool Dir,
-          typename Comp,
-          bool FullMerge>
+template <
+        int NumThreads,
+        typename K,
+        typename V,
+        int N,
+        int L,
+        bool AllThreads,
+        bool Dir,
+        typename Comp,
+        bool FullMerge>
 inline __device__ void blockMergeSmall(K* listK, V* listV) {
-  static_assert(utils::isPowerOf2(L), "L must be a power-of-2");
-  static_assert(utils::isPowerOf2(NumThreads),
-                "NumThreads must be a power-of-2");
-  static_assert(L <= NumThreads, "merge list size must be <= NumThreads");
+    static_assert(utils::isPowerOf2(L), "L must be a power-of-2");
+    static_assert(
+            utils::isPowerOf2(NumThreads), "NumThreads must be a power-of-2");
+    static_assert(L <= NumThreads, "merge list size must be <= NumThreads");
 
-  // Which pair of lists we are merging
-  int mergeId = threadIdx.x / L;
+    // Which pair of lists we are merging
+    int mergeId = threadIdx.x / L;
 
-  // Which thread we are within the merge
-  int tid = threadIdx.x % L;
+    // Which thread we are within the merge
+    int tid = threadIdx.x % L;
 
-  // listK points to a region of size N * 2 * L
-  listK += 2 * L * mergeId;
-  listV += 2 * L * mergeId;
+    // listK points to a region of size N * 2 * L
+    listK += 2 * L * mergeId;
+    listV += 2 * L * mergeId;
 
-  // It's not a bitonic merge, both lists are in the same direction,
-  // so handle the first swap assuming the second list is reversed
-  int pos = L - 1 - tid;
-  int stride = 2 * tid + 1;
-
-  if (AllThreads || (threadIdx.x < N * L)) {
-    K ka = listK[pos];
-    K kb = listK[pos + stride];
-
-    bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-    listK[pos] = swap ? kb : ka;
-    listK[pos + stride] = swap ? ka : kb;
-
-    V va = listV[pos];
-    V vb = listV[pos + stride];
-    listV[pos] = swap ? vb : va;
-    listV[pos + stride] = swap ? va : vb;
-
-    // FIXME: is this a CUDA 9 compiler bug?
-    // K& ka = listK[pos];
-    // K& kb = listK[pos + stride];
-
-    // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-    // swap(s, ka, kb);
-
-    // V& va = listV[pos];
-    // V& vb = listV[pos + stride];
-    // swap(s, va, vb);
-  }
-
-  __syncthreads();
-
-#pragma unroll
-  for (int stride = L / 2; stride > 0; stride /= 2) {
-    int pos = 2 * tid - (tid & (stride - 1));
-
-    if (AllThreads || (threadIdx.x < N * L)) {
-      K ka = listK[pos];
-      K kb = listK[pos + stride];
-
-      bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-      listK[pos] = swap ? kb : ka;
-      listK[pos + stride] = swap ? ka : kb;
-
-      V va = listV[pos];
-      V vb = listV[pos + stride];
-      listV[pos] = swap ? vb : va;
-      listV[pos + stride] = swap ? va : vb;
-
-      // FIXME: is this a CUDA 9 compiler bug?
-      // K& ka = listK[pos];
-      // K& kb = listK[pos + stride];
-
-      // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-      // swap(s, ka, kb);
-
-      // V& va = listV[pos];
-      // V& vb = listV[pos + stride];
-      // swap(s, va, vb);
-    }
-
-    __syncthreads();
-  }
-}
-
-// Merge pairs of sorted lists larger than blockDim.x (NumThreads)
-template <int NumThreads,
-          typename K,
-          typename V,
-          int L,
-          bool Dir,
-          typename Comp,
-          bool FullMerge>
-inline __device__ void blockMergeLarge(K* listK, V* listV) {
-  static_assert(utils::isPowerOf2(L), "L must be a power-of-2");
-  static_assert(L >= kWarpSize, "merge list size must be >= 32");
-  static_assert(utils::isPowerOf2(NumThreads),
-                "NumThreads must be a power-of-2");
-  static_assert(L >= NumThreads, "merge list size must be >= NumThreads");
-
-  // For L > NumThreads, each thread has to perform more work
-  // per each stride.
-  constexpr int kLoopPerThread = L / NumThreads;
-
-  // It's not a bitonic merge, both lists are in the same direction,
-  // so handle the first swap assuming the second list is reversed
-#pragma unroll
-  for (int loop = 0; loop < kLoopPerThread; ++loop) {
-    int tid = loop * NumThreads + threadIdx.x;
+    // It's not a bitonic merge, both lists are in the same direction,
+    // so handle the first swap assuming the second list is reversed
     int pos = L - 1 - tid;
     int stride = 2 * tid + 1;
 
-   K ka = listK[pos];
-    K kb = listK[pos + stride];
+    if (AllThreads || (threadIdx.x < N * L)) {
+        K ka = listK[pos];
+        K kb = listK[pos + stride];
 
-    bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-    listK[pos] = swap ? kb : ka;
-    listK[pos + stride] = swap ? ka : kb;
+        bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+        listK[pos] = swap ? kb : ka;
+        listK[pos + stride] = swap ? ka : kb;
 
-    V va = listV[pos];
-    V vb = listV[pos + stride];
-    listV[pos] = swap ? vb : va;
-    listV[pos + stride] = swap ? va : vb;
+        V va = listV[pos];
+        V vb = listV[pos + stride];
+        listV[pos] = swap ? vb : va;
+        listV[pos + stride] = swap ? va : vb;
 
-    // FIXME: is this a CUDA 9 compiler bug?
-    // K& ka = listK[pos];
-    // K& kb = listK[pos + stride];
+        // FIXME: is this a CUDA 9 compiler bug?
+        // K& ka = listK[pos];
+        // K& kb = listK[pos + stride];
 
-    // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-    // swap(s, ka, kb);
+        // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+        // swap(s, ka, kb);
 
-    // V& va = listV[pos];
-    // V& vb = listV[pos + stride];
-    // swap(s, va, vb);
-  }
-
-  __syncthreads();
-
-  constexpr int kSecondLoopPerThread =
-    FullMerge ? kLoopPerThread : kLoopPerThread / 2;
-
-#pragma unroll
-  for (int stride = L / 2; stride > 0; stride /= 2) {
-#pragma unroll
-    for (int loop = 0; loop < kSecondLoopPerThread; ++loop) {
-      int tid = loop * NumThreads + threadIdx.x;
-      int pos = 2 * tid - (tid & (stride - 1));
-
-      K ka = listK[pos];
-      K kb = listK[pos + stride];
-
-      bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-      listK[pos] = swap ? kb : ka;
-      listK[pos + stride] = swap ? ka : kb;
-
-      V va = listV[pos];
-      V vb = listV[pos + stride];
-      listV[pos] = swap ? vb : va;
-      listV[pos + stride] = swap ? va : vb;
-
-      // FIXME: is this a CUDA 9 compiler bug?
-      // K& ka = listK[pos];
-      // K& kb = listK[pos + stride];
-
-      // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-      // swap(s, ka, kb);
-
-      // V& va = listV[pos];
-      // V& vb = listV[pos + stride];
-      // swap(s, va, vb);
+        // V& va = listV[pos];
+        // V& vb = listV[pos + stride];
+        // swap(s, va, vb);
     }
 
     __syncthreads();
-  }
+
+#pragma unroll
+    for (int stride = L / 2; stride > 0; stride /= 2) {
+        int pos = 2 * tid - (tid & (stride - 1));
+
+        if (AllThreads || (threadIdx.x < N * L)) {
+            K ka = listK[pos];
+            K kb = listK[pos + stride];
+
+            bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+            listK[pos] = swap ? kb : ka;
+            listK[pos + stride] = swap ? ka : kb;
+
+            V va = listV[pos];
+            V vb = listV[pos + stride];
+            listV[pos] = swap ? vb : va;
+            listV[pos + stride] = swap ? va : vb;
+
+            // FIXME: is this a CUDA 9 compiler bug?
+            // K& ka = listK[pos];
+            // K& kb = listK[pos + stride];
+
+            // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+            // swap(s, ka, kb);
+
+            // V& va = listV[pos];
+            // V& vb = listV[pos + stride];
+            // swap(s, va, vb);
+        }
+
+        __syncthreads();
+    }
+}
+
+// Merge pairs of sorted lists larger than blockDim.x (NumThreads)
+template <
+        int NumThreads,
+        typename K,
+        typename V,
+        int L,
+        bool Dir,
+        typename Comp,
+        bool FullMerge>
+inline __device__ void blockMergeLarge(K* listK, V* listV) {
+    static_assert(utils::isPowerOf2(L), "L must be a power-of-2");
+    static_assert(L >= kWarpSize, "merge list size must be >= 32");
+    static_assert(
+            utils::isPowerOf2(NumThreads), "NumThreads must be a power-of-2");
+    static_assert(L >= NumThreads, "merge list size must be >= NumThreads");
+
+    // For L > NumThreads, each thread has to perform more work
+    // per each stride.
+    constexpr int kLoopPerThread = L / NumThreads;
+
+    // It's not a bitonic merge, both lists are in the same direction,
+    // so handle the first swap assuming the second list is reversed
+#pragma unroll
+    for (int loop = 0; loop < kLoopPerThread; ++loop) {
+        int tid = loop * NumThreads + threadIdx.x;
+        int pos = L - 1 - tid;
+        int stride = 2 * tid + 1;
+
+        K ka = listK[pos];
+        K kb = listK[pos + stride];
+
+        bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+        listK[pos] = swap ? kb : ka;
+        listK[pos + stride] = swap ? ka : kb;
+
+        V va = listV[pos];
+        V vb = listV[pos + stride];
+        listV[pos] = swap ? vb : va;
+        listV[pos + stride] = swap ? va : vb;
+
+        // FIXME: is this a CUDA 9 compiler bug?
+        // K& ka = listK[pos];
+        // K& kb = listK[pos + stride];
+
+        // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+        // swap(s, ka, kb);
+
+        // V& va = listV[pos];
+        // V& vb = listV[pos + stride];
+        // swap(s, va, vb);
+    }
+
+    __syncthreads();
+
+    constexpr int kSecondLoopPerThread =
+            FullMerge ? kLoopPerThread : kLoopPerThread / 2;
+
+#pragma unroll
+    for (int stride = L / 2; stride > 0; stride /= 2) {
+#pragma unroll
+        for (int loop = 0; loop < kSecondLoopPerThread; ++loop) {
+            int tid = loop * NumThreads + threadIdx.x;
+            int pos = 2 * tid - (tid & (stride - 1));
+
+            K ka = listK[pos];
+            K kb = listK[pos + stride];
+
+            bool swap = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+            listK[pos] = swap ? kb : ka;
+            listK[pos + stride] = swap ? ka : kb;
+
+            V va = listV[pos];
+            V vb = listV[pos + stride];
+            listV[pos] = swap ? vb : va;
+            listV[pos + stride] = swap ? va : vb;
+
+            // FIXME: is this a CUDA 9 compiler bug?
+            // K& ka = listK[pos];
+            // K& kb = listK[pos + stride];
+
+            // bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+            // swap(s, ka, kb);
+
+            // V& va = listV[pos];
+            // V& vb = listV[pos + stride];
+            // swap(s, va, vb);
+        }
+
+        __syncthreads();
+    }
 }
 
 /// Class template to prevent static_assert from firing for
 /// mixing smaller/larger than block cases
-template <int NumThreads,
-          typename K,
-          typename V,
-          int N,
-          int L,
-          bool Dir,
-          typename Comp,
-          bool SmallerThanBlock,
-          bool FullMerge>
-struct BlockMerge {
-};
+template <
+        int NumThreads,
+        typename K,
+        typename V,
+        int N,
+        int L,
+        bool Dir,
+        typename Comp,
+        bool SmallerThanBlock,
+        bool FullMerge>
+struct BlockMerge {};
 
 /// Merging lists smaller than a block
-template <int NumThreads,
-          typename K,
-          typename V,
-          int N,
-          int L,
-          bool Dir,
-          typename Comp,
-          bool FullMerge>
+template <
+        int NumThreads,
+        typename K,
+        typename V,
+        int N,
+        int L,
+        bool Dir,
+        typename Comp,
+        bool FullMerge>
 struct BlockMerge<NumThreads, K, V, N, L, Dir, Comp, true, FullMerge> {
-  static inline __device__ void merge(K* listK, V* listV) {
-    constexpr int kNumParallelMerges = NumThreads / L;
-    constexpr int kNumIterations = N / kNumParallelMerges;
+    static inline __device__ void merge(K* listK, V* listV) {
+        constexpr int kNumParallelMerges = NumThreads / L;
+        constexpr int kNumIterations = N / kNumParallelMerges;
 
-    static_assert(L <= NumThreads, "list must be <= NumThreads");
-    static_assert((N < kNumParallelMerges) ||
-                  (kNumIterations * kNumParallelMerges == N),
-                  "improper selection of N and L");
+        static_assert(L <= NumThreads, "list must be <= NumThreads");
+        static_assert(
+                (N < kNumParallelMerges) ||
+                        (kNumIterations * kNumParallelMerges == N),
+                "improper selection of N and L");
 
-    if (N < kNumParallelMerges) {
-      // We only need L threads per each list to perform the merge
-      blockMergeSmall<NumThreads, K, V, N, L, false, Dir, Comp, FullMerge>(
-        listK, listV);
-    } else {
-      // All threads participate
+        if (N < kNumParallelMerges) {
+            // We only need L threads per each list to perform the merge
+            blockMergeSmall<
+                    NumThreads,
+                    K,
+                    V,
+                    N,
+                    L,
+                    false,
+                    Dir,
+                    Comp,
+                    FullMerge>(listK, listV);
+        } else {
+            // All threads participate
 #pragma unroll
-      for (int i = 0; i < kNumIterations; ++i) {
-        int start = i * kNumParallelMerges * 2 * L;
+            for (int i = 0; i < kNumIterations; ++i) {
+                int start = i * kNumParallelMerges * 2 * L;
 
-        blockMergeSmall<NumThreads, K, V, N, L, true, Dir, Comp, FullMerge>(
-          listK + start, listV + start);
-      }
+                blockMergeSmall<
+                        NumThreads,
+                        K,
+                        V,
+                        N,
+                        L,
+                        true,
+                        Dir,
+                        Comp,
+                        FullMerge>(listK + start, listV + start);
+            }
+        }
     }
-  }
 };
 
 /// Merging lists larger than a block
-template <int NumThreads,
-          typename K,
-          typename V,
-          int N,
-          int L,
-          bool Dir,
-          typename Comp,
-          bool FullMerge>
+template <
+        int NumThreads,
+        typename K,
+        typename V,
+        int N,
+        int L,
+        bool Dir,
+        typename Comp,
+        bool FullMerge>
 struct BlockMerge<NumThreads, K, V, N, L, Dir, Comp, false, FullMerge> {
-  static inline __device__ void merge(K* listK, V* listV) {
-    // Each pair of lists is merged sequentially
+    static inline __device__ void merge(K* listK, V* listV) {
+        // Each pair of lists is merged sequentially
 #pragma unroll
-    for (int i = 0; i < N; ++i) {
-      int start = i * 2 * L;
+        for (int i = 0; i < N; ++i) {
+            int start = i * 2 * L;
 
-      blockMergeLarge<NumThreads, K, V, L, Dir, Comp, FullMerge>(
-        listK + start, listV + start);
+            blockMergeLarge<NumThreads, K, V, L, Dir, Comp, FullMerge>(
+                    listK + start, listV + start);
+        }
     }
-  }
 };
 
-template <int NumThreads,
-          typename K,
-          typename V,
-          int N,
-          int L,
-          bool Dir,
-          typename Comp,
-          bool FullMerge = true>
+template <
+        int NumThreads,
+        typename K,
+        typename V,
+        int N,
+        int L,
+        bool Dir,
+        typename Comp,
+        bool FullMerge = true>
 inline __device__ void blockMerge(K* listK, V* listV) {
-  constexpr bool kSmallerThanBlock = (L <= NumThreads);
+    constexpr bool kSmallerThanBlock = (L <= NumThreads);
 
-  BlockMerge<NumThreads, K, V, N, L, Dir, Comp, kSmallerThanBlock, FullMerge>::
-    merge(listK, listV);
+    BlockMerge<
+            NumThreads,
+            K,
+            V,
+            N,
+            L,
+            Dir,
+            Comp,
+            kSmallerThanBlock,
+            FullMerge>::merge(listK, listV);
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/MergeNetworkUtils.cuh

@@ -7,18 +7,20 @@
 
 #pragma once
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T>
 inline __device__ void swap(bool swap, T& x, T& y) {
-  T tmp = x;
-  x = swap ? y : x;
-  y = swap ? tmp : y;
+    T tmp = x;
+    x = swap ? y : x;
+    y = swap ? tmp : y;
 }
 
 template <typename T>
 inline __device__ void assign(bool assign, T& x, T y) {
-  x = assign ? y : x;
+    x = assign ? y : x;
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/MergeNetworkWarp.cuh

@@ -7,13 +7,14 @@
 
 #pragma once
 
+#include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
 #include <faiss/gpu/utils/MergeNetworkUtils.cuh>
 #include <faiss/gpu/utils/PtxUtils.cuh>
-#include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/gpu/utils/WarpShuffles.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 //
 // This file contains functions to:
@@ -83,66 +84,76 @@ namespace faiss { namespace gpu {
 //
 // If IsBitonic is false, the first stage is reversed, so we don't
 // need to sort directionally. It's still technically a bitonic sort.
-template <typename K, typename V, int L,
-          bool Dir, typename Comp, bool IsBitonic>
+template <
+        typename K,
+        typename V,
+        int L,
+        bool Dir,
+        typename Comp,
+        bool IsBitonic>
 inline __device__ void warpBitonicMergeLE16(K& k, V& v) {
-  static_assert(utils::isPowerOf2(L), "L must be a power-of-2");
-  static_assert(L <= kWarpSize / 2, "merge list size must be <= 16");
+    static_assert(utils::isPowerOf2(L), "L must be a power-of-2");
+    static_assert(L <= kWarpSize / 2, "merge list size must be <= 16");
 
-  int laneId = getLaneId();
+    int laneId = getLaneId();
 
-  if (!IsBitonic) {
-    // Reverse the first comparison stage.
-    // For example, merging a list of size 8 has the exchanges:
-    // 0 <-> 15, 1 <-> 14, ...
-    K otherK = shfl_xor(k, 2 * L - 1);
-    V otherV = shfl_xor(v, 2 * L - 1);
+    if (!IsBitonic) {
+        // Reverse the first comparison stage.
+        // For example, merging a list of size 8 has the exchanges:
+        // 0 <-> 15, 1 <-> 14, ...
+        K otherK = shfl_xor(k, 2 * L - 1);
+        V otherV = shfl_xor(v, 2 * L - 1);
 
-    // Whether we are the lesser thread in the exchange
-    bool small = !(laneId & L);
+        // Whether we are the lesser thread in the exchange
+        bool small = !(laneId & L);
 
-    if (Dir) {
-      // See the comment above how performing both of these
-      // comparisons in the warp seems to win out over the
-      // alternatives in practice
-      bool s = small ? Comp::gt(k, otherK) : Comp::lt(k, otherK);
-      assign(s, k, otherK);
-      assign(s, v, otherV);
+        if (Dir) {
+            // See the comment above how performing both of these
+            // comparisons in the warp seems to win out over the
+            // alternatives in practice
+            bool s = small ? Comp::gt(k, otherK) : Comp::lt(k, otherK);
+            assign(s, k, otherK);
+            assign(s, v, otherV);
 
-    } else {
-      bool s = small ? Comp::lt(k, otherK) : Comp::gt(k, otherK);
-      assign(s, k, otherK);
-      assign(s, v, otherV);
+        } else {
+            bool s = small ? Comp::lt(k, otherK) : Comp::gt(k, otherK);
+            assign(s, k, otherK);
+            assign(s, v, otherV);
+        }
     }
-  }
 
 #pragma unroll
-  for (int stride = IsBitonic ? L : L / 2; stride > 0; stride /= 2) {
-    K otherK = shfl_xor(k, stride);
-    V otherV = shfl_xor(v, stride);
+    for (int stride = IsBitonic ? L : L / 2; stride > 0; stride /= 2) {
+        K otherK = shfl_xor(k, stride);
+        V otherV = shfl_xor(v, stride);
 
-    // Whether we are the lesser thread in the exchange
-    bool small = !(laneId & stride);
+        // Whether we are the lesser thread in the exchange
+        bool small = !(laneId & stride);
 
-    if (Dir) {
-      bool s = small ? Comp::gt(k, otherK) : Comp::lt(k, otherK);
-      assign(s, k, otherK);
-      assign(s, v, otherV);
+        if (Dir) {
+            bool s = small ? Comp::gt(k, otherK) : Comp::lt(k, otherK);
+            assign(s, k, otherK);
+            assign(s, v, otherV);
 
-    } else {
-      bool s = small ? Comp::lt(k, otherK) : Comp::gt(k, otherK);
-      assign(s, k, otherK);
-      assign(s, v, otherV);
+        } else {
+            bool s = small ? Comp::lt(k, otherK) : Comp::gt(k, otherK);
+            assign(s, k, otherK);
+            assign(s, v, otherV);
+        }
     }
-  }
 }
 
 // Template for performing a bitonic merge of an arbitrary set of
 // registers
-template <typename K, typename V, int N,
-          bool Dir, typename Comp, bool Low, bool Pow2>
-struct BitonicMergeStep {
-};
+template <
+        typename K,
+        typename V,
+        int N,
+        bool Dir,
+        typename Comp,
+        bool Low,
+        bool Pow2>
+struct BitonicMergeStep {};
 
 //
 // Power-of-2 merge specialization
@@ -151,69 +162,71 @@ struct BitonicMergeStep {
 // All merges eventually call this
 template <typename K, typename V, bool Dir, typename Comp, bool Low>
 struct BitonicMergeStep<K, V, 1, Dir, Comp, Low, true> {
-  static inline __device__ void merge(K k[1], V v[1]) {
-    // Use warp shuffles
-    warpBitonicMergeLE16<K, V, 16, Dir, Comp, true>(k[0], v[0]);
-  }
+    static inline __device__ void merge(K k[1], V v[1]) {
+        // Use warp shuffles
+        warpBitonicMergeLE16<K, V, 16, Dir, Comp, true>(k[0], v[0]);
+    }
 };
 
 template <typename K, typename V, int N, bool Dir, typename Comp, bool Low>
 struct BitonicMergeStep<K, V, N, Dir, Comp, Low, true> {
-  static inline __device__ void merge(K k[N], V v[N]) {
-    static_assert(utils::isPowerOf2(N), "must be power of 2");
-    static_assert(N > 1, "must be N > 1");
+    static inline __device__ void merge(K k[N], V v[N]) {
+        static_assert(utils::isPowerOf2(N), "must be power of 2");
+        static_assert(N > 1, "must be N > 1");
 
 #pragma unroll
-    for (int i = 0; i < N / 2; ++i) {
-      K& ka = k[i];
-      V& va = v[i];
+        for (int i = 0; i < N / 2; ++i) {
+            K& ka = k[i];
+            V& va = v[i];
 
-      K& kb = k[i + N / 2];
-      V& vb = v[i + N / 2];
+            K& kb = k[i + N / 2];
+            V& vb = v[i + N / 2];
 
-      bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-      swap(s, ka, kb);
-      swap(s, va, vb);
+            bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+            swap(s, ka, kb);
+            swap(s, va, vb);
+        }
+
+        {
+            K newK[N / 2];
+            V newV[N / 2];
+
+#pragma unroll
+            for (int i = 0; i < N / 2; ++i) {
+                newK[i] = k[i];
+                newV[i] = v[i];
+            }
+
+            BitonicMergeStep<K, V, N / 2, Dir, Comp, true, true>::merge(
+                    newK, newV);
+
+#pragma unroll
+            for (int i = 0; i < N / 2; ++i) {
+                k[i] = newK[i];
+                v[i] = newV[i];
+            }
+        }
+
+        {
+            K newK[N / 2];
+            V newV[N / 2];
+
+#pragma unroll
+            for (int i = 0; i < N / 2; ++i) {
+                newK[i] = k[i + N / 2];
+                newV[i] = v[i + N / 2];
+            }
+
+            BitonicMergeStep<K, V, N / 2, Dir, Comp, false, true>::merge(
+                    newK, newV);
+
+#pragma unroll
+            for (int i = 0; i < N / 2; ++i) {
+                k[i + N / 2] = newK[i];
+                v[i + N / 2] = newV[i];
+            }
+        }
     }
-
-    {
-      K newK[N / 2];
-      V newV[N / 2];
-
-#pragma unroll
-      for (int i = 0; i < N / 2; ++i) {
-        newK[i] = k[i];
-        newV[i] = v[i];
-      }
-
-      BitonicMergeStep<K, V, N / 2, Dir, Comp, true, true>::merge(newK, newV);
-
-#pragma unroll
-      for (int i = 0; i < N / 2; ++i) {
-        k[i] = newK[i];
-        v[i] = newV[i];
-      }
-    }
-
-    {
-      K newK[N / 2];
-      V newV[N / 2];
-
-#pragma unroll
-      for (int i = 0; i < N / 2; ++i) {
-        newK[i] = k[i + N / 2];
-        newV[i] = v[i + N / 2];
-      }
-
-      BitonicMergeStep<K, V, N / 2, Dir, Comp, false, true>::merge(newK, newV);
-
-#pragma unroll
-      for (int i = 0; i < N / 2; ++i) {
-        k[i + N / 2] = newK[i];
-        v[i + N / 2] = newV[i];
-      }
-    }
-  }
 };
 
 //
@@ -223,288 +236,315 @@ struct BitonicMergeStep<K, V, N, Dir, Comp, Low, true> {
 // Low recursion
 template <typename K, typename V, int N, bool Dir, typename Comp>
 struct BitonicMergeStep<K, V, N, Dir, Comp, true, false> {
-  static inline __device__ void merge(K k[N], V v[N]) {
-    static_assert(!utils::isPowerOf2(N), "must be non-power-of-2");
-    static_assert(N >= 3, "must be N >= 3");
+    static inline __device__ void merge(K k[N], V v[N]) {
+        static_assert(!utils::isPowerOf2(N), "must be non-power-of-2");
+        static_assert(N >= 3, "must be N >= 3");
 
-    constexpr int kNextHighestPowerOf2 = utils::nextHighestPowerOf2(N);
+        constexpr int kNextHighestPowerOf2 = utils::nextHighestPowerOf2(N);
 
 #pragma unroll
-    for (int i = 0; i < N - kNextHighestPowerOf2 / 2; ++i) {
-      K& ka = k[i];
-      V& va = v[i];
+        for (int i = 0; i < N - kNextHighestPowerOf2 / 2; ++i) {
+            K& ka = k[i];
+            V& va = v[i];
 
-      K& kb = k[i + kNextHighestPowerOf2 / 2];
-      V& vb = v[i + kNextHighestPowerOf2 / 2];
+            K& kb = k[i + kNextHighestPowerOf2 / 2];
+            V& vb = v[i + kNextHighestPowerOf2 / 2];
 
-      bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-      swap(s, ka, kb);
-      swap(s, va, vb);
+            bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+            swap(s, ka, kb);
+            swap(s, va, vb);
+        }
+
+        constexpr int kLowSize = N - kNextHighestPowerOf2 / 2;
+        constexpr int kHighSize = kNextHighestPowerOf2 / 2;
+        {
+            K newK[kLowSize];
+            V newV[kLowSize];
+
+#pragma unroll
+            for (int i = 0; i < kLowSize; ++i) {
+                newK[i] = k[i];
+                newV[i] = v[i];
+            }
+
+            constexpr bool kLowIsPowerOf2 =
+                    utils::isPowerOf2(N - kNextHighestPowerOf2 / 2);
+            // FIXME: compiler doesn't like this expression? compiler bug?
+            //      constexpr bool kLowIsPowerOf2 = utils::isPowerOf2(kLowSize);
+            BitonicMergeStep<
+                    K,
+                    V,
+                    kLowSize,
+                    Dir,
+                    Comp,
+                    true, // low
+                    kLowIsPowerOf2>::merge(newK, newV);
+
+#pragma unroll
+            for (int i = 0; i < kLowSize; ++i) {
+                k[i] = newK[i];
+                v[i] = newV[i];
+            }
+        }
+
+        {
+            K newK[kHighSize];
+            V newV[kHighSize];
+
+#pragma unroll
+            for (int i = 0; i < kHighSize; ++i) {
+                newK[i] = k[i + kLowSize];
+                newV[i] = v[i + kLowSize];
+            }
+
+            constexpr bool kHighIsPowerOf2 =
+                    utils::isPowerOf2(kNextHighestPowerOf2 / 2);
+            // FIXME: compiler doesn't like this expression? compiler bug?
+            //      constexpr bool kHighIsPowerOf2 =
+            //      utils::isPowerOf2(kHighSize);
+            BitonicMergeStep<
+                    K,
+                    V,
+                    kHighSize,
+                    Dir,
+                    Comp,
+                    false, // high
+                    kHighIsPowerOf2>::merge(newK, newV);
+
+#pragma unroll
+            for (int i = 0; i < kHighSize; ++i) {
+                k[i + kLowSize] = newK[i];
+                v[i + kLowSize] = newV[i];
+            }
+        }
     }
-
-    constexpr int kLowSize = N - kNextHighestPowerOf2 / 2;
-    constexpr int kHighSize = kNextHighestPowerOf2 / 2;
-    {
-      K newK[kLowSize];
-      V newV[kLowSize];
-
-#pragma unroll
-      for (int i = 0; i < kLowSize; ++i) {
-        newK[i] = k[i];
-        newV[i] = v[i];
-      }
-
-      constexpr bool kLowIsPowerOf2 =
-        utils::isPowerOf2(N - kNextHighestPowerOf2 / 2);
-      // FIXME: compiler doesn't like this expression? compiler bug?
-//      constexpr bool kLowIsPowerOf2 = utils::isPowerOf2(kLowSize);
-      BitonicMergeStep<K, V, kLowSize, Dir, Comp,
-                       true, // low
-                       kLowIsPowerOf2>::merge(newK, newV);
-
-#pragma unroll
-      for (int i = 0; i < kLowSize; ++i) {
-        k[i] = newK[i];
-        v[i] = newV[i];
-      }
-    }
-
-    {
-      K newK[kHighSize];
-      V newV[kHighSize];
-
-#pragma unroll
-      for (int i = 0; i < kHighSize; ++i) {
-        newK[i] = k[i + kLowSize];
-        newV[i] = v[i + kLowSize];
-      }
-
-      constexpr bool kHighIsPowerOf2 =
-        utils::isPowerOf2(kNextHighestPowerOf2 / 2);
-      // FIXME: compiler doesn't like this expression? compiler bug?
-//      constexpr bool kHighIsPowerOf2 = utils::isPowerOf2(kHighSize);
-      BitonicMergeStep<K, V, kHighSize, Dir, Comp,
-                       false, // high
-                       kHighIsPowerOf2>::merge(newK, newV);
-
-#pragma unroll
-      for (int i = 0; i < kHighSize; ++i) {
-        k[i + kLowSize] = newK[i];
-        v[i + kLowSize] = newV[i];
-      }
-    }
-  }
 };
 
 // High recursion
 template <typename K, typename V, int N, bool Dir, typename Comp>
 struct BitonicMergeStep<K, V, N, Dir, Comp, false, false> {
-  static inline __device__ void merge(K k[N], V v[N]) {
-    static_assert(!utils::isPowerOf2(N), "must be non-power-of-2");
-    static_assert(N >= 3, "must be N >= 3");
+    static inline __device__ void merge(K k[N], V v[N]) {
+        static_assert(!utils::isPowerOf2(N), "must be non-power-of-2");
+        static_assert(N >= 3, "must be N >= 3");
 
-    constexpr int kNextHighestPowerOf2 = utils::nextHighestPowerOf2(N);
+        constexpr int kNextHighestPowerOf2 = utils::nextHighestPowerOf2(N);
 
 #pragma unroll
-    for (int i = 0; i < N - kNextHighestPowerOf2 / 2; ++i) {
-      K& ka = k[i];
-      V& va = v[i];
+        for (int i = 0; i < N - kNextHighestPowerOf2 / 2; ++i) {
+            K& ka = k[i];
+            V& va = v[i];
 
-      K& kb = k[i + kNextHighestPowerOf2 / 2];
-      V& vb = v[i + kNextHighestPowerOf2 / 2];
+            K& kb = k[i + kNextHighestPowerOf2 / 2];
+            V& vb = v[i + kNextHighestPowerOf2 / 2];
 
-      bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
-      swap(s, ka, kb);
-      swap(s, va, vb);
+            bool s = Dir ? Comp::gt(ka, kb) : Comp::lt(ka, kb);
+            swap(s, ka, kb);
+            swap(s, va, vb);
+        }
+
+        constexpr int kLowSize = kNextHighestPowerOf2 / 2;
+        constexpr int kHighSize = N - kNextHighestPowerOf2 / 2;
+        {
+            K newK[kLowSize];
+            V newV[kLowSize];
+
+#pragma unroll
+            for (int i = 0; i < kLowSize; ++i) {
+                newK[i] = k[i];
+                newV[i] = v[i];
+            }
+
+            constexpr bool kLowIsPowerOf2 =
+                    utils::isPowerOf2(kNextHighestPowerOf2 / 2);
+            // FIXME: compiler doesn't like this expression? compiler bug?
+            //      constexpr bool kLowIsPowerOf2 = utils::isPowerOf2(kLowSize);
+            BitonicMergeStep<
+                    K,
+                    V,
+                    kLowSize,
+                    Dir,
+                    Comp,
+                    true, // low
+                    kLowIsPowerOf2>::merge(newK, newV);
+
+#pragma unroll
+            for (int i = 0; i < kLowSize; ++i) {
+                k[i] = newK[i];
+                v[i] = newV[i];
+            }
+        }
+
+        {
+            K newK[kHighSize];
+            V newV[kHighSize];
+
+#pragma unroll
+            for (int i = 0; i < kHighSize; ++i) {
+                newK[i] = k[i + kLowSize];
+                newV[i] = v[i + kLowSize];
+            }
+
+            constexpr bool kHighIsPowerOf2 =
+                    utils::isPowerOf2(N - kNextHighestPowerOf2 / 2);
+            // FIXME: compiler doesn't like this expression? compiler bug?
+            //      constexpr bool kHighIsPowerOf2 =
+            //      utils::isPowerOf2(kHighSize);
+            BitonicMergeStep<
+                    K,
+                    V,
+                    kHighSize,
+                    Dir,
+                    Comp,
+                    false, // high
+                    kHighIsPowerOf2>::merge(newK, newV);
+
+#pragma unroll
+            for (int i = 0; i < kHighSize; ++i) {
+                k[i + kLowSize] = newK[i];
+                v[i + kLowSize] = newV[i];
+            }
+        }
     }
-
-    constexpr int kLowSize = kNextHighestPowerOf2 / 2;
-    constexpr int kHighSize = N - kNextHighestPowerOf2 / 2;
-    {
-      K newK[kLowSize];
-      V newV[kLowSize];
-
-#pragma unroll
-      for (int i = 0; i < kLowSize; ++i) {
-        newK[i] = k[i];
-        newV[i] = v[i];
-      }
-
-      constexpr bool kLowIsPowerOf2 =
-        utils::isPowerOf2(kNextHighestPowerOf2 / 2);
-      // FIXME: compiler doesn't like this expression? compiler bug?
-//      constexpr bool kLowIsPowerOf2 = utils::isPowerOf2(kLowSize);
-      BitonicMergeStep<K, V, kLowSize, Dir, Comp,
-                       true, // low
-                       kLowIsPowerOf2>::merge(newK, newV);
-
-#pragma unroll
-      for (int i = 0; i < kLowSize; ++i) {
-        k[i] = newK[i];
-        v[i] = newV[i];
-      }
-    }
-
-    {
-      K newK[kHighSize];
-      V newV[kHighSize];
-
-#pragma unroll
-      for (int i = 0; i < kHighSize; ++i) {
-        newK[i] = k[i + kLowSize];
-        newV[i] = v[i + kLowSize];
-      }
-
-      constexpr bool kHighIsPowerOf2 =
-        utils::isPowerOf2(N - kNextHighestPowerOf2 / 2);
-      // FIXME: compiler doesn't like this expression? compiler bug?
-//      constexpr bool kHighIsPowerOf2 = utils::isPowerOf2(kHighSize);
-      BitonicMergeStep<K, V, kHighSize, Dir, Comp,
-                       false, // high
-                       kHighIsPowerOf2>::merge(newK, newV);
-
-#pragma unroll
-      for (int i = 0; i < kHighSize; ++i) {
-        k[i + kLowSize] = newK[i];
-        v[i + kLowSize] = newV[i];
-      }
-    }
-  }
 };
 
 /// Merges two sets of registers across the warp of any size;
 /// i.e., merges a sorted k/v list of size kWarpSize * N1 with a
 /// sorted k/v list of size kWarpSize * N2, where N1 and N2 are any
 /// value >= 1
-template <typename K,
-          typename V,
-          int N1,
-          int N2,
-          bool Dir,
-          typename Comp,
-          bool FullMerge = true>
-inline __device__ void warpMergeAnyRegisters(K k1[N1], V v1[N1],
-                                             K k2[N2], V v2[N2]) {
-  constexpr int kSmallestN = N1 < N2 ? N1 : N2;
+template <
+        typename K,
+        typename V,
+        int N1,
+        int N2,
+        bool Dir,
+        typename Comp,
+        bool FullMerge = true>
+inline __device__ void warpMergeAnyRegisters(
+        K k1[N1],
+        V v1[N1],
+        K k2[N2],
+        V v2[N2]) {
+    constexpr int kSmallestN = N1 < N2 ? N1 : N2;
 
 #pragma unroll
-  for (int i = 0; i < kSmallestN; ++i) {
-    K& ka = k1[N1 - 1 - i];
-    V& va = v1[N1 - 1 - i];
+    for (int i = 0; i < kSmallestN; ++i) {
+        K& ka = k1[N1 - 1 - i];
+        V& va = v1[N1 - 1 - i];
 
-    K& kb = k2[i];
-    V& vb = v2[i];
+        K& kb = k2[i];
+        V& vb = v2[i];
 
-    K otherKa;
-    V otherVa;
+        K otherKa;
+        V otherVa;
 
-    if (FullMerge) {
-      // We need the other values
-      otherKa = shfl_xor(ka, kWarpSize - 1);
-      otherVa = shfl_xor(va, kWarpSize - 1);
+        if (FullMerge) {
+            // We need the other values
+            otherKa = shfl_xor(ka, kWarpSize - 1);
+            otherVa = shfl_xor(va, kWarpSize - 1);
+        }
+
+        K otherKb = shfl_xor(kb, kWarpSize - 1);
+        V otherVb = shfl_xor(vb, kWarpSize - 1);
+
+        // ka is always first in the list, so we needn't use our lane
+        // in this comparison
+        bool swapa = Dir ? Comp::gt(ka, otherKb) : Comp::lt(ka, otherKb);
+        assign(swapa, ka, otherKb);
+        assign(swapa, va, otherVb);
+
+        // kb is always second in the list, so we needn't use our lane
+        // in this comparison
+        if (FullMerge) {
+            bool swapb = Dir ? Comp::lt(kb, otherKa) : Comp::gt(kb, otherKa);
+            assign(swapb, kb, otherKa);
+            assign(swapb, vb, otherVa);
+
+        } else {
+            // We don't care about updating elements in the second list
+        }
     }
 
-    K otherKb = shfl_xor(kb, kWarpSize - 1);
-    V otherVb = shfl_xor(vb, kWarpSize - 1);
-
-    // ka is always first in the list, so we needn't use our lane
-    // in this comparison
-    bool swapa = Dir ? Comp::gt(ka, otherKb) : Comp::lt(ka, otherKb);
-    assign(swapa, ka, otherKb);
-    assign(swapa, va, otherVb);
-
-    // kb is always second in the list, so we needn't use our lane
-    // in this comparison
+    BitonicMergeStep<K, V, N1, Dir, Comp, true, utils::isPowerOf2(N1)>::merge(
+            k1, v1);
     if (FullMerge) {
-      bool swapb = Dir ? Comp::lt(kb, otherKa) : Comp::gt(kb, otherKa);
-      assign(swapb, kb, otherKa);
-      assign(swapb, vb, otherVa);
-
-    } else {
-      // We don't care about updating elements in the second list
+        // Only if we care about N2 do we need to bother merging it fully
+        BitonicMergeStep<K, V, N2, Dir, Comp, false, utils::isPowerOf2(N2)>::
+                merge(k2, v2);
     }
-  }
-
-  BitonicMergeStep<K, V, N1, Dir, Comp,
-                   true, utils::isPowerOf2(N1)>::merge(k1, v1);
-  if (FullMerge) {
-    // Only if we care about N2 do we need to bother merging it fully
-    BitonicMergeStep<K, V, N2, Dir, Comp,
-                     false, utils::isPowerOf2(N2)>::merge(k2, v2);
-  }
 }
 
 // Recursive template that uses the above bitonic merge to perform a
 // bitonic sort
 template <typename K, typename V, int N, bool Dir, typename Comp>
 struct BitonicSortStep {
-  static inline __device__ void sort(K k[N], V v[N]) {
-    static_assert(N > 1, "did not hit specialized case");
+    static inline __device__ void sort(K k[N], V v[N]) {
+        static_assert(N > 1, "did not hit specialized case");
 
-    // Sort recursively
-    constexpr int kSizeA = N / 2;
-    constexpr int kSizeB = N - kSizeA;
+        // Sort recursively
+        constexpr int kSizeA = N / 2;
+        constexpr int kSizeB = N - kSizeA;
 
-    K aK[kSizeA];
-    V aV[kSizeA];
+        K aK[kSizeA];
+        V aV[kSizeA];
 
 #pragma unroll
-    for (int i = 0; i < kSizeA; ++i) {
-      aK[i] = k[i];
-      aV[i] = v[i];
-    }
+        for (int i = 0; i < kSizeA; ++i) {
+            aK[i] = k[i];
+            aV[i] = v[i];
+        }
 
-    BitonicSortStep<K, V, kSizeA, Dir, Comp>::sort(aK, aV);
+        BitonicSortStep<K, V, kSizeA, Dir, Comp>::sort(aK, aV);
 
-    K bK[kSizeB];
-    V bV[kSizeB];
+        K bK[kSizeB];
+        V bV[kSizeB];
 
 #pragma unroll
-    for (int i = 0; i < kSizeB; ++i) {
-      bK[i] = k[i + kSizeA];
-      bV[i] = v[i + kSizeA];
-    }
+        for (int i = 0; i < kSizeB; ++i) {
+            bK[i] = k[i + kSizeA];
+            bV[i] = v[i + kSizeA];
+        }
 
-    BitonicSortStep<K, V, kSizeB, Dir, Comp>::sort(bK, bV);
+        BitonicSortStep<K, V, kSizeB, Dir, Comp>::sort(bK, bV);
 
-    // Merge halves
-    warpMergeAnyRegisters<K, V, kSizeA, kSizeB, Dir, Comp>(aK, aV, bK, bV);
+        // Merge halves
+        warpMergeAnyRegisters<K, V, kSizeA, kSizeB, Dir, Comp>(aK, aV, bK, bV);
 
 #pragma unroll
-    for (int i = 0; i < kSizeA; ++i) {
-      k[i] = aK[i];
-      v[i] = aV[i];
-    }
+        for (int i = 0; i < kSizeA; ++i) {
+            k[i] = aK[i];
+            v[i] = aV[i];
+        }
 
 #pragma unroll
-    for (int i = 0; i < kSizeB; ++i) {
-      k[i + kSizeA] = bK[i];
-      v[i + kSizeA] = bV[i];
+        for (int i = 0; i < kSizeB; ++i) {
+            k[i + kSizeA] = bK[i];
+            v[i + kSizeA] = bV[i];
+        }
     }
-  }
 };
 
 // Single warp (N == 1) sorting specialization
 template <typename K, typename V, bool Dir, typename Comp>
 struct BitonicSortStep<K, V, 1, Dir, Comp> {
-  static inline __device__ void sort(K k[1], V v[1]) {
-    // Update this code if this changes
-    // should go from 1 -> kWarpSize in multiples of 2
-    static_assert(kWarpSize == 32, "unexpected warp size");
+    static inline __device__ void sort(K k[1], V v[1]) {
+        // Update this code if this changes
+        // should go from 1 -> kWarpSize in multiples of 2
+        static_assert(kWarpSize == 32, "unexpected warp size");
 
-    warpBitonicMergeLE16<K, V, 1, Dir, Comp, false>(k[0], v[0]);
-    warpBitonicMergeLE16<K, V, 2, Dir, Comp, false>(k[0], v[0]);
-    warpBitonicMergeLE16<K, V, 4, Dir, Comp, false>(k[0], v[0]);
-    warpBitonicMergeLE16<K, V, 8, Dir, Comp, false>(k[0], v[0]);
-    warpBitonicMergeLE16<K, V, 16, Dir, Comp, false>(k[0], v[0]);
-  }
+        warpBitonicMergeLE16<K, V, 1, Dir, Comp, false>(k[0], v[0]);
+        warpBitonicMergeLE16<K, V, 2, Dir, Comp, false>(k[0], v[0]);
+        warpBitonicMergeLE16<K, V, 4, Dir, Comp, false>(k[0], v[0]);
+        warpBitonicMergeLE16<K, V, 8, Dir, Comp, false>(k[0], v[0]);
+        warpBitonicMergeLE16<K, V, 16, Dir, Comp, false>(k[0], v[0]);
+    }
 };
 
 /// Sort a list of kWarpSize * N elements in registers, where N is an
 /// arbitrary >= 1
 template <typename K, typename V, int N, bool Dir, typename Comp>
 inline __device__ void warpSortAnyRegisters(K k[N], V v[N]) {
-  BitonicSortStep<K, V, N, Dir, Comp>::sort(k, v);
+    BitonicSortStep<K, V, N, Dir, Comp>::sort(k, v);
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/NoTypeTensor.cuh

@@ -5,119 +5,116 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/gpu/utils/Tensor.cuh>
 #include <initializer_list>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <int Dim, bool InnerContig = false, typename IndexT = int>
 class NoTypeTensor {
- public:
-  NoTypeTensor()
-      : mem_(nullptr),
-        typeSize_(0) {
-  }
+   public:
+    NoTypeTensor() : mem_(nullptr), typeSize_(0) {}
 
-  template <typename T>
-  NoTypeTensor(Tensor<T, Dim, InnerContig, IndexT>& t)
-      : mem_(t.data()),
-        typeSize_(sizeof(T)) {
-    for (int i = 0; i < Dim; ++i) {
-      size_[i] = t.getSize(i);
-      stride_[i] = t.getStride(i);
-    }
-  }
-
-  NoTypeTensor(void* mem, int typeSize, std::initializer_list<IndexT> sizes)
-      : mem_(mem),
-        typeSize_(typeSize) {
-
-    int i = 0;
-    for (auto s : sizes) {
-      size_[i++] = s;
+    template <typename T>
+    NoTypeTensor(Tensor<T, Dim, InnerContig, IndexT>& t)
+            : mem_(t.data()), typeSize_(sizeof(T)) {
+        for (int i = 0; i < Dim; ++i) {
+            size_[i] = t.getSize(i);
+            stride_[i] = t.getStride(i);
+        }
     }
 
-    stride_[Dim - 1] = (IndexT) 1;
-    for (int j = Dim - 2; j >= 0; --j) {
-      stride_[j] = stride_[j + 1] * size_[j + 1];
-    }
-  }
+    NoTypeTensor(void* mem, int typeSize, std::initializer_list<IndexT> sizes)
+            : mem_(mem), typeSize_(typeSize) {
+        int i = 0;
+        for (auto s : sizes) {
+            size_[i++] = s;
+        }
 
-  NoTypeTensor(void* mem, int typeSize, int sizes[Dim])
-      : mem_(mem),
-        typeSize_(typeSize) {
-    for (int i = 0; i < Dim; ++i) {
-      size_[i] = sizes[i];
+        stride_[Dim - 1] = (IndexT)1;
+        for (int j = Dim - 2; j >= 0; --j) {
+            stride_[j] = stride_[j + 1] * size_[j + 1];
+        }
     }
 
-    stride_[Dim - 1] = (IndexT) 1;
-    for (int i = Dim - 2; i >= 0; --i) {
-      stride_[i] = stride_[i + 1] * sizes[i + 1];
-    }
-  }
+    NoTypeTensor(void* mem, int typeSize, int sizes[Dim])
+            : mem_(mem), typeSize_(typeSize) {
+        for (int i = 0; i < Dim; ++i) {
+            size_[i] = sizes[i];
+        }
 
-  NoTypeTensor(void* mem, int typeSize,
-               IndexT sizes[Dim], IndexT strides[Dim])
-    : mem_(mem),
-      typeSize_(typeSize) {
-    for (int i = 0; i < Dim; ++i) {
-      size_[i] = sizes[i];
-      stride_[i] = strides[i];
-    }
-  }
-
-  int getTypeSize() const {
-    return typeSize_;
-  }
-
-  IndexT getSize(int dim) const {
-    FAISS_ASSERT(dim < Dim);
-    return size_[dim];
-  }
-
-  IndexT getStride(int dim) const {
-    FAISS_ASSERT(dim < Dim);
-    return stride_[dim];
-  }
-
-  template <typename T>
-  Tensor<T, Dim, InnerContig, IndexT> toTensor() {
-    FAISS_ASSERT(sizeof(T) == typeSize_);
-
-    return Tensor<T, Dim, InnerContig, IndexT>((T*) mem_, size_, stride_);
-  }
-
-  NoTypeTensor<Dim, InnerContig, IndexT> narrowOutermost(IndexT start,
-                                                         IndexT size) {
-    char* newPtr = (char*) mem_;
-
-    if (start > 0) {
-      newPtr += typeSize_ * start * stride_[0];
+        stride_[Dim - 1] = (IndexT)1;
+        for (int i = Dim - 2; i >= 0; --i) {
+            stride_[i] = stride_[i + 1] * sizes[i + 1];
+        }
     }
 
-    IndexT newSize[Dim];
-    for (int i = 0; i < Dim; ++i) {
-      if (i == 0) {
-        assert(start + size <= size_[0]);
-        newSize[i] = size;
-      } else {
-        newSize[i] = size_[i];
-      }
+    NoTypeTensor(
+            void* mem,
+            int typeSize,
+            IndexT sizes[Dim],
+            IndexT strides[Dim])
+            : mem_(mem), typeSize_(typeSize) {
+        for (int i = 0; i < Dim; ++i) {
+            size_[i] = sizes[i];
+            stride_[i] = strides[i];
+        }
     }
 
-    return NoTypeTensor<Dim, InnerContig, IndexT>(
-      newPtr, typeSize_, newSize, stride_);
-  }
+    int getTypeSize() const {
+        return typeSize_;
+    }
 
- private:
-  void* mem_;
-  int typeSize_;
-  IndexT size_[Dim];
-  IndexT stride_[Dim];
+    IndexT getSize(int dim) const {
+        FAISS_ASSERT(dim < Dim);
+        return size_[dim];
+    }
+
+    IndexT getStride(int dim) const {
+        FAISS_ASSERT(dim < Dim);
+        return stride_[dim];
+    }
+
+    template <typename T>
+    Tensor<T, Dim, InnerContig, IndexT> toTensor() {
+        FAISS_ASSERT(sizeof(T) == typeSize_);
+
+        return Tensor<T, Dim, InnerContig, IndexT>((T*)mem_, size_, stride_);
+    }
+
+    NoTypeTensor<Dim, InnerContig, IndexT> narrowOutermost(
+            IndexT start,
+            IndexT size) {
+        char* newPtr = (char*)mem_;
+
+        if (start > 0) {
+            newPtr += typeSize_ * start * stride_[0];
+        }
+
+        IndexT newSize[Dim];
+        for (int i = 0; i < Dim; ++i) {
+            if (i == 0) {
+                assert(start + size <= size_[0]);
+                newSize[i] = size;
+            } else {
+                newSize[i] = size_[i];
+            }
+        }
+
+        return NoTypeTensor<Dim, InnerContig, IndexT>(
+                newPtr, typeSize_, newSize, stride_);
+    }
+
+   private:
+    void* mem_;
+    int typeSize_;
+    IndexT size_[Dim];
+    IndexT stride_[Dim];
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/Pair.cuh

@@ -5,65 +5,62 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cuda.h>
 #include <faiss/gpu/utils/MathOperators.cuh>
 #include <faiss/gpu/utils/WarpShuffles.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// A simple pair type for CUDA device usage
 template <typename K, typename V>
 struct Pair {
-  constexpr __device__ inline Pair() {
-  }
+    constexpr __device__ inline Pair() {}
 
-  constexpr __device__ inline Pair(K key, V value)
-      : k(key), v(value) {
-  }
+    constexpr __device__ inline Pair(K key, V value) : k(key), v(value) {}
 
-  __device__ inline bool
-  operator==(const Pair<K, V>& rhs) const {
-    return Math<K>::eq(k, rhs.k) && Math<V>::eq(v, rhs.v);
-  }
+    __device__ inline bool operator==(const Pair<K, V>& rhs) const {
+        return Math<K>::eq(k, rhs.k) && Math<V>::eq(v, rhs.v);
+    }
 
-  __device__ inline bool
-  operator!=(const Pair<K, V>& rhs) const {
-    return !operator==(rhs);
-  }
+    __device__ inline bool operator!=(const Pair<K, V>& rhs) const {
+        return !operator==(rhs);
+    }
 
-  __device__ inline bool
-  operator<(const Pair<K, V>& rhs) const {
-    return Math<K>::lt(k, rhs.k) ||
-      (Math<K>::eq(k, rhs.k) && Math<V>::lt(v, rhs.v));
-  }
+    __device__ inline bool operator<(const Pair<K, V>& rhs) const {
+        return Math<K>::lt(k, rhs.k) ||
+                (Math<K>::eq(k, rhs.k) && Math<V>::lt(v, rhs.v));
+    }
 
-  __device__ inline bool
-  operator>(const Pair<K, V>& rhs) const {
-    return Math<K>::gt(k, rhs.k) ||
-      (Math<K>::eq(k, rhs.k) && Math<V>::gt(v, rhs.v));
-  }
+    __device__ inline bool operator>(const Pair<K, V>& rhs) const {
+        return Math<K>::gt(k, rhs.k) ||
+                (Math<K>::eq(k, rhs.k) && Math<V>::gt(v, rhs.v));
+    }
 
-  K k;
-  V v;
+    K k;
+    V v;
 };
 
 template <typename T, typename U>
-inline __device__ Pair<T, U> shfl_up(const Pair<T, U>& pair,
-                                     unsigned int delta,
-                                     int width = kWarpSize) {
-  return Pair<T, U>(shfl_up(pair.k, delta, width),
-                    shfl_up(pair.v, delta, width));
+inline __device__ Pair<T, U> shfl_up(
+        const Pair<T, U>& pair,
+        unsigned int delta,
+        int width = kWarpSize) {
+    return Pair<T, U>(
+            shfl_up(pair.k, delta, width), shfl_up(pair.v, delta, width));
 }
 
 template <typename T, typename U>
-inline __device__ Pair<T, U> shfl_xor(const Pair<T, U>& pair,
-                                      int laneMask,
-                                      int width = kWarpSize) {
-  return Pair<T, U>(shfl_xor(pair.k, laneMask, width),
-                    shfl_xor(pair.v, laneMask, width));
+inline __device__ Pair<T, U> shfl_xor(
+        const Pair<T, U>& pair,
+        int laneMask,
+        int width = kWarpSize) {
+    return Pair<T, U>(
+            shfl_xor(pair.k, laneMask, width),
+            shfl_xor(pair.v, laneMask, width));
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/PtxUtils.cuh

@@ -5,79 +5,88 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cuda.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // defines to simplify the SASS assembly structure file/line in the profiler
-#define GET_BITFIELD_U32(OUT, VAL, POS, LEN)                            \
-  asm("bfe.u32 %0, %1, %2, %3;" : "=r"(OUT) : "r"(VAL), "r"(POS), "r"(LEN));
+#define GET_BITFIELD_U32(OUT, VAL, POS, LEN) \
+    asm("bfe.u32 %0, %1, %2, %3;" : "=r"(OUT) : "r"(VAL), "r"(POS), "r"(LEN));
 
-#define GET_BITFIELD_U64(OUT, VAL, POS, LEN)                            \
-  asm("bfe.u64 %0, %1, %2, %3;" : "=l"(OUT) : "l"(VAL), "r"(POS), "r"(LEN));
+#define GET_BITFIELD_U64(OUT, VAL, POS, LEN) \
+    asm("bfe.u64 %0, %1, %2, %3;" : "=l"(OUT) : "l"(VAL), "r"(POS), "r"(LEN));
 
-__device__ __forceinline__
-unsigned int getBitfield(unsigned int val, int pos, int len) {
-  unsigned int ret;
-  asm("bfe.u32 %0, %1, %2, %3;" : "=r"(ret) : "r"(val), "r"(pos), "r"(len));
-  return ret;
+__device__ __forceinline__ unsigned int getBitfield(
+        unsigned int val,
+        int pos,
+        int len) {
+    unsigned int ret;
+    asm("bfe.u32 %0, %1, %2, %3;" : "=r"(ret) : "r"(val), "r"(pos), "r"(len));
+    return ret;
 }
 
-__device__ __forceinline__
-uint64_t getBitfield(uint64_t val, int pos, int len) {
-  uint64_t ret;
-  asm("bfe.u64 %0, %1, %2, %3;" : "=l"(ret) : "l"(val), "r"(pos), "r"(len));
-  return ret;
+__device__ __forceinline__ uint64_t
+getBitfield(uint64_t val, int pos, int len) {
+    uint64_t ret;
+    asm("bfe.u64 %0, %1, %2, %3;" : "=l"(ret) : "l"(val), "r"(pos), "r"(len));
+    return ret;
 }
 
-__device__ __forceinline__
-unsigned int setBitfield(unsigned int val,
-                         unsigned int toInsert, int pos, int len) {
-  unsigned int ret;
-  asm("bfi.b32 %0, %1, %2, %3, %4;" :
-      "=r"(ret) : "r"(toInsert), "r"(val), "r"(pos), "r"(len));
-  return ret;
+__device__ __forceinline__ unsigned int setBitfield(
+        unsigned int val,
+        unsigned int toInsert,
+        int pos,
+        int len) {
+    unsigned int ret;
+    asm("bfi.b32 %0, %1, %2, %3, %4;"
+        : "=r"(ret)
+        : "r"(toInsert), "r"(val), "r"(pos), "r"(len));
+    return ret;
 }
 
 __device__ __forceinline__ int getLaneId() {
-  int laneId;
-  asm("mov.u32 %0, %%laneid;" : "=r"(laneId) );
-  return laneId;
+    int laneId;
+    asm("mov.u32 %0, %%laneid;" : "=r"(laneId));
+    return laneId;
 }
 
 __device__ __forceinline__ unsigned getLaneMaskLt() {
-  unsigned mask;
-  asm("mov.u32 %0, %%lanemask_lt;" : "=r"(mask));
-  return mask;
+    unsigned mask;
+    asm("mov.u32 %0, %%lanemask_lt;" : "=r"(mask));
+    return mask;
 }
 
 __device__ __forceinline__ unsigned getLaneMaskLe() {
-  unsigned mask;
-  asm("mov.u32 %0, %%lanemask_le;" : "=r"(mask));
-  return mask;
+    unsigned mask;
+    asm("mov.u32 %0, %%lanemask_le;" : "=r"(mask));
+    return mask;
 }
 
 __device__ __forceinline__ unsigned getLaneMaskGt() {
-  unsigned mask;
-  asm("mov.u32 %0, %%lanemask_gt;" : "=r"(mask));
-  return mask;
+    unsigned mask;
+    asm("mov.u32 %0, %%lanemask_gt;" : "=r"(mask));
+    return mask;
 }
 
 __device__ __forceinline__ unsigned getLaneMaskGe() {
-  unsigned mask;
-  asm("mov.u32 %0, %%lanemask_ge;" : "=r"(mask));
-  return mask;
+    unsigned mask;
+    asm("mov.u32 %0, %%lanemask_ge;" : "=r"(mask));
+    return mask;
 }
 
 __device__ __forceinline__ void namedBarrierWait(int name, int numThreads) {
-  asm volatile("bar.sync %0, %1;" : : "r"(name), "r"(numThreads) : "memory");
+    asm volatile("bar.sync %0, %1;" : : "r"(name), "r"(numThreads) : "memory");
 }
 
 __device__ __forceinline__ void namedBarrierArrived(int name, int numThreads) {
-  asm volatile("bar.arrive %0, %1;" : : "r"(name), "r"(numThreads) : "memory");
+    asm volatile("bar.arrive %0, %1;"
+                 :
+                 : "r"(name), "r"(numThreads)
+                 : "memory");
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/ReductionOperators.cuh

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cuda.h>
@@ -13,61 +12,60 @@
 #include <faiss/gpu/utils/MathOperators.cuh>
 #include <faiss/gpu/utils/Pair.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T>
 struct Sum {
-  __device__ inline T operator()(T a, T b) const {
-    return Math<T>::add(a, b);
-  }
+    __device__ inline T operator()(T a, T b) const {
+        return Math<T>::add(a, b);
+    }
 
-  inline __device__ T identity() const {
-    return Math<T>::zero();
-  }
+    inline __device__ T identity() const {
+        return Math<T>::zero();
+    }
 };
 
 template <typename T>
 struct Min {
-  __device__ inline T operator()(T a, T b) const {
-    return Math<T>::lt(a, b) ? a : b;
-  }
+    __device__ inline T operator()(T a, T b) const {
+        return Math<T>::lt(a, b) ? a : b;
+    }
 
-  inline __device__ T identity() const {
-    return Limits<T>::getMax();
-  }
+    inline __device__ T identity() const {
+        return Limits<T>::getMax();
+    }
 };
 
 template <typename T>
 struct Max {
-  __device__ inline T operator()(T a, T b) const {
-    return Math<T>::gt(a, b) ? a : b;
-  }
+    __device__ inline T operator()(T a, T b) const {
+        return Math<T>::gt(a, b) ? a : b;
+    }
 
-  inline __device__ T identity() const {
-    return Limits<T>::getMin();
-  }
+    inline __device__ T identity() const {
+        return Limits<T>::getMin();
+    }
 };
 
 /// Used for producing segmented prefix scans; the value of the Pair
 /// denotes the start of a new segment for the scan
 template <typename T, typename ReduceOp>
 struct SegmentedReduce {
-  inline __device__ SegmentedReduce(const ReduceOp& o)
-      : op(o) {
-  }
+    inline __device__ SegmentedReduce(const ReduceOp& o) : op(o) {}
 
-  __device__
-  inline Pair<T, bool>
-  operator()(const Pair<T, bool>& a, const Pair<T, bool>& b) const {
-    return Pair<T, bool>(b.v ? b.k : op(a.k, b.k),
-                         a.v || b.v);
-  }
+    __device__ inline Pair<T, bool> operator()(
+            const Pair<T, bool>& a,
+            const Pair<T, bool>& b) const {
+        return Pair<T, bool>(b.v ? b.k : op(a.k, b.k), a.v || b.v);
+    }
 
-  inline __device__ Pair<T, bool> identity() const {
-    return Pair<T, bool>(op.identity(), false);
-  }
+    inline __device__ Pair<T, bool> identity() const {
+        return Pair<T, bool>(op.identity(), false);
+    }
 
-  ReduceOp op;
+    ReduceOp op;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/Reductions.cuh

@@ -5,138 +5,143 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
+#include <cuda.h>
+#include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
 #include <faiss/gpu/utils/PtxUtils.cuh>
 #include <faiss/gpu/utils/ReductionOperators.cuh>
-#include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/gpu/utils/WarpShuffles.cuh>
-#include <cuda.h>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T, typename Op, int ReduceWidth = kWarpSize>
 __device__ inline T warpReduceAll(T val, Op op) {
 #pragma unroll
-  for (int mask = ReduceWidth / 2; mask > 0; mask >>= 1) {
-    val = op(val, shfl_xor(val, mask));
-  }
+    for (int mask = ReduceWidth / 2; mask > 0; mask >>= 1) {
+        val = op(val, shfl_xor(val, mask));
+    }
 
-  return val;
+    return val;
 }
 
 /// Sums a register value across all warp threads
 template <typename T, int ReduceWidth = kWarpSize>
 __device__ inline T warpReduceAllSum(T val) {
-  return warpReduceAll<T, Sum<T>, ReduceWidth>(val, Sum<T>());
+    return warpReduceAll<T, Sum<T>, ReduceWidth>(val, Sum<T>());
 }
 
 /// Performs a block-wide reduction
 template <typename T, typename Op, bool BroadcastAll, bool KillWARDependency>
 __device__ inline T blockReduceAll(T val, Op op, T* smem) {
-  int laneId = getLaneId();
-  int warpId = threadIdx.x / kWarpSize;
+    int laneId = getLaneId();
+    int warpId = threadIdx.x / kWarpSize;
 
-  val = warpReduceAll<T, Op>(val, op);
-  if (laneId == 0) {
-    smem[warpId] = val;
-  }
-  __syncthreads();
-
-  if (warpId == 0) {
-    val = laneId < utils::divUp(blockDim.x, kWarpSize) ? smem[laneId] :
-      op.identity();
     val = warpReduceAll<T, Op>(val, op);
+    if (laneId == 0) {
+        smem[warpId] = val;
+    }
+    __syncthreads();
+
+    if (warpId == 0) {
+        val = laneId < utils::divUp(blockDim.x, kWarpSize) ? smem[laneId]
+                                                           : op.identity();
+        val = warpReduceAll<T, Op>(val, op);
+
+        if (BroadcastAll) {
+            __threadfence_block();
+
+            if (laneId == 0) {
+                smem[0] = val;
+            }
+        }
+    }
 
     if (BroadcastAll) {
-      __threadfence_block();
-
-      if (laneId == 0) {
-        smem[0] = val;
-      }
+        __syncthreads();
+        val = smem[0];
     }
-  }
 
-  if (BroadcastAll) {
-    __syncthreads();
-    val = smem[0];
-  }
+    if (KillWARDependency) {
+        __syncthreads();
+    }
 
-  if (KillWARDependency) {
-    __syncthreads();
-  }
-
-  return val;
+    return val;
 }
 
 /// Performs a block-wide reduction of multiple values simultaneously
-template <int Num, typename T, typename Op, bool BroadcastAll, bool KillWARDependency>
+template <
+        int Num,
+        typename T,
+        typename Op,
+        bool BroadcastAll,
+        bool KillWARDependency>
 __device__ inline void blockReduceAll(T val[Num], Op op, T* smem) {
-  int laneId = getLaneId();
-  int warpId = threadIdx.x / kWarpSize;
+    int laneId = getLaneId();
+    int warpId = threadIdx.x / kWarpSize;
 
-#pragma unroll
-  for (int i = 0; i < Num; ++i) {
-    val[i] = warpReduceAll<T, Op>(val[i], op);
-  }
-
-  if (laneId == 0) {
 #pragma unroll
     for (int i = 0; i < Num; ++i) {
-      smem[warpId * Num + i] = val[i];
+        val[i] = warpReduceAll<T, Op>(val[i], op);
     }
-  }
 
-  __syncthreads();
-
-  if (warpId == 0) {
+    if (laneId == 0) {
 #pragma unroll
-    for (int i = 0; i < Num; ++i) {
-      val[i] =
-        laneId < utils::divUp(blockDim.x, kWarpSize) ? smem[laneId * Num + i] :
-        op.identity();
-      val[i] = warpReduceAll<T, Op>(val[i], op);
+        for (int i = 0; i < Num; ++i) {
+            smem[warpId * Num + i] = val[i];
+        }
+    }
+
+    __syncthreads();
+
+    if (warpId == 0) {
+#pragma unroll
+        for (int i = 0; i < Num; ++i) {
+            val[i] = laneId < utils::divUp(blockDim.x, kWarpSize)
+                    ? smem[laneId * Num + i]
+                    : op.identity();
+            val[i] = warpReduceAll<T, Op>(val[i], op);
+        }
+
+        if (BroadcastAll) {
+            __threadfence_block();
+
+            if (laneId == 0) {
+#pragma unroll
+                for (int i = 0; i < Num; ++i) {
+                    smem[i] = val[i];
+                }
+            }
+        }
     }
 
     if (BroadcastAll) {
-      __threadfence_block();
-
-      if (laneId == 0) {
+        __syncthreads();
 #pragma unroll
         for (int i = 0; i < Num; ++i) {
-          smem[i] = val[i];
+            val[i] = smem[i];
         }
-      }
     }
-  }
 
-  if (BroadcastAll) {
-    __syncthreads();
-#pragma unroll
-    for (int i = 0; i < Num; ++i) {
-      val[i] = smem[i];
+    if (KillWARDependency) {
+        __syncthreads();
     }
-  }
-
-  if (KillWARDependency) {
-    __syncthreads();
-  }
 }
 
-
 /// Sums a register value across the entire block
 template <typename T, bool BroadcastAll, bool KillWARDependency>
 __device__ inline T blockReduceAllSum(T val, T* smem) {
-  return blockReduceAll<T, Sum<T>, BroadcastAll, KillWARDependency>(
-    val, Sum<T>(), smem);
+    return blockReduceAll<T, Sum<T>, BroadcastAll, KillWARDependency>(
+            val, Sum<T>(), smem);
 }
 
 template <int Num, typename T, bool BroadcastAll, bool KillWARDependency>
 __device__ inline void blockReduceAllSum(T vals[Num], T* smem) {
-  return blockReduceAll<Num, T, Sum<T>, BroadcastAll, KillWARDependency>(
-    vals, Sum<T>(), smem);
+    return blockReduceAll<Num, T, Sum<T>, BroadcastAll, KillWARDependency>(
+            vals, Sum<T>(), smem);
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/ThrustAllocator.cuh

@@ -5,77 +5,77 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/gpu/GpuResources.h>
 #include <cuda.h>
+#include <faiss/gpu/GpuResources.h>
 #include <unordered_set>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// Allocator for Thrust that comes out of a specified memory space
 class GpuResourcesThrustAllocator {
- public:
-  typedef char value_type;
+   public:
+    typedef char value_type;
 
-  inline GpuResourcesThrustAllocator(GpuResources* res,
-                                     cudaStream_t stream,
-                                     void* mem,
-                                     size_t size)
-      : res_(res),
-        stream_(stream),
-        start_((char*) mem),
-        cur_((char*) mem),
-        end_((char*) mem + size) {
-  }
+    inline GpuResourcesThrustAllocator(
+            GpuResources* res,
+            cudaStream_t stream,
+            void* mem,
+            size_t size)
+            : res_(res),
+              stream_(stream),
+              start_((char*)mem),
+              cur_((char*)mem),
+              end_((char*)mem + size) {}
 
-  inline ~GpuResourcesThrustAllocator() {
-    // In the case of an exception being thrown, we may not have called
-    // deallocate on all of our sub-allocations. Free them here
-    for (auto p : mallocAllocs_) {
-      res_->deallocMemory(getCurrentDevice(), p);
+    inline ~GpuResourcesThrustAllocator() {
+        // In the case of an exception being thrown, we may not have called
+        // deallocate on all of our sub-allocations. Free them here
+        for (auto p : mallocAllocs_) {
+            res_->deallocMemory(getCurrentDevice(), p);
+        }
     }
-  }
 
-  inline char* allocate(std::ptrdiff_t size) {
-    if (size <= (end_ - cur_)) {
-      char* p = cur_;
-      cur_ += size;
-      FAISS_ASSERT(cur_ <= end_);
+    inline char* allocate(std::ptrdiff_t size) {
+        if (size <= (end_ - cur_)) {
+            char* p = cur_;
+            cur_ += size;
+            FAISS_ASSERT(cur_ <= end_);
 
-      return p;
-    } else {
-      // FIXME: we cannot use temporary memory for new requests because the
-      // current temporary memory allocator cannot handle stream synchronization
-      // at present, so just allocate through the general device
-      char* p =
-        (char*) res_->allocMemory(
-          AllocRequest(makeDevAlloc(AllocType::Other, stream_), size));
+            return p;
+        } else {
+            // FIXME: we cannot use temporary memory for new requests because
+            // the current temporary memory allocator cannot handle stream
+            // synchronization at present, so just allocate through the general
+            // device
+            char* p = (char*)res_->allocMemory(AllocRequest(
+                    makeDevAlloc(AllocType::Other, stream_), size));
 
-      mallocAllocs_.insert(p);
-      return p;
+            mallocAllocs_.insert(p);
+            return p;
+        }
     }
-  }
 
-  inline void deallocate(char* p, size_t size) {
-    // Allocations could be returned out-of-order; ignore those we
-    // didn't cudaMalloc
-    auto it = mallocAllocs_.find(p);
-    if (it != mallocAllocs_.end()) {
-      res_->deallocMemory(getCurrentDevice(), p);
-      mallocAllocs_.erase(it);
+    inline void deallocate(char* p, size_t size) {
+        // Allocations could be returned out-of-order; ignore those we
+        // didn't cudaMalloc
+        auto it = mallocAllocs_.find(p);
+        if (it != mallocAllocs_.end()) {
+            res_->deallocMemory(getCurrentDevice(), p);
+            mallocAllocs_.erase(it);
+        }
     }
-  }
 
- private:
-  GpuResources* res_;
-  cudaStream_t stream_;
-  char* start_;
-  char* cur_;
-  char* end_;
-  std::unordered_set<char*> mallocAllocs_;
+   private:
+    GpuResources* res_;
+    cudaStream_t stream_;
+    char* start_;
+    char* cur_;
+    char* end_;
+    std::unordered_set<char*> mallocAllocs_;
 };
 
-
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/Transpose.cuh

@@ -5,105 +5,111 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/gpu/utils/Tensor.cuh>
+#include <cuda.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
 #include <faiss/gpu/utils/StaticUtils.h>
-#include <cuda.h>
+#include <faiss/impl/FaissAssert.h>
 #include <stdint.h>
+#include <faiss/gpu/utils/Tensor.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 template <typename T, typename IndexT>
 struct TensorInfo {
-  static constexpr int kMaxDims = 8;
+    static constexpr int kMaxDims = 8;
 
-  T* data;
-  IndexT sizes[kMaxDims];
-  IndexT strides[kMaxDims];
-  int dims;
+    T* data;
+    IndexT sizes[kMaxDims];
+    IndexT strides[kMaxDims];
+    int dims;
 };
 
 template <typename T, typename IndexT, int Dim>
 struct TensorInfoOffset {
-  __device__ inline static unsigned int get(const TensorInfo<T, IndexT>& info,
-                                            IndexT linearId) {
-    IndexT offset = 0;
+    __device__ inline static unsigned int get(
+            const TensorInfo<T, IndexT>& info,
+            IndexT linearId) {
+        IndexT offset = 0;
 
 #pragma unroll
-    for (int i = Dim - 1; i >= 0; --i) {
-      IndexT curDimIndex = linearId % info.sizes[i];
-      IndexT curDimOffset = curDimIndex * info.strides[i];
+        for (int i = Dim - 1; i >= 0; --i) {
+            IndexT curDimIndex = linearId % info.sizes[i];
+            IndexT curDimOffset = curDimIndex * info.strides[i];
 
-      offset += curDimOffset;
+            offset += curDimOffset;
 
-      if (i > 0) {
-        linearId /= info.sizes[i];
-      }
+            if (i > 0) {
+                linearId /= info.sizes[i];
+            }
+        }
+
+        return offset;
     }
-
-    return offset;
-  }
 };
 
 template <typename T, typename IndexT>
 struct TensorInfoOffset<T, IndexT, -1> {
-  __device__ inline static unsigned int get(const TensorInfo<T, IndexT>& info,
-                                            IndexT linearId) {
-    return linearId;
-  }
+    __device__ inline static unsigned int get(
+            const TensorInfo<T, IndexT>& info,
+            IndexT linearId) {
+        return linearId;
+    }
 };
 
 template <typename T, typename IndexT, int Dim>
 TensorInfo<T, IndexT> getTensorInfo(const Tensor<T, Dim, true>& t) {
-  TensorInfo<T, IndexT> info;
+    TensorInfo<T, IndexT> info;
 
-  for (int i = 0; i < Dim; ++i) {
-    info.sizes[i] = (IndexT) t.getSize(i);
-    info.strides[i] = (IndexT) t.getStride(i);
-  }
+    for (int i = 0; i < Dim; ++i) {
+        info.sizes[i] = (IndexT)t.getSize(i);
+        info.strides[i] = (IndexT)t.getStride(i);
+    }
 
-  info.data = t.data();
-  info.dims = Dim;
+    info.data = t.data();
+    info.dims = Dim;
 
-  return info;
+    return info;
 }
 
 template <typename T, typename IndexT, int DimInput, int DimOutput>
-__global__ void transposeAny(TensorInfo<T, IndexT> input,
-                             TensorInfo<T, IndexT> output,
-                             IndexT totalSize) {
-  for (IndexT i = blockIdx.x * blockDim.x + threadIdx.x;
-       i < totalSize;
-       i += gridDim.x * blockDim.x) {
-    auto inputOffset = TensorInfoOffset<T, IndexT, DimInput>::get(input, i);
-    auto outputOffset = TensorInfoOffset<T, IndexT, DimOutput>::get(output, i);
+__global__ void transposeAny(
+        TensorInfo<T, IndexT> input,
+        TensorInfo<T, IndexT> output,
+        IndexT totalSize) {
+    for (IndexT i = blockIdx.x * blockDim.x + threadIdx.x; i < totalSize;
+         i += gridDim.x * blockDim.x) {
+        auto inputOffset = TensorInfoOffset<T, IndexT, DimInput>::get(input, i);
+        auto outputOffset =
+                TensorInfoOffset<T, IndexT, DimOutput>::get(output, i);
 
 #if __CUDA_ARCH__ >= 350
-    output.data[outputOffset] = __ldg(&input.data[inputOffset]);
+        output.data[outputOffset] = __ldg(&input.data[inputOffset]);
 #else
-    output.data[outputOffset] = input.data[inputOffset];
+        output.data[outputOffset] = input.data[inputOffset];
 #endif
-  }
+    }
 }
 
 // Transpose contiguous t1 t2 i1 -> t2 t1 i1
 template <typename T, typename IndexT>
-__global__ void transposeOuter(const T* in,
-                               T* out,
-                               IndexT t1, IndexT t2, IndexT i1) {
-  IndexT gt1 = blockIdx.y;
-  IndexT gt2 = blockIdx.x;
+__global__ void transposeOuter(
+        const T* in,
+        T* out,
+        IndexT t1,
+        IndexT t2,
+        IndexT i1) {
+    IndexT gt1 = blockIdx.y;
+    IndexT gt2 = blockIdx.x;
 
-  in += i1 * (gt1 * t2 + gt2);
-  out += i1 * (gt2 * t1 + gt1);
+    in += i1 * (gt1 * t2 + gt2);
+    out += i1 * (gt2 * t1 + gt1);
 
-  for (IndexT i = threadIdx.x; i < i1; i += blockDim.x) {
-    out[i] = in[i];
-  }
+    for (IndexT i = threadIdx.x; i < i1; i += blockDim.x) {
+        out[i] = in[i];
+    }
 }
 
 /// Performs an out-of-place transposition between any two dimensions.
@@ -116,95 +122,101 @@ __global__ void transposeOuter(const T* in,
 /// especially for cases that we care about (outer dimension
 /// transpositions).
 template <typename T, int Dim>
-void runTransposeAny(Tensor<T, Dim, true>& in,
-                     int dim1, int dim2,
-                     Tensor<T, Dim, true>& out,
-                     cudaStream_t stream) {
-  static_assert(Dim <= TensorInfo<T, unsigned int>::kMaxDims,
-                "too many dimensions");
+void runTransposeAny(
+        Tensor<T, Dim, true>& in,
+        int dim1,
+        int dim2,
+        Tensor<T, Dim, true>& out,
+        cudaStream_t stream) {
+    static_assert(
+            Dim <= TensorInfo<T, unsigned int>::kMaxDims,
+            "too many dimensions");
 
-  FAISS_ASSERT(dim1 != dim2);
-  FAISS_ASSERT(dim1 < Dim && dim2 < Dim);
+    FAISS_ASSERT(dim1 != dim2);
+    FAISS_ASSERT(dim1 < Dim && dim2 < Dim);
 
-  // Rearrange dim1 and dim2 in increasing order in order to see if this is an
-  // outer dimension transposition (below)
-  if (dim1 > dim2) {
-    std::swap(dim1, dim2);
-  }
-
-  int outSize[Dim];
-
-  for (int i = 0; i < Dim; ++i) {
-    outSize[i] = in.getSize(i);
-  }
-
-  std::swap(outSize[dim1], outSize[dim2]);
-
-  for (int i = 0; i < Dim; ++i) {
-    FAISS_ASSERT(out.getSize(i) == outSize[i]);
-  }
-
-  auto maxThreads = getMaxThreadsCurrentDevice();
-  auto totalSize = in.numElements();
-
-  // Is this a transposition of the two outer dimensions?
-  bool isTransposeOuter = (Dim >= 3) && (dim1 == 0) && (dim2 == 1);
-  if (isTransposeOuter) {
-    // Outer dimension transposition only (there is a contiguous inner
-    // dimension)
-    size_t innerSize = 1;
-    for (int i = 2; i < Dim; ++i) {
-      innerSize *= in.getSize(i);
+    // Rearrange dim1 and dim2 in increasing order in order to see if this is an
+    // outer dimension transposition (below)
+    if (dim1 > dim2) {
+        std::swap(dim1, dim2);
     }
 
-    auto grid = dim3(in.getSize(1), in.getSize(0));
-    int block = (innerSize < maxThreads) ? innerSize : maxThreads;
+    int outSize[Dim];
 
-    if (totalSize <= (size_t) std::numeric_limits<int>::max()) {
-      transposeOuter<T, int32_t><<<grid, block, 0, stream>>>(in.data(),
-                                                             out.data(),
-                                                             in.getSize(0),
-                                                             in.getSize(1),
-                                                             innerSize);
+    for (int i = 0; i < Dim; ++i) {
+        outSize[i] = in.getSize(i);
+    }
+
+    std::swap(outSize[dim1], outSize[dim2]);
+
+    for (int i = 0; i < Dim; ++i) {
+        FAISS_ASSERT(out.getSize(i) == outSize[i]);
+    }
+
+    auto maxThreads = getMaxThreadsCurrentDevice();
+    auto totalSize = in.numElements();
+
+    // Is this a transposition of the two outer dimensions?
+    bool isTransposeOuter = (Dim >= 3) && (dim1 == 0) && (dim2 == 1);
+    if (isTransposeOuter) {
+        // Outer dimension transposition only (there is a contiguous inner
+        // dimension)
+        size_t innerSize = 1;
+        for (int i = 2; i < Dim; ++i) {
+            innerSize *= in.getSize(i);
+        }
+
+        auto grid = dim3(in.getSize(1), in.getSize(0));
+        int block = (innerSize < maxThreads) ? innerSize : maxThreads;
+
+        if (totalSize <= (size_t)std::numeric_limits<int>::max()) {
+            transposeOuter<T, int32_t><<<grid, block, 0, stream>>>(
+                    in.data(),
+                    out.data(),
+                    in.getSize(0),
+                    in.getSize(1),
+                    innerSize);
+        } else {
+            transposeOuter<T, int64_t><<<grid, block, 0, stream>>>(
+                    in.data(),
+                    out.data(),
+                    in.getSize(0),
+                    in.getSize(1),
+                    innerSize);
+        }
     } else {
-      transposeOuter<T, int64_t><<<grid, block, 0, stream>>>(in.data(),
-                                                             out.data(),
-                                                             in.getSize(0),
-                                                             in.getSize(1),
-                                                             innerSize);
+        int block = (totalSize < maxThreads) ? totalSize : maxThreads;
+
+        // Non-outer transposition
+        if (totalSize <= (size_t)std::numeric_limits<int>::max()) {
+            // General transposition
+            // div/mod seems faster with unsigned types
+            auto inInfo = getTensorInfo<T, uint32_t, Dim>(in);
+            auto outInfo = getTensorInfo<T, uint32_t, Dim>(out);
+
+            std::swap(inInfo.sizes[dim1], inInfo.sizes[dim2]);
+            std::swap(inInfo.strides[dim1], inInfo.strides[dim2]);
+
+            auto grid = std::min(utils::divUp(totalSize, block), (size_t)4096);
+
+            transposeAny<T, uint32_t, Dim, -1>
+                    <<<grid, block, 0, stream>>>(inInfo, outInfo, totalSize);
+        } else {
+            auto inInfo = getTensorInfo<T, uint64_t, Dim>(in);
+            auto outInfo = getTensorInfo<T, uint64_t, Dim>(out);
+
+            std::swap(inInfo.sizes[dim1], inInfo.sizes[dim2]);
+            std::swap(inInfo.strides[dim1], inInfo.strides[dim2]);
+
+            auto grid = std::min(utils::divUp(totalSize, block), (size_t)4096);
+
+            transposeAny<T, uint64_t, Dim, -1>
+                    <<<grid, block, 0, stream>>>(inInfo, outInfo, totalSize);
+        }
     }
-  } else {
-    int block = (totalSize < maxThreads) ? totalSize : maxThreads;
 
-    // Non-outer transposition
-    if (totalSize <= (size_t) std::numeric_limits<int>::max()) {
-      // General transposition
-      // div/mod seems faster with unsigned types
-      auto inInfo = getTensorInfo<T, uint32_t, Dim>(in);
-      auto outInfo = getTensorInfo<T, uint32_t, Dim>(out);
-
-      std::swap(inInfo.sizes[dim1], inInfo.sizes[dim2]);
-      std::swap(inInfo.strides[dim1], inInfo.strides[dim2]);
-
-      auto grid = std::min(utils::divUp(totalSize, block), (size_t) 4096);
-
-      transposeAny<T, uint32_t, Dim, -1>
-        <<<grid, block, 0, stream>>>(inInfo, outInfo, totalSize);
-    } else {
-      auto inInfo = getTensorInfo<T, uint64_t, Dim>(in);
-      auto outInfo = getTensorInfo<T, uint64_t, Dim>(out);
-
-      std::swap(inInfo.sizes[dim1], inInfo.sizes[dim2]);
-      std::swap(inInfo.strides[dim1], inInfo.strides[dim2]);
-
-      auto grid = std::min(utils::divUp(totalSize, block), (size_t) 4096);
-
-      transposeAny<T, uint64_t, Dim, -1>
-        <<<grid, block, 0, stream>>>(inInfo, outInfo, totalSize);
-    }
-  }
-
-  CUDA_TEST_ERROR();
+    CUDA_TEST_ERROR();
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/WarpPackedBits.cuh

@@ -5,13 +5,13 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <faiss/gpu/utils/PtxUtils.cuh>
 #include <faiss/gpu/utils/WarpShuffles.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 //
 // Warp-coalesced parallel reading and writing of packed bits
@@ -20,245 +20,262 @@ namespace faiss { namespace gpu {
 // Read/write native word sizes
 template <typename WordT, int Bits>
 struct WarpPackedBits {
-  static __device__ void write(int laneId, WordT v, bool valid, WordT* out) {
-    static_assert(sizeof(WordT) == Bits / 8 &&
-                  (Bits % 8) == 0, "");
-    // We can just write directly
-    if (valid) {
-      out[laneId] = v;
+    static __device__ void write(int laneId, WordT v, bool valid, WordT* out) {
+        static_assert(sizeof(WordT) == Bits / 8 && (Bits % 8) == 0, "");
+        // We can just write directly
+        if (valid) {
+            out[laneId] = v;
+        }
     }
-  }
 
-  static inline __device__ WordT read(int laneId, WordT* in) {
-    return in[laneId];
-  }
+    static inline __device__ WordT read(int laneId, WordT* in) {
+        return in[laneId];
+    }
 
-  static inline __device__ WordT postRead(int laneId, WordT v) {
-    return v;
-  }
+    static inline __device__ WordT postRead(int laneId, WordT v) {
+        return v;
+    }
 };
 
 // Read/write 6 bit fields, packed across the warp into 24 bytes
 template <>
 struct WarpPackedBits<uint8_t, 6> {
-  static __device__ void write(int laneId, uint8_t v, bool valid, uint8_t* out) {
-    // Lower 24 lanes wwrite out packed data
-    int laneFrom = (laneId * 8) / 6;
+    static __device__ void write(
+            int laneId,
+            uint8_t v,
+            bool valid,
+            uint8_t* out) {
+        // Lower 24 lanes wwrite out packed data
+        int laneFrom = (laneId * 8) / 6;
 
-    v = valid ? v : 0;
-    v &= 0x3f; // ensure we have only 6 bits
+        v = valid ? v : 0;
+        v &= 0x3f; // ensure we have only 6 bits
 
-    uint8_t vLower = (uint8_t) SHFL_SYNC((unsigned int) v, laneFrom, kWarpSize);
-    uint8_t vUpper = (uint8_t) SHFL_SYNC((unsigned int) v, laneFrom + 1, kWarpSize);
+        uint8_t vLower =
+                (uint8_t)SHFL_SYNC((unsigned int)v, laneFrom, kWarpSize);
+        uint8_t vUpper =
+                (uint8_t)SHFL_SYNC((unsigned int)v, laneFrom + 1, kWarpSize);
 
-    // lsb     ...    msb
-    // 0: 0 0 0 0 0 0 1 1
-    // 1: 1 1 1 1 2 2 2 2
-    // 2: 2 2 3 3 3 3 3 3
-    int typeLane = laneId % 3;
-    uint8_t vOut = 0;
-    switch (typeLane) {
-      case 0:
-        // 6 msbs of lower as vOut lsbs
-        // 2 lsbs of upper as vOut msbs
-        vOut = vLower | (vUpper << 6);
-        break;
-      case 1:
-        // 4 msbs of lower as vOut lsbs
-        // 4 lsbs of upper as vOut msbs
-        vOut = (vLower >> 2) | (vUpper << 4);
-        break;
-      case 2:
-        // 2 msbs of lower as vOut lsbs
-        // 6 lsbs of upper as vOut msbs
-        vOut = (vLower >> 4) | (vUpper << 2);
-        break;
+        // lsb     ...    msb
+        // 0: 0 0 0 0 0 0 1 1
+        // 1: 1 1 1 1 2 2 2 2
+        // 2: 2 2 3 3 3 3 3 3
+        int typeLane = laneId % 3;
+        uint8_t vOut = 0;
+        switch (typeLane) {
+            case 0:
+                // 6 msbs of lower as vOut lsbs
+                // 2 lsbs of upper as vOut msbs
+                vOut = vLower | (vUpper << 6);
+                break;
+            case 1:
+                // 4 msbs of lower as vOut lsbs
+                // 4 lsbs of upper as vOut msbs
+                vOut = (vLower >> 2) | (vUpper << 4);
+                break;
+            case 2:
+                // 2 msbs of lower as vOut lsbs
+                // 6 lsbs of upper as vOut msbs
+                vOut = (vLower >> 4) | (vUpper << 2);
+                break;
+        }
+
+        if (laneId < 24) {
+            // There could be prior data
+            out[laneId] |= vOut;
+        }
     }
 
-    if (laneId < 24) {
-      // There could be prior data
-      out[laneId] |= vOut;
-    }
-  }
+    static inline __device__ uint8_t read(int laneId, uint8_t* in) {
+        uint8_t v = 0;
 
-  static inline __device__ uint8_t read(int laneId, uint8_t* in) {
-    uint8_t v = 0;
+        if (laneId < 24) {
+            v = in[laneId];
+        }
 
-    if (laneId < 24) {
-      v = in[laneId];
+        return v;
     }
 
-    return v;
-  }
+    static inline __device__ uint8_t postRead(int laneId, uint8_t v) {
+        int laneFrom = (laneId * 6) / 8;
 
-  static inline __device__ uint8_t postRead(int laneId, uint8_t v) {
-    int laneFrom = (laneId * 6) / 8;
+        auto vLower = SHFL_SYNC((unsigned int)v, laneFrom, kWarpSize);
+        auto vUpper = SHFL_SYNC((unsigned int)v, laneFrom + 1, kWarpSize);
+        auto vConcat = (vUpper << 8) | vLower;
 
-    auto vLower = SHFL_SYNC((unsigned int) v, laneFrom, kWarpSize);
-    auto vUpper = SHFL_SYNC((unsigned int) v, laneFrom + 1, kWarpSize);
-    auto vConcat = (vUpper << 8) | vLower;
+        // Now, this is weird. Each lane reads two uint8, but we wish to use the
+        // bfe.u32 instruction to read a 6 bit value from the concatenated
+        // uint32. The offset in which we wish to read in the concatenated word
+        // is the following:
+        //
+        // 0: 0, 1: offset 0 size 6
+        // 1: 0, 1: offset 6 size 6
+        // 2: 1, 2: offset 4 size 6
+        // 3: 2, 3: offset 2 size 6
+        //
+        // The offsets are the following (concatenated together):
+        // 0x2460
+        // We can thus use bfe.u32 as a lookup table for the above sequence.
+        unsigned int pos;
+        GET_BITFIELD_U32(pos, 0x2460, (laneId & 0x3) * 4, 4);
 
-    // Now, this is weird. Each lane reads two uint8, but we wish to use the
-    // bfe.u32 instruction to read a 6 bit value from the concatenated uint32.
-    // The offset in which we wish to read in the concatenated word is the
-    // following:
-    //
-    // 0: 0, 1: offset 0 size 6
-    // 1: 0, 1: offset 6 size 6
-    // 2: 1, 2: offset 4 size 6
-    // 3: 2, 3: offset 2 size 6
-    //
-    // The offsets are the following (concatenated together):
-    // 0x2460
-    // We can thus use bfe.u32 as a lookup table for the above sequence.
-    unsigned int pos;
-    GET_BITFIELD_U32(pos, 0x2460, (laneId & 0x3) * 4, 4);
+        unsigned int out;
+        GET_BITFIELD_U32(out, vConcat, pos, 6);
 
-    unsigned int out;
-    GET_BITFIELD_U32(out, vConcat, pos, 6);
-
-    return out;
-  }
+        return out;
+    }
 };
 
-
 // Read/write 5 bit fields, packed across the warp into 20 bytes
 template <>
 struct WarpPackedBits<uint8_t, 5> {
-  static __device__ void write(int laneId, uint8_t v, bool valid, uint8_t* out) {
-    // Lower 24 lanes wwrite out packed data
-    int laneFrom = (laneId * 8) / 5;
+    static __device__ void write(
+            int laneId,
+            uint8_t v,
+            bool valid,
+            uint8_t* out) {
+        // Lower 24 lanes wwrite out packed data
+        int laneFrom = (laneId * 8) / 5;
 
-    v = valid ? v : 0;
-    v &= 0x1f; // ensure we have only 6 bits
+        v = valid ? v : 0;
+        v &= 0x1f; // ensure we have only 6 bits
 
-    uint8_t lo = (uint8_t) SHFL_SYNC((unsigned int) v, laneFrom, kWarpSize);
-    uint8_t hi = (uint8_t) SHFL_SYNC((unsigned int) v, laneFrom + 1, kWarpSize);
-    uint8_t hi2 = (uint8_t) SHFL_SYNC((unsigned int) v, laneFrom + 2, kWarpSize);
+        uint8_t lo = (uint8_t)SHFL_SYNC((unsigned int)v, laneFrom, kWarpSize);
+        uint8_t hi =
+                (uint8_t)SHFL_SYNC((unsigned int)v, laneFrom + 1, kWarpSize);
+        uint8_t hi2 =
+                (uint8_t)SHFL_SYNC((unsigned int)v, laneFrom + 2, kWarpSize);
 
-    uint8_t vOut = 0;
+        uint8_t vOut = 0;
 
-    // lsb     ...    msb
-    // 0: 0 0 0 0 0 1 1 1
-    // 1: 1 1 2 2 2 2 2 3
-    // 2: 3 3 3 3 4 4 4 4
-    // 3: 4 5 5 5 5 5 6 6
-    // 4: 6 6 6 7 7 7 7 7
-    switch (laneId % 5) {
-      case 0:
-        // 5 msbs of lower as vOut lsbs
-        // 3 lsbs of upper as vOut msbs
-        vOut = (lo & 0x1f) | (hi << 5);
-        break;
-      case 1:
-        // 2 msbs of lower as vOut lsbs
-        // 5 lsbs of upper as vOut msbs
-        // 1 lsbs of upper2 as vOut msb
-        vOut = (lo >> 3) | (hi << 2) | (hi2 << 7);
-        break;
-      case 2:
-        // 4 msbs of lower as vOut lsbs
-        // 4 lsbs of upper as vOut msbs
-        vOut = (lo >> 1) | (hi << 4);
-        break;
-      case 3:
-        // 1 msbs of lower as vOut lsbs
-        // 5 lsbs of upper as vOut msbs
-        // 2 lsbs of upper2 as vOut msb
-        vOut = (lo >> 4) | (hi << 1) | (hi2 << 6);
-        break;
-      case 4:
-        // 3 msbs of lower as vOut lsbs
-        // 5 lsbs of upper as vOut msbs
-        vOut = (lo >> 2) | (hi << 3);
-        break;
+        // lsb     ...    msb
+        // 0: 0 0 0 0 0 1 1 1
+        // 1: 1 1 2 2 2 2 2 3
+        // 2: 3 3 3 3 4 4 4 4
+        // 3: 4 5 5 5 5 5 6 6
+        // 4: 6 6 6 7 7 7 7 7
+        switch (laneId % 5) {
+            case 0:
+                // 5 msbs of lower as vOut lsbs
+                // 3 lsbs of upper as vOut msbs
+                vOut = (lo & 0x1f) | (hi << 5);
+                break;
+            case 1:
+                // 2 msbs of lower as vOut lsbs
+                // 5 lsbs of upper as vOut msbs
+                // 1 lsbs of upper2 as vOut msb
+                vOut = (lo >> 3) | (hi << 2) | (hi2 << 7);
+                break;
+            case 2:
+                // 4 msbs of lower as vOut lsbs
+                // 4 lsbs of upper as vOut msbs
+                vOut = (lo >> 1) | (hi << 4);
+                break;
+            case 3:
+                // 1 msbs of lower as vOut lsbs
+                // 5 lsbs of upper as vOut msbs
+                // 2 lsbs of upper2 as vOut msb
+                vOut = (lo >> 4) | (hi << 1) | (hi2 << 6);
+                break;
+            case 4:
+                // 3 msbs of lower as vOut lsbs
+                // 5 lsbs of upper as vOut msbs
+                vOut = (lo >> 2) | (hi << 3);
+                break;
+        }
+
+        if (laneId < 20) {
+            // There could be prior data
+            out[laneId] |= vOut;
+        }
     }
 
-    if (laneId < 20) {
-      // There could be prior data
-      out[laneId] |= vOut;
-    }
-  }
+    static inline __device__ uint8_t read(int laneId, uint8_t* in) {
+        uint8_t v = 0;
 
-  static inline __device__ uint8_t read(int laneId, uint8_t* in) {
-    uint8_t v = 0;
+        if (laneId < 20) {
+            v = in[laneId];
+        }
 
-    if (laneId < 20) {
-      v = in[laneId];
+        return v;
     }
 
-    return v;
-  }
+    static inline __device__ uint8_t postRead(int laneId, uint8_t v) {
+        int laneFrom = (laneId * 5) / 8;
 
-  static inline __device__ uint8_t postRead(int laneId, uint8_t v) {
-    int laneFrom = (laneId * 5) / 8;
+        auto vLower = SHFL_SYNC((unsigned int)v, laneFrom, kWarpSize);
+        auto vUpper = SHFL_SYNC((unsigned int)v, laneFrom + 1, kWarpSize);
+        auto vConcat = (vUpper << 8) | vLower;
 
-    auto vLower = SHFL_SYNC((unsigned int) v, laneFrom, kWarpSize);
-    auto vUpper = SHFL_SYNC((unsigned int) v, laneFrom + 1, kWarpSize);
-    auto vConcat = (vUpper << 8) | vLower;
+        // Now, this is weird. Each lane reads two uint8, but we wish to use the
+        // bfe.u32 instruction to read a 5 bit value from the concatenated
+        // uint32. The offset in which we wish to read in the concatenated word
+        // is the following:
+        //
+        // 0: 0, 1: offset 0 size 5
+        // 1: 0, 1: offset 5 size 5
+        // 2: 1, 2: offset 2 size 5
+        // 3: 1, 2: offset 7 size 5
+        // 4: 2, 3: offset 4 size 5
+        // 5: 3, 4: offset 1 size 5
+        // 6: 3, 4: offset 6 size 5
+        // 7: 4, 5: offset 3 size 5
+        //
+        // The offsets are the following (concatenated together):
+        // 0x36147250
+        // We can thus use bfe.u32 as a lookup table for the above sequence.
+        unsigned int pos;
+        GET_BITFIELD_U32(pos, 0x36147250, (laneId & 0x7) * 4, 4);
 
-    // Now, this is weird. Each lane reads two uint8, but we wish to use the
-    // bfe.u32 instruction to read a 5 bit value from the concatenated uint32.
-    // The offset in which we wish to read in the concatenated word is the
-    // following:
-    //
-    // 0: 0, 1: offset 0 size 5
-    // 1: 0, 1: offset 5 size 5
-    // 2: 1, 2: offset 2 size 5
-    // 3: 1, 2: offset 7 size 5
-    // 4: 2, 3: offset 4 size 5
-    // 5: 3, 4: offset 1 size 5
-    // 6: 3, 4: offset 6 size 5
-    // 7: 4, 5: offset 3 size 5
-    //
-    // The offsets are the following (concatenated together):
-    // 0x36147250
-    // We can thus use bfe.u32 as a lookup table for the above sequence.
-    unsigned int pos;
-    GET_BITFIELD_U32(pos, 0x36147250, (laneId & 0x7) * 4, 4);
+        unsigned int out;
+        GET_BITFIELD_U32(out, vConcat, pos, 5);
 
-    unsigned int out;
-    GET_BITFIELD_U32(out, vConcat, pos, 5);
-
-    return out;
-  }
+        return out;
+    }
 };
 
 // Read/write 4 bit fields, packed across the warp into 16 bytes
 template <>
 struct WarpPackedBits<uint8_t, 4> {
-  static __device__ void write(int laneId, uint8_t v, bool valid, uint8_t* out) {
-    // Lower 16 lanes write out packed data
-    int laneFrom = laneId * 2;
+    static __device__ void write(
+            int laneId,
+            uint8_t v,
+            bool valid,
+            uint8_t* out) {
+        // Lower 16 lanes write out packed data
+        int laneFrom = laneId * 2;
 
-    v = valid ? v : 0;
+        v = valid ? v : 0;
 
-    uint8_t vLower = (uint8_t) SHFL_SYNC((unsigned int) v, laneFrom, kWarpSize);
-    uint8_t vUpper = (uint8_t) SHFL_SYNC((unsigned int) v, laneFrom + 1, kWarpSize);
+        uint8_t vLower =
+                (uint8_t)SHFL_SYNC((unsigned int)v, laneFrom, kWarpSize);
+        uint8_t vUpper =
+                (uint8_t)SHFL_SYNC((unsigned int)v, laneFrom + 1, kWarpSize);
 
-    uint8_t vOut = (vLower & 0xf) | (vUpper << 4);
+        uint8_t vOut = (vLower & 0xf) | (vUpper << 4);
 
-    if (laneId < 16) {
-      // There could be prior data
-      out[laneId] |= vOut;
-    }
-  }
-
-  static inline __device__ uint8_t read(int laneId, uint8_t* in) {
-    uint8_t v = 0;
-
-    if (laneId < 16) {
-      v = in[laneId];
+        if (laneId < 16) {
+            // There could be prior data
+            out[laneId] |= vOut;
+        }
     }
 
-    return v;
-  }
+    static inline __device__ uint8_t read(int laneId, uint8_t* in) {
+        uint8_t v = 0;
 
-  static inline __device__ uint8_t postRead(int laneId, uint8_t v) {
-    int laneFrom = laneId / 2;
-    auto v2 = shfl((unsigned int) v, laneFrom);
-    return getBitfield(v2, (laneId & 0x1) * 4, 4);
-  }
+        if (laneId < 16) {
+            v = in[laneId];
+        }
+
+        return v;
+    }
+
+    static inline __device__ uint8_t postRead(int laneId, uint8_t v) {
+        int laneFrom = laneId / 2;
+        auto v2 = shfl((unsigned int)v, laneFrom);
+        return getBitfield(v2, (laneId & 0x1) * 4, 4);
+    }
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/WarpSelectKernel.cuh

@@ -9,62 +9,77 @@
 
 #include <faiss/gpu/utils/Select.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
-template <typename K,
-          typename IndexType,
-          bool Dir,
-          int NumWarpQ,
-          int NumThreadQ,
-          int ThreadsPerBlock>
-__global__ void warpSelect(Tensor<K, 2, true> in,
-                           Tensor<K, 2, true> outK,
-                           Tensor<IndexType, 2, true> outV,
-                           K initK,
-                           IndexType initV,
-                           int k) {
-  constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
+template <
+        typename K,
+        typename IndexType,
+        bool Dir,
+        int NumWarpQ,
+        int NumThreadQ,
+        int ThreadsPerBlock>
+__global__ void warpSelect(
+        Tensor<K, 2, true> in,
+        Tensor<K, 2, true> outK,
+        Tensor<IndexType, 2, true> outV,
+        K initK,
+        IndexType initV,
+        int k) {
+    constexpr int kNumWarps = ThreadsPerBlock / kWarpSize;
 
-  WarpSelect<K, IndexType, Dir, Comparator<K>,
-                NumWarpQ, NumThreadQ, ThreadsPerBlock>
-    heap(initK, initV, k);
+    WarpSelect<
+            K,
+            IndexType,
+            Dir,
+            Comparator<K>,
+            NumWarpQ,
+            NumThreadQ,
+            ThreadsPerBlock>
+            heap(initK, initV, k);
 
-  int warpId = threadIdx.x / kWarpSize;
-  int row = blockIdx.x * kNumWarps + warpId;
+    int warpId = threadIdx.x / kWarpSize;
+    int row = blockIdx.x * kNumWarps + warpId;
 
-  if (row >= in.getSize(0)) {
-    return;
-  }
+    if (row >= in.getSize(0)) {
+        return;
+    }
 
-  int i = getLaneId();
-  K* inStart = in[row][i].data();
+    int i = getLaneId();
+    K* inStart = in[row][i].data();
 
-  // Whole warps must participate in the selection
-  int limit = utils::roundDown(in.getSize(1), kWarpSize);
+    // Whole warps must participate in the selection
+    int limit = utils::roundDown(in.getSize(1), kWarpSize);
 
-  for (; i < limit; i += kWarpSize) {
-    heap.add(*inStart, (IndexType) i);
-    inStart += kWarpSize;
-  }
+    for (; i < limit; i += kWarpSize) {
+        heap.add(*inStart, (IndexType)i);
+        inStart += kWarpSize;
+    }
 
-  // Handle non-warp multiple remainder
-  if (i < in.getSize(1)) {
-    heap.addThreadQ(*inStart, (IndexType) i);
-  }
+    // Handle non-warp multiple remainder
+    if (i < in.getSize(1)) {
+        heap.addThreadQ(*inStart, (IndexType)i);
+    }
 
-  heap.reduce();
-  heap.writeOut(outK[row].data(),
-                outV[row].data(), k);
+    heap.reduce();
+    heap.writeOut(outK[row].data(), outV[row].data(), k);
 }
 
-void runWarpSelect(Tensor<float, 2, true>& in,
-                   Tensor<float, 2, true>& outKeys,
-                   Tensor<int, 2, true>& outIndices,
-                   bool dir, int k, cudaStream_t stream);
+void runWarpSelect(
+        Tensor<float, 2, true>& in,
+        Tensor<float, 2, true>& outKeys,
+        Tensor<int, 2, true>& outIndices,
+        bool dir,
+        int k,
+        cudaStream_t stream);
 
-void runWarpSelect(Tensor<half, 2, true>& in,
-                   Tensor<half, 2, true>& outKeys,
-                   Tensor<int, 2, true>& outIndices,
-                   bool dir, int k, cudaStream_t stream);
+void runWarpSelect(
+        Tensor<half, 2, true>& in,
+        Tensor<half, 2, true>& outKeys,
+        Tensor<int, 2, true>& outIndices,
+        bool dir,
+        int k,
+        cudaStream_t stream);
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/WarpShuffles.cuh

@@ -5,122 +5,123 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
 #include <cuda.h>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 // defines to simplify the SASS assembly structure file/line in the profiler
 #if CUDA_VERSION >= 9000
-#define SHFL_SYNC(VAL, SRC_LANE, WIDTH)         \
-  __shfl_sync(0xffffffff, VAL, SRC_LANE, WIDTH)
+#define SHFL_SYNC(VAL, SRC_LANE, WIDTH) \
+    __shfl_sync(0xffffffff, VAL, SRC_LANE, WIDTH)
 #else
-#define SHFL_SYNC(VAL, SRC_LANE, WIDTH)         \
-  __shfl(VAL, SRC_LANE, WIDTH)
+#define SHFL_SYNC(VAL, SRC_LANE, WIDTH) __shfl(VAL, SRC_LANE, WIDTH)
 #endif
 
 template <typename T>
-inline __device__ T shfl(const T val,
-                         int srcLane, int width = kWarpSize) {
+inline __device__ T shfl(const T val, int srcLane, int width = kWarpSize) {
 #if CUDA_VERSION >= 9000
-  return __shfl_sync(0xffffffff, val, srcLane, width);
+    return __shfl_sync(0xffffffff, val, srcLane, width);
 #else
-  return __shfl(val, srcLane, width);
+    return __shfl(val, srcLane, width);
 #endif
 }
 
 // CUDA SDK does not provide specializations for T*
 template <typename T>
-inline __device__ T* shfl(T* const val,
-                         int srcLane, int width = kWarpSize) {
-  static_assert(sizeof(T*) == sizeof(long long), "pointer size");
-  long long v = (long long) val;
+inline __device__ T* shfl(T* const val, int srcLane, int width = kWarpSize) {
+    static_assert(sizeof(T*) == sizeof(long long), "pointer size");
+    long long v = (long long)val;
 
-  return (T*) shfl(v, srcLane, width);
+    return (T*)shfl(v, srcLane, width);
 }
 
 template <typename T>
-inline __device__ T shfl_up(const T val,
-                            unsigned int delta, int width = kWarpSize) {
+inline __device__ T
+shfl_up(const T val, unsigned int delta, int width = kWarpSize) {
 #if CUDA_VERSION >= 9000
-  return __shfl_up_sync(0xffffffff, val, delta, width);
+    return __shfl_up_sync(0xffffffff, val, delta, width);
 #else
-  return __shfl_up(val, delta, width);
+    return __shfl_up(val, delta, width);
 #endif
 }
 
 // CUDA SDK does not provide specializations for T*
 template <typename T>
-inline __device__ T* shfl_up(T* const val,
-                             unsigned int delta, int width = kWarpSize) {
-  static_assert(sizeof(T*) == sizeof(long long), "pointer size");
-  long long v = (long long) val;
+inline __device__ T* shfl_up(
+        T* const val,
+        unsigned int delta,
+        int width = kWarpSize) {
+    static_assert(sizeof(T*) == sizeof(long long), "pointer size");
+    long long v = (long long)val;
 
-  return (T*) shfl_up(v, delta, width);
+    return (T*)shfl_up(v, delta, width);
 }
 
 template <typename T>
-inline __device__ T shfl_down(const T val,
-                              unsigned int delta, int width = kWarpSize) {
+inline __device__ T
+shfl_down(const T val, unsigned int delta, int width = kWarpSize) {
 #if CUDA_VERSION >= 9000
-  return __shfl_down_sync(0xffffffff, val, delta, width);
+    return __shfl_down_sync(0xffffffff, val, delta, width);
 #else
-  return __shfl_down(val, delta, width);
+    return __shfl_down(val, delta, width);
 #endif
 }
 
 // CUDA SDK does not provide specializations for T*
 template <typename T>
-inline __device__ T* shfl_down(T* const val,
-                              unsigned int delta, int width = kWarpSize) {
-  static_assert(sizeof(T*) == sizeof(long long), "pointer size");
-  long long v = (long long) val;
-  return (T*) shfl_down(v, delta, width);
+inline __device__ T* shfl_down(
+        T* const val,
+        unsigned int delta,
+        int width = kWarpSize) {
+    static_assert(sizeof(T*) == sizeof(long long), "pointer size");
+    long long v = (long long)val;
+    return (T*)shfl_down(v, delta, width);
 }
 
 template <typename T>
-inline __device__ T shfl_xor(const T val,
-                             int laneMask, int width = kWarpSize) {
+inline __device__ T shfl_xor(const T val, int laneMask, int width = kWarpSize) {
 #if CUDA_VERSION >= 9000
-  return __shfl_xor_sync(0xffffffff, val, laneMask, width);
+    return __shfl_xor_sync(0xffffffff, val, laneMask, width);
 #else
-  return __shfl_xor(val, laneMask, width);
+    return __shfl_xor(val, laneMask, width);
 #endif
 }
 
 // CUDA SDK does not provide specializations for T*
 template <typename T>
-inline __device__ T* shfl_xor(T* const val,
-                              int laneMask, int width = kWarpSize) {
-  static_assert(sizeof(T*) == sizeof(long long), "pointer size");
-  long long v = (long long) val;
-  return (T*) shfl_xor(v, laneMask, width);
+inline __device__ T* shfl_xor(
+        T* const val,
+        int laneMask,
+        int width = kWarpSize) {
+    static_assert(sizeof(T*) == sizeof(long long), "pointer size");
+    long long v = (long long)val;
+    return (T*)shfl_xor(v, laneMask, width);
 }
 
 // CUDA 9.0+ has half shuffle
 #if CUDA_VERSION < 9000
-inline __device__ half shfl(half v,
-                            int srcLane, int width = kWarpSize) {
-  unsigned int vu = v.x;
-  vu = __shfl(vu, srcLane, width);
+inline __device__ half shfl(half v, int srcLane, int width = kWarpSize) {
+    unsigned int vu = v.x;
+    vu = __shfl(vu, srcLane, width);
 
-  half h;
-  h.x = (unsigned short) vu;
-  return h;
+    half h;
+    h.x = (unsigned short)vu;
+    return h;
 }
 
-inline __device__ half shfl_xor(half v,
-                                int laneMask, int width = kWarpSize) {
-  unsigned int vu = v.x;
-  vu = __shfl_xor(vu, laneMask, width);
+inline __device__ half shfl_xor(half v, int laneMask, int width = kWarpSize) {
+    unsigned int vu = v.x;
+    vu = __shfl_xor(vu, laneMask, width);
 
-  half h;
-  h.x = (unsigned short) vu;
-  return h;
+    half h;
+    h.x = (unsigned short)vu;
+    return h;
 }
 #endif // CUDA_VERSION
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

faiss/gpu/utils/blockselect/BlockSelectImpl.cuh

@@ -10,85 +10,89 @@
 #include <faiss/gpu/utils/BlockSelectKernel.cuh>
 #include <faiss/gpu/utils/Limits.cuh>
 
-#define BLOCK_SELECT_DECL(TYPE, DIR, WARP_Q)                            \
-  extern void runBlockSelect_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(  \
-    Tensor<TYPE, 2, true>& in,                                          \
-    Tensor<TYPE, 2, true>& outK,                                        \
-    Tensor<int, 2, true>& outV,                                         \
-    bool dir,                                                           \
-    int k,                                                              \
-    cudaStream_t stream);                                               \
-                                                                        \
-  extern void runBlockSelectPair_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _( \
-    Tensor<TYPE, 2, true>& inK,                                         \
-    Tensor<int, 2, true>& inV,                                          \
-    Tensor<TYPE, 2, true>& outK,                                        \
-    Tensor<int, 2, true>& outV,                                         \
-    bool dir,                                                           \
-    int k,                                                              \
-    cudaStream_t stream)
+#define BLOCK_SELECT_DECL(TYPE, DIR, WARP_Q)                     \
+    extern void runBlockSelect_##TYPE##_##DIR##_##WARP_Q##_(     \
+            Tensor<TYPE, 2, true>& in,                           \
+            Tensor<TYPE, 2, true>& outK,                         \
+            Tensor<int, 2, true>& outV,                          \
+            bool dir,                                            \
+            int k,                                               \
+            cudaStream_t stream);                                \
+                                                                 \
+    extern void runBlockSelectPair_##TYPE##_##DIR##_##WARP_Q##_( \
+            Tensor<TYPE, 2, true>& inK,                          \
+            Tensor<int, 2, true>& inV,                           \
+            Tensor<TYPE, 2, true>& outK,                         \
+            Tensor<int, 2, true>& outV,                          \
+            bool dir,                                            \
+            int k,                                               \
+            cudaStream_t stream)
 
-#define BLOCK_SELECT_IMPL(TYPE, DIR, WARP_Q, THREAD_Q)                  \
-  void runBlockSelect_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(         \
-    Tensor<TYPE, 2, true>& in,                                          \
-    Tensor<TYPE, 2, true>& outK,                                        \
-    Tensor<int, 2, true>& outV,                                         \
-    bool dir,                                                           \
-    int k,                                                              \
-    cudaStream_t stream) {                                              \
-    FAISS_ASSERT(in.getSize(0) == outK.getSize(0));                     \
-    FAISS_ASSERT(in.getSize(0) == outV.getSize(0));                     \
-    FAISS_ASSERT(outK.getSize(1) == k);                                 \
-    FAISS_ASSERT(outV.getSize(1) == k);                                 \
-                                                                        \
-    auto grid = dim3(in.getSize(0));                                    \
-                                                                        \
-    constexpr int kBlockSelectNumThreads = (WARP_Q <= 1024) ? 128 : 64; \
-    auto block = dim3(kBlockSelectNumThreads);                          \
-                                                                        \
-    FAISS_ASSERT(k <= WARP_Q);                                          \
-    FAISS_ASSERT(dir == DIR);                                           \
-                                                                        \
-    auto kInit = dir ? Limits<TYPE>::getMin() : Limits<TYPE>::getMax(); \
-    auto vInit = -1;                                                    \
-                                                                        \
-    blockSelect<TYPE, int, DIR, WARP_Q, THREAD_Q, kBlockSelectNumThreads> \
-      <<<grid, block, 0, stream>>>(in, outK, outV, kInit, vInit, k);    \
-    CUDA_TEST_ERROR();                                                  \
-  }                                                                     \
-                                                                        \
-  void runBlockSelectPair_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(     \
-    Tensor<TYPE, 2, true>& inK,                                         \
-    Tensor<int, 2, true>& inV,                                          \
-    Tensor<TYPE, 2, true>& outK,                                        \
-    Tensor<int, 2, true>& outV,                                         \
-    bool dir,                                                           \
-    int k,                                                              \
-    cudaStream_t stream) {                                              \
-    FAISS_ASSERT(inK.isSameSize(inV));                                  \
-    FAISS_ASSERT(outK.isSameSize(outV));                                \
-                                                                        \
-    auto grid = dim3(inK.getSize(0));                                   \
-                                                                        \
-    constexpr int kBlockSelectNumThreads = (WARP_Q <= 1024) ? 128 : 64; \
-    auto block = dim3(kBlockSelectNumThreads);                          \
-                                                                        \
-    FAISS_ASSERT(k <= WARP_Q);                                          \
-    FAISS_ASSERT(dir == DIR);                                           \
-                                                                        \
-    auto kInit = dir ? Limits<TYPE>::getMin() : Limits<TYPE>::getMax(); \
-    auto vInit = -1;                                                    \
-                                                                        \
-    blockSelectPair<TYPE, int, DIR, WARP_Q, THREAD_Q, kBlockSelectNumThreads> \
-      <<<grid, block, 0, stream>>>(inK, inV, outK, outV, kInit, vInit, k); \
-    CUDA_TEST_ERROR();                                                  \
-  }
+#define BLOCK_SELECT_IMPL(TYPE, DIR, WARP_Q, THREAD_Q)                         \
+    void runBlockSelect_##TYPE##_##DIR##_##WARP_Q##_(                          \
+            Tensor<TYPE, 2, true>& in,                                         \
+            Tensor<TYPE, 2, true>& outK,                                       \
+            Tensor<int, 2, true>& outV,                                        \
+            bool dir,                                                          \
+            int k,                                                             \
+            cudaStream_t stream) {                                             \
+        FAISS_ASSERT(in.getSize(0) == outK.getSize(0));                        \
+        FAISS_ASSERT(in.getSize(0) == outV.getSize(0));                        \
+        FAISS_ASSERT(outK.getSize(1) == k);                                    \
+        FAISS_ASSERT(outV.getSize(1) == k);                                    \
+                                                                               \
+        auto grid = dim3(in.getSize(0));                                       \
+                                                                               \
+        constexpr int kBlockSelectNumThreads = (WARP_Q <= 1024) ? 128 : 64;    \
+        auto block = dim3(kBlockSelectNumThreads);                             \
+                                                                               \
+        FAISS_ASSERT(k <= WARP_Q);                                             \
+        FAISS_ASSERT(dir == DIR);                                              \
+                                                                               \
+        auto kInit = dir ? Limits<TYPE>::getMin() : Limits<TYPE>::getMax();    \
+        auto vInit = -1;                                                       \
+                                                                               \
+        blockSelect<TYPE, int, DIR, WARP_Q, THREAD_Q, kBlockSelectNumThreads>  \
+                <<<grid, block, 0, stream>>>(in, outK, outV, kInit, vInit, k); \
+        CUDA_TEST_ERROR();                                                     \
+    }                                                                          \
+                                                                               \
+    void runBlockSelectPair_##TYPE##_##DIR##_##WARP_Q##_(                      \
+            Tensor<TYPE, 2, true>& inK,                                        \
+            Tensor<int, 2, true>& inV,                                         \
+            Tensor<TYPE, 2, true>& outK,                                       \
+            Tensor<int, 2, true>& outV,                                        \
+            bool dir,                                                          \
+            int k,                                                             \
+            cudaStream_t stream) {                                             \
+        FAISS_ASSERT(inK.isSameSize(inV));                                     \
+        FAISS_ASSERT(outK.isSameSize(outV));                                   \
+                                                                               \
+        auto grid = dim3(inK.getSize(0));                                      \
+                                                                               \
+        constexpr int kBlockSelectNumThreads = (WARP_Q <= 1024) ? 128 : 64;    \
+        auto block = dim3(kBlockSelectNumThreads);                             \
+                                                                               \
+        FAISS_ASSERT(k <= WARP_Q);                                             \
+        FAISS_ASSERT(dir == DIR);                                              \
+                                                                               \
+        auto kInit = dir ? Limits<TYPE>::getMin() : Limits<TYPE>::getMax();    \
+        auto vInit = -1;                                                       \
+                                                                               \
+        blockSelectPair<                                                       \
+                TYPE,                                                          \
+                int,                                                           \
+                DIR,                                                           \
+                WARP_Q,                                                        \
+                THREAD_Q,                                                      \
+                kBlockSelectNumThreads><<<grid, block, 0, stream>>>(           \
+                inK, inV, outK, outV, kInit, vInit, k);                        \
+        CUDA_TEST_ERROR();                                                     \
+    }
 
+#define BLOCK_SELECT_CALL(TYPE, DIR, WARP_Q) \
+    runBlockSelect_##TYPE##_##DIR##_##WARP_Q##_(in, outK, outV, dir, k, stream)
 
-#define BLOCK_SELECT_CALL(TYPE, DIR, WARP_Q)                    \
-  runBlockSelect_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(      \
-    in, outK, outV, dir, k, stream)
-
-#define BLOCK_SELECT_PAIR_CALL(TYPE, DIR, WARP_Q)               \
-  runBlockSelectPair_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(  \
-    inK, inV, outK, outV, dir, k, stream)
+#define BLOCK_SELECT_PAIR_CALL(TYPE, DIR, WARP_Q)    \
+    runBlockSelectPair_##TYPE##_##DIR##_##WARP_Q##_( \
+            inK, inV, outK, outV, dir, k, stream)

faiss/gpu/utils/warpselect/WarpSelectImpl.cuh

@@ -5,43 +5,41 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <faiss/gpu/utils/WarpSelectKernel.cuh>
 #include <faiss/gpu/utils/Limits.cuh>
+#include <faiss/gpu/utils/WarpSelectKernel.cuh>
 
-#define WARP_SELECT_DECL(TYPE, DIR, WARP_Q)                             \
-  extern void runWarpSelect_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(   \
-    Tensor<TYPE, 2, true>& in,                                          \
-    Tensor<TYPE, 2, true>& outK,                                        \
-    Tensor<int, 2, true>& outV,                                         \
-    bool dir,                                                           \
-    int k,                                                              \
-    cudaStream_t stream)
+#define WARP_SELECT_DECL(TYPE, DIR, WARP_Q)                 \
+    extern void runWarpSelect_##TYPE##_##DIR##_##WARP_Q##_( \
+            Tensor<TYPE, 2, true>& in,                      \
+            Tensor<TYPE, 2, true>& outK,                    \
+            Tensor<int, 2, true>& outV,                     \
+            bool dir,                                       \
+            int k,                                          \
+            cudaStream_t stream)
 
-#define WARP_SELECT_IMPL(TYPE, DIR, WARP_Q, THREAD_Q)                   \
-  void runWarpSelect_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(          \
-    Tensor<TYPE, 2, true>& in,                                          \
-    Tensor<TYPE, 2, true>& outK,                                        \
-    Tensor<int, 2, true>& outV,                                         \
-    bool dir,                                                           \
-    int k,                                                              \
-    cudaStream_t stream) {                                              \
-                                                                        \
-    constexpr int kWarpSelectNumThreads = 128;                          \
-    auto grid = dim3(utils::divUp(in.getSize(0),                        \
-                                  (kWarpSelectNumThreads / kWarpSize))); \
-    auto block = dim3(kWarpSelectNumThreads);                           \
-                                                                        \
-    FAISS_ASSERT(k <= WARP_Q);                                          \
-    FAISS_ASSERT(dir == DIR);                                           \
-                                                                        \
-    auto kInit = dir ? Limits<TYPE>::getMin() : Limits<TYPE>::getMax(); \
-    auto vInit = -1;                                                    \
-                                                                        \
-    warpSelect<TYPE, int, DIR, WARP_Q, THREAD_Q, kWarpSelectNumThreads> \
-      <<<grid, block, 0, stream>>>(in, outK, outV, kInit, vInit, k);    \
-    CUDA_TEST_ERROR();                                                  \
-  }
+#define WARP_SELECT_IMPL(TYPE, DIR, WARP_Q, THREAD_Q)                          \
+    void runWarpSelect_##TYPE##_##DIR##_##WARP_Q##_(                           \
+            Tensor<TYPE, 2, true>& in,                                         \
+            Tensor<TYPE, 2, true>& outK,                                       \
+            Tensor<int, 2, true>& outV,                                        \
+            bool dir,                                                          \
+            int k,                                                             \
+            cudaStream_t stream) {                                             \
+        constexpr int kWarpSelectNumThreads = 128;                             \
+        auto grid = dim3(utils::divUp(                                         \
+                in.getSize(0), (kWarpSelectNumThreads / kWarpSize)));          \
+        auto block = dim3(kWarpSelectNumThreads);                              \
+                                                                               \
+        FAISS_ASSERT(k <= WARP_Q);                                             \
+        FAISS_ASSERT(dir == DIR);                                              \
+                                                                               \
+        auto kInit = dir ? Limits<TYPE>::getMin() : Limits<TYPE>::getMax();    \
+        auto vInit = -1;                                                       \
+                                                                               \
+        warpSelect<TYPE, int, DIR, WARP_Q, THREAD_Q, kWarpSelectNumThreads>    \
+                <<<grid, block, 0, stream>>>(in, outK, outV, kInit, vInit, k); \
+        CUDA_TEST_ERROR();                                                     \
+    }
 
-#define WARP_SELECT_CALL(TYPE, DIR, WARP_Q)                     \
-  runWarpSelect_ ## TYPE ## _ ## DIR ## _ ## WARP_Q ## _(       \
-    in, outK, outV, dir, k, stream)
+#define WARP_SELECT_CALL(TYPE, DIR, WARP_Q) \
+    runWarpSelect_##TYPE##_##DIR##_##WARP_Q##_(in, outK, outV, dir, k, stream)