Sprinkled some OpenMP SIMD directives for obviously vectorized loops.

tatami-inc · Jul 26, 2024 · d924315 · d924315
1 parent 07089e8
commit d924315
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/include/tatami_mult/dense_column.hpp b/include/tatami_mult/dense_column.hpp
@@ -25,6 +25,9 @@ void dense_column_vector(const tatami::Matrix<Value_, Index_>& matrix, const Rig
         for (Index_ c = 0; c < NC; ++c) {
             auto ptr = ext->fetch(buffer.data());
             Output_ mult = rhs[c];
+#ifdef _OPENMP
+            #pragma omp simd
+#endif
             for (Index_ r = 0; r < length; ++r) {
                 optr[r] += mult * ptr[r];
             }
@@ -51,6 +54,9 @@ void dense_column_vectors(const tatami::Matrix<Value_, Index_>& matrix, const st
             for (size_t j = 0; j < num_rhs; ++j) {
                 auto optr = stores[j].data();
                 Output_ mult = rhs[j][c];
+#ifdef _OPENMP
+                #pragma omp simd
+#endif
                 for (Index_ r = 0; r < length; ++r) {
                     optr[r] += mult * ptr[r];
                 }
@@ -86,6 +92,9 @@ void dense_column_tatami_dense(const tatami::Matrix<Value_, Index_>& matrix, con
             for (RightIndex_ j = 0; j < rhs_col; ++j) {
                 auto optr = stores[j].data();
                 Output_ mult = rptr[j];
+#ifdef _OPENMP
+                #pragma omp simd
+#endif
                 for (Index_ r = 0; r < length; ++r) {
                     optr[r] += mult * ptr[r];
                 }
@@ -136,6 +145,9 @@ void dense_column_tatami_sparse(const tatami::Matrix<Value_, Index_>& matrix, co
                         auto optr = stores[j].data();
                         if (k < range.number && j == range.index[k]) {
                             Output_ mult = range.value[k];
+#ifdef _OPENMP
+                            #pragma omp simd
+#endif
                             for (Index_ r = 0; r < length; ++r) {
                                 optr[r] += mult * ptr[r];
                             }
@@ -153,6 +165,9 @@ void dense_column_tatami_sparse(const tatami::Matrix<Value_, Index_>& matrix, co
             for (RightIndex_ k = 0; k < range.number; ++k) {
                 auto optr = stores[range.index[k]].data();
                 Output_ mult = range.value[k];
+#ifdef _OPENMP
+                #pragma omp simd
+#endif
                 for (Index_ r = 0; r < length; ++r) {
                     optr[r] += mult * ptr[r];
                 }

diff --git a/include/tatami_mult/sparse_column.hpp b/include/tatami_mult/sparse_column.hpp
@@ -14,20 +14,29 @@ namespace internal {
 
 template<typename Value_, typename Index_, typename Output_>
 void sparse_multiply_add(const tatami::SparseRange<Value_, Index_>& range, Index_ start, Output_ mult, Output_* optr) {
+#ifdef _OPENMP
+    #pragma omp simd
+#endif
     for (Index_ r = 0; r < range.number; ++r) {
         optr[range.index[r] - start] += mult * range.value[r];
     }
 }
 
 template<typename Value_, typename Index_>
 void expand_sparse_range(const tatami::SparseRange<Value_, Index_>& range, Index_ start, std::vector<Value_>& expanded) {
+#ifdef _OPENMP
+    #pragma omp simd
+#endif
     for (Index_ k = 0; k < range.number; ++k) {
         expanded[range.index[k] - start] = range.value[k];
     }
 }
 
 template<typename Value_, typename Index_>
 void reset_expanded_sparse_range(const tatami::SparseRange<Value_, Index_>& range, Index_ start, std::vector<Value_>& expanded) {
+#ifdef _OPENMP
+    #pragma omp simd
+#endif
     for (Index_ k = 0; k < range.number; ++k) {
         expanded[range.index[k] - start] = 0;
     }
@@ -59,6 +68,9 @@ void sparse_column_vector(const tatami::Matrix<Value_, Index_>& matrix, const Ri
                     expanded.resize(length);
                     expand_sparse_range(range, start, expanded);
 
+#ifdef _OPENMP
+                    #pragma omp simd
+#endif
                     for (Index_ r = 0; r < length; ++r) {
                         optr[r] += expanded[r] * mult;
                     }
@@ -107,6 +119,9 @@ void sparse_column_vectors(const tatami::Matrix<Value_, Index_>& matrix, const s
                             has_expanded = true;
                         }
 
+#ifdef _OPENMP
+                        #pragma omp simd
+#endif
                         for (Index_ r = 0; r < length; ++r) {
                             optr[r] += expanded[r] * mult;
                         }
@@ -168,6 +183,9 @@ void sparse_column_tatami_dense(const tatami::Matrix<Value_, Index_>& matrix, co
                             has_expanded = true;
                         }
 
+#ifdef _OPENMP
+                        #pragma omp simd
+#endif
                         for (Index_ r = 0; r < length; ++r) {
                             optr[r] += expanded[r] * mult;
                         }
@@ -239,6 +257,9 @@ void sparse_column_tatami_sparse(const tatami::Matrix<Value_, Index_>& matrix, c
                         auto optr = stores[j].data();
                         if (rhs_k < rhs_range.number && j == rhs_range.index[rhs_k]) {
                             Output_ mult = rhs_range.value[rhs_k];
+#ifdef _OPENMP
+                            #pragma omp simd
+#endif
                             for (Index_ k = 0; k < range.number; ++k) {
                                 optr[range.index[k] - start] += mult * range.value[k];
                             }
@@ -256,6 +277,9 @@ void sparse_column_tatami_sparse(const tatami::Matrix<Value_, Index_>& matrix, c
             for (RightIndex_ rhs_k = 0; rhs_k < rhs_range.number; ++rhs_k) {
                 auto optr = stores[rhs_range.index[rhs_k]].data();
                 Output_ mult = rhs_range.value[rhs_k];
+#ifdef _OPENMP
+                #pragma omp simd
+#endif
                 for (Index_ k = 0; k < range.number; ++k) {
                     optr[range.index[k] - start] += mult * range.value[k];
                 }