CategoricalMatrix A.Tb reproducibility. (#348)

* transpose matmul categorical bit reproducibility * Remove reproducibility test * Drop redundant atomic * changelog --------- Co-authored-by: Marc-Antoine Schmidt <[email protected]>
Quantco · Feb 28, 2024 · 4ac3dc3 · 4ac3dc3
1 parent 0f2e648
commit 4ac3dc3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -14,10 +14,10 @@ Unreleased
 
 - Added cython compiler directive legacy_implicit_noexcept = True to fix performance regression with cython 3.
 
-
 **Other changes:**
 
 - Refactored the pre-commit hooks to use ruff.
+- Refactored CategoricalMatrix's transpose_matvec to be deterministic when using OpenMP.
 
 3.1.13 - 2023-10-17
 -------------------

diff --git a/src/tabmat/ext/cat_split_helpers-tmpl.cpp b/src/tabmat/ext/cat_split_helpers-tmpl.cpp
@@ -1,5 +1,5 @@
 #include <vector>
-
+#include <omp.h>
 
 <%def name="transpose_matvec(dropfirst)">
 template <typename Int, typename F>
@@ -10,24 +10,29 @@ void _transpose_matvec_${dropfirst}(
     F* res,
     Int res_size
 ) {
-    #pragma omp parallel
+    int num_threads = omp_get_max_threads();
+    std::vector<F> all_res(num_threads * res_size, 0.0);
+    #pragma omp parallel shared(all_res) 
     {
-        std::vector<F> restemp(res_size, 0.0);
-        #pragma omp for
+	int tid = omp_get_thread_num();
+	F* res_slice = &all_res[tid * res_size]; 
+	#pragma omp for
         for (Py_ssize_t i = 0; i < n_rows; i++) {
             % if dropfirst == 'all_rows_drop_first':
                 Py_ssize_t col_idx = indices[i] - 1;
                 if (col_idx != -1) {
-                    restemp[col_idx] += other[i];
+                    res_slice[col_idx] += other[i];
                 }
             % else:
-                restemp[indices[i]] += other[i];
+  	        res_slice[indices[i]] += other[i];
             % endif
         }
-        for (Py_ssize_t i = 0; i < res_size; i++) {
-            # pragma omp atomic
-            res[i] += restemp[i];
-        }
+	#pragma omp for
+	for (Py_ssize_t i = 0; i < res_size; ++i) {
+	    for (int tid = 0; tid < num_threads; ++tid) {
+	        res[i] += all_res[tid * res_size + i];	    
+	    }
+	}	
     }
 }
 </%def>