From 8c2de998f92f8add315e67bfaa1d331d8e21f011 Mon Sep 17 00:00:00 2001
From: Aditya Goel <agoel4512@gmail.com>
Date: Fri, 23 Feb 2024 01:11:54 +0000
Subject: [PATCH] transpose matmul categorical bit reproducibility

---
 src/tabmat/ext/cat_split_helpers-tmpl.cpp | 26 ++++++++++++++---------
 tests/test_reproducibility.py             | 24 +++++++++++++++++++++
 2 files changed, 40 insertions(+), 10 deletions(-)
 create mode 100644 tests/test_reproducibility.py

diff --git a/src/tabmat/ext/cat_split_helpers-tmpl.cpp b/src/tabmat/ext/cat_split_helpers-tmpl.cpp
index c40f851c..0ef80be6 100644
--- 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,30 @@ 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) {
+		#pragma omp atomic
+	        res[i] += all_res[tid * res_size + i];	    
+	    }
+	}	
     }
 }
 </%def>
diff --git a/tests/test_reproducibility.py b/tests/test_reproducibility.py
new file mode 100644
index 00000000..4dc06ee0
--- /dev/null
+++ b/tests/test_reproducibility.py
@@ -0,0 +1,24 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+import tabmat as tm
+
+N = 100
+K = 5
+
+
+@pytest.fixture
+def df():
+    rng = np.random.default_rng(1234)
+    return pd.DataFrame(
+        pd.Categorical(rng.integers(low=0, high=K - 1, size=N), categories=range(K))
+    )
+
+
+@pytest.mark.parametrize("cat_threshold", [K, K + 1])
+def test_mat_transpose_vec(df, cat_threshold):
+    rng = np.random.default_rng(1234)
+    vec = rng.normal(size=N)
+    mat = tm.from_pandas(df, cat_threshold=cat_threshold)
+    np.testing.assert_equal(mat.transpose_matvec(vec), mat.transpose_matvec(vec))