From a60e47fd0230e6bf2e09ab2cbb0b4f1ad741dc17 Mon Sep 17 00:00:00 2001
From: Maria Lomeli <marialomeli@meta.com>
Date: Tue, 15 Oct 2024 11:11:49 -0700
Subject: [PATCH] Cache device major version value to avoid multiple calls of
 getCudaDeviceProperties

Summary: This diff enables to cache the device major version value so getCudaDeviceProperties() doesn't need to be called multiple times.

Differential Revision: D64047778
---
 faiss/gpu/GpuDistance.cu | 14 ++++++++++----
 faiss/gpu/GpuIndex.cu    | 12 +++++++++---
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/faiss/gpu/GpuDistance.cu b/faiss/gpu/GpuDistance.cu
index 38a62f03bb..599f4a3072 100644
--- a/faiss/gpu/GpuDistance.cu
+++ b/faiss/gpu/GpuDistance.cu
@@ -51,12 +51,18 @@ using namespace raft::distance;
 using namespace raft::neighbors;
 #endif
 
+/// Caches device major version
+int device_major_version = -1;
+
 bool should_use_raft(GpuDistanceParams args) {
-    cudaDeviceProp prop;
-    int dev = args.device >= 0 ? args.device : getCurrentDevice();
-    cudaGetDeviceProperties(&prop, dev);
+    if (device_major_version < 0) {
+        cudaDeviceProp prop;
+        int dev = args.device >= 0 ? args.device : getCurrentDevice();
+        cudaGetDeviceProperties(&prop, dev);
+        device_major_version = prop.major;
+    }
 
-    if (prop.major < 7)
+    if (device_major_version < 7)
         return false;
 
     return args.use_raft;
diff --git a/faiss/gpu/GpuIndex.cu b/faiss/gpu/GpuIndex.cu
index d1ae3b5384..f91b7dc9c5 100644
--- a/faiss/gpu/GpuIndex.cu
+++ b/faiss/gpu/GpuIndex.cu
@@ -42,11 +42,17 @@ constexpr idx_t kAddVecSize = (idx_t)512 * 1024;
 // FIXME: parameterize based on algorithm need
 constexpr idx_t kSearchVecSize = (idx_t)32 * 1024;
 
+/// Caches device major version
+extern int device_major_version;
+
 bool should_use_raft(GpuIndexConfig config_) {
-    cudaDeviceProp prop;
-    cudaGetDeviceProperties(&prop, config_.device);
+    if (device_major_version < 0) {
+        cudaDeviceProp prop;
+        cudaGetDeviceProperties(&prop, config_.device);
+        device_major_version = prop.major;
+    }
 
-    if (prop.major < 7)
+    if (device_major_version < 7)
         return false;
 
     return config_.use_raft;