Make simd a default behavior

intellistream · Sep 20, 2024 · faae468 · faae468
1 parent 375c499
commit faae468
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/src/CANDY/DPGIndex.cpp b/src/CANDY/DPGIndex.cpp
@@ -165,7 +165,7 @@ double DPGIndex::calcDist(const torch::Tensor &ta, const torch::Tensor &tb) {
     //for (size_t i = 0; i < vecDim; ++i) ans -= taPtr[i] * tbPtr[i];
     // Calculate the inner (dot) product
     auto inner_product = torch::dot(ta, tb);
-    ans = inner_product.item<float>();
+    ans = -inner_product.item<float>();
     // Convert the result back to float and return
   }
   return ans;

diff --git a/src/CANDY/NNDescentIndex.cpp b/src/CANDY/NNDescentIndex.cpp
@@ -128,16 +128,19 @@ bool NNDescentIndex::updateNN(size_t i, size_t j, double dist) {
 
 double NNDescentIndex::calcDist(const torch::Tensor &ta,
                                 const torch::Tensor &tb) {
-  auto taPtr = ta.contiguous().data_ptr<float>(),
-       tbPtr = tb.contiguous().data_ptr<float>();
   double ans = 0;
   if (faissMetric == faiss::METRIC_L2) {
-    for (size_t i = 0; i < vecDim; ++i) {
-      auto diff = taPtr[i] - tbPtr[i];
-      ans += diff * diff;
-    }
+    // Calculate the squared L2 distance as ||v1 - v2||^2
+    auto diff = ta-tb;
+    auto l2_sqr = torch::sum(diff * diff);
+    // Convert the result back to float and return
+    ans = l2_sqr.item<float>();
   } else {
-    for (size_t i = 0; i < vecDim; ++i) ans -= taPtr[i] * tbPtr[i];
+    //for (size_t i = 0; i < vecDim; ++i) ans -= taPtr[i] * tbPtr[i];
+    // Calculate the inner (dot) product
+    auto inner_product = torch::dot(ta, tb);
+    ans = -inner_product.item<float>();
+    // Convert the result back to float and return
   }
   return ans;
 }