Use BlockSize in spmv

Author: schnellerhase

cpp/dolfinx/la/MatrixCSR.h

@@ -16,6 +16,7 @@
 #include <mpi.h>
 #include <numeric>
 #include <span>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -788,13 +789,13 @@ void MatrixCSR<Scalar, V, W, X>::mult(la::Vector<Scalar>& x,
   // yi[0] += Ai[0] * xi[0]
   if (_bs[1] == 1)
   {
-    impl::spmv<Scalar, 1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
-                          _bs[0], 1);
+    impl::spmv<Scalar, BS<1>>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x,
+                              _y, _bs[0], BS<1>());
   }
   else
   {
-    impl::spmv<Scalar, -1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
-                           _bs[0], _bs[1]);
+    impl::spmv<Scalar, int>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x,
+                            _y, _bs[0], _bs[1]);
   }
 
   // finalize ghost update
@@ -804,13 +805,13 @@ void MatrixCSR<Scalar, V, W, X>::mult(la::Vector<Scalar>& x,
   // yi[0] += Ai[1] * xi[1]
   if (_bs[1] == 1)
   {
-    impl::spmv<Scalar, 1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,
-                          _bs[0], 1);
+    impl::spmv<Scalar, BS<1>>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x,
+                              _y, _bs[0], BS<1>());
   }
   else
   {
-    impl::spmv<Scalar, -1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,
-                           _bs[0], _bs[1]);
+    impl::spmv<Scalar, int>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,
+                            _bs[0], _bs[1]);
   }
 }
 

cpp/dolfinx/la/matrix_csr_impl.h

@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include "dolfinx/common/types.h"
 #include <numeric>
 #include <span>
 #include <utility>
@@ -222,12 +223,13 @@ void insert_nonblocked_csr(U&& data, const V& cols, const W& row_ptr,
 /// @param y
 /// @param bs0
 /// @param bs1
-template <typename T, int BS1>
+template <typename T, BlockSize BS1>
 void spmv(std::span<const T> values, std::span<const std::int64_t> row_begin,
           std::span<const std::int64_t> row_end,
           std::span<const std::int32_t> indices, std::span<const T> x,
-          std::span<T> y, int bs0, int bs1)
+          std::span<T> y, int bs0, BS1 _bs1)
 {
+  int bs1 = block_size(_bs1);
   assert(row_begin.size() == row_end.size());
   for (int k0 = 0; k0 < bs0; ++k0)
   {
@@ -236,21 +238,10 @@ void spmv(std::span<const T> values, std::span<const std::int64_t> row_begin,
       T vi{0};
       for (std::int32_t j = row_begin[i]; j < row_end[i]; j++)
       {
-        if constexpr (BS1 == -1)
+        for (int k1 = 0; k1 < bs1; ++k1)
         {
-          for (int k1 = 0; k1 < bs1; ++k1)
-          {
-            vi += values[j * bs1 * bs0 + k1 * bs0 + k0]
-                  * x[indices[j] * bs1 + k1];
-          }
-        }
-        else
-        {
-          for (int k1 = 0; k1 < BS1; ++k1)
-          {
-            vi += values[j * BS1 * bs0 + k1 * bs0 + k0]
-                  * x[indices[j] * BS1 + k1];
-          }
+          vi += values[j * bs1 * bs0 + k1 * bs0 + k0]
+                * x[indices[j] * bs1 + k1];
         }
       }