c++sparse-matrixconstant-expression

It is possible to compute both the size and values for a std array in a constexpr function?


I am trying to express the structure of a sparse matrix (and the structure resulting from sparse matrix operations) at compile time with template parameters and constexpr functions.

I've defined this struct to represent the structure of CSC matrix.

template <size_t n_rows, 
          size_t n_cols, 
          size_t nonzeros, 
          std::array<size_t, n_cols + 1> column_ptrs,
          std::array<size_t, nonzeros> row_indices>
struct CscStructure {
  // ...
};

Then I'd like to be able to have constexpr operations like matrix-vector multiplication produce CSC structure outputs that represent the sparsity structure of the result.

For that operation one needs to compute the number of nonzero elements and std::array of row indices as constexpr values as they are template params.

This is possible by first computing the nnz elements and then the row indices e.g.:

constexpr size_t nnz = get_mat_vec_nnz<lhs, rhs>();
constexpr std::array<size_t, nnz> row_indices = get_mat_vec_row_indices<lhs, rhs, nnz>();
...
CscStructure<lhs.rows(), rhs.cols(), nnz, column_ptrs, row_indices> mat_vec_structure;

However there is significant duplication between get_mat_vec_nnz and get_mat_vec_row_indices because the first function does all the same work as the second to compute the nnz, just without recording the indices.

Is there a way to simplify this code while preserving the constexprs? For example, it seems one cannot use a std::vector because nnz = vector_row_indices.size() would not satisfy the requirements for a template param.


Solution

  • You can compute both the array content and its actual size in a single constexpr function and return a std::pair holding these two information. For instance, suppose we have a compute function that keeps only items greater than 0

    template<typename T, std::size_t N>
    constexpr auto compute (std::array<T,N> const& arr) {
        return  [&] {
            std::array<T,N> res = {};
            std::size_t i=0;
            for (auto item : arr)  {
                if (item>0)  {  res[i++] = item; }
            }
            return std::make_pair(res,i);
        } ();
    }
    

    Then you can write a generic truncate function that will truncate the array to the actual size (needs c++20 for non-type template parameter)

    template<auto info>
    constexpr auto truncate () {
        return  [&] <std::size_t...Is> (std::index_sequence<Is...> ) {
            return std::array<typename decltype(info.first)::value_type,info.second> { info.first[Is]... };
        } (std::make_index_sequence<info.second>());
    }
    

    So you don't have to write "twice" your function but instead call both truncate and compute

    constexpr std::array input  = { 1,2,-4,5,0,9};
    constexpr std::array output = truncate<compute(input)> ();
    

    Demo