Added thread control (#3)

Co-authored-by: Jozef Cernansky <[email protected]>

Author: jozefCer

task1.py

@@ -1,5 +1,11 @@
 from __future__ import annotations
 
+import os
+os.environ['MKL_NUM_THREADS'] = '27'
+os.environ['OMP_NUM_THREADS'] = '27'
+os.environ['OMP_DYNAMIC'] = 'FALSE'
+os.environ['MKL_DYNAMIC'] = 'FALSE'
+
 import argparse
 import gc
 import time
@@ -20,6 +26,8 @@
 
 import utils
 
+torch.set_num_threads(27)
+faiss.omp_set_num_threads(27)
 SEED = 42
 torch.manual_seed(SEED)
 
@@ -121,7 +129,10 @@ def search(self, queries: Tensor, k: int, nprobe: int = 100) -> tuple[np.ndarray
         D = np.empty((n_queries, k), dtype=np.float16)
         I = np.empty((n_queries, k), dtype=np.int32)
 
-        with ThreadPoolExecutor() as executor:
+        torch.set_num_threads(3)
+        faiss.omp_set_num_threads(3)
+
+        with ThreadPoolExecutor(max_workers=9) as executor: 
             results = executor.map(
                 lambda i: self._visit_buckets(k, predicted_bucket_ids[i], queries[i : i + 1], i, nprobe),
                 range(n_queries),

task2.py

@@ -1,5 +1,11 @@
 from __future__ import annotations
 
+import os
+os.environ['MKL_NUM_THREADS'] = '4'
+os.environ['OMP_NUM_THREADS'] = '4'
+os.environ['OMP_DYNAMIC'] = 'FALSE'
+os.environ['MKL_DYNAMIC'] = 'FALSE'
+
 import argparse
 import gc
 import time
@@ -21,6 +27,8 @@
 
 import utils
 
+torch.set_num_threads(4)
+faiss.omp_set_num_threads(4)
 SEED = 42
 torch.manual_seed(SEED)
 
@@ -143,7 +151,10 @@ def search(
         D = np.empty((n_queries, k), dtype=np.float16)
         I = np.empty((n_queries, k), dtype=np.int32)
 
-        with ThreadPoolExecutor() as executor:
+        torch.set_num_threads(2)
+        faiss.omp_set_num_threads(2)
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
             results = executor.map(
                 lambda i: self._visit_buckets(
                     k,
@@ -202,7 +213,7 @@ def _chunk_sort(self, dataset: Path, classes: Tensor, chunk_i: int, offsets: Off
 
         del chunk
 
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(max_workers=12) as executor:
             executor.map(
                 lambda x: self._sort_data(
                     decomposed_chunk,
@@ -244,7 +255,7 @@ def _create_buckets(self, dataset: Path, n_data: int, chunk_size: int) -> float:
         offsets = self._create_offsets(classes, n_chunks, chunk_size)
 
         logger.debug('Started bucket_init')
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(max_workers=12) as executor:
             executor.map(lambda x: self._bucket_init(classes, x), range(self.n_buckets))
 
         logger.debug('First part done')

task3.py

@@ -1,5 +1,11 @@
 from __future__ import annotations
 
+import os
+os.environ['MKL_NUM_THREADS'] = '27'
+os.environ['OMP_NUM_THREADS'] = '27'
+os.environ['OMP_DYNAMIC'] = 'FALSE'
+os.environ['MKL_DYNAMIC'] = 'FALSE'
+
 import argparse
 import gc
 import time
@@ -21,6 +27,8 @@
 
 import utils
 
+torch.set_num_threads(27)
+faiss.omp_set_num_threads(27)
 SEED = 42
 torch.manual_seed(SEED)
 
@@ -130,7 +138,10 @@ def search(
         D = np.empty((n_queries, k), dtype=np.float16)
         I = np.empty((n_queries, k), dtype=np.int32)
 
-        with ThreadPoolExecutor() as executor:
+        torch.set_num_threads(3)
+        faiss.omp_set_num_threads(3)
+
+        with ThreadPoolExecutor(max_workers=9) as executor: # max_workers=9
             results = executor.map(
                 lambda i: self._visit_buckets(k, predicted_bucket_ids[i], decomposed_queries[i : i + 1], i, nprobe),
                 range(n_queries),
@@ -181,7 +192,7 @@ def _chunk_sort(self, dataset: Path, classes: Tensor, chunk_i: int, offsets: Off
 
         del chunk
 
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(max_workers=32) as executor:
             executor.map(
                 lambda x: self._sort_data(
                     decomposed_chunk,  # noqa: F821
@@ -222,7 +233,7 @@ def _create_buckets(self, dataset: Path, n_data: int, chunk_size: int) -> float:
 
         offsets = self._create_offsets(classes, n_chunks, chunk_size)
 
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(max_workers=32) as executor:
             executor.map(lambda x: self._bucket_init(classes, x), range(self.n_buckets))
 
         logger.debug('First part done')