sparse_matrix_simd.h

Go to the documentation of this file.

//
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// Copyright (C) 2020 - 2024 by the ryujin authors
//
 
#pragma once
 
#include <deal.II/base/aligned_vector.h>
#include <deal.II/base/partitioner.h>
#include <deal.II/lac/dynamic_sparsity_pattern.h>
 
#include "openmp.h"
#include "simd.h"
 
namespace ryujin
{
  namespace
  {
    template <typename Number, int n_components>
    struct EntryType {
      using type = dealii::Tensor<1, n_components, Number>;
    };
 
    template <typename Number>
    struct EntryType<Number, 1> {
      using type = Number;
    };
  } // namespace
 
 
  template <typename Number,
            int n_components = 1,
            int simd_length = dealii::VectorizedArray<Number>::size()>
  class SparseMatrixSIMD;
 
  template <int simd_length>
  class SparsityPatternSIMD
  {
  public:
    SparsityPatternSIMD();
 
    SparsityPatternSIMD(
        const unsigned int n_internal_dofs,
        const dealii::DynamicSparsityPattern &sparsity,
        const std::shared_ptr<const dealii::Utilities::MPI::Partitioner>
            &partitioner);
 
 
    void reinit(const unsigned int n_internal_dofs,
                const dealii::DynamicSparsityPattern &sparsity,
                const std::shared_ptr<const dealii::Utilities::MPI::Partitioner>
                    &partitioner);
 
    unsigned int stride_of_row(const unsigned int row) const;
 
    const unsigned int *columns(const unsigned int row) const;
 
    unsigned int row_length(const unsigned int row) const;
 
    unsigned int n_rows() const;
 
    std::size_t n_nonzero_elements() const;
 
  private:
    unsigned int n_internal_dofs;
    unsigned int n_locally_owned_dofs;
    std::shared_ptr<const dealii::Utilities::MPI::Partitioner> partitioner;
 
    dealii::AlignedVector<std::size_t> row_starts;
    dealii::AlignedVector<unsigned int> column_indices;
    dealii::AlignedVector<unsigned int> indices_transposed;
 
    std::vector<std::pair<unsigned int, unsigned int>> entries_to_be_sent;
 
    std::vector<std::pair<unsigned int, unsigned int>> send_targets;
 
    std::vector<std::pair<unsigned int, unsigned int>> receive_targets;
 
    MPI_Comm mpi_communicator;
 
    template <typename, int, int>
    friend class SparseMatrixSIMD;
  };
 
 
  template <typename Number, int n_components, int simd_length>
  class SparseMatrixSIMD
  {
  public:
    SparseMatrixSIMD();
 
    SparseMatrixSIMD(const SparsityPatternSIMD<simd_length> &sparsity);
 
    void reinit(const SparsityPatternSIMD<simd_length> &sparsity);
 
    template <typename SparseMatrix>
    void read_in(const std::array<SparseMatrix, n_components> &sparse_matrix,
                 bool locally_indexed = true);
 
    template <typename SparseMatrix>
    void read_in(const SparseMatrix &sparse_matrix,
                 bool locally_indexed = true);
 
    using VectorizedArray = dealii::VectorizedArray<Number, simd_length>;
 
    template <typename Number2>
    using EntryType = typename ryujin::EntryType<Number2, n_components>::type;
 
    /* Get scalar or tensor-valued entry: */
 
    template <typename Number2 = Number>
    EntryType<Number2>
    get_entry(const unsigned int row,
              const unsigned int position_within_column) const;
 
    template <typename Number2 = Number>
    dealii::Tensor<1, n_components, Number2>
    get_tensor(const unsigned int row,
               const unsigned int position_within_column) const;
 
    /* Get transposed scalar or tensor-valued entry: */
 
    template <typename Number2 = Number>
    EntryType<Number2>
    get_transposed_entry(const unsigned int row,
                         const unsigned int position_within_column) const;
 
    template <typename Number2 = Number>
    dealii::Tensor<1, n_components, Number2>
    get_transposed_tensor(const unsigned int row,
                          const unsigned int position_within_column) const;
 
    /* Write scalar or tensor entry: */
 
    template <typename Number2 = Number>
    void write_entry(const Number2 entry,
                     const unsigned int row,
                     const unsigned int position_within_column,
                     const bool do_streaming_store = false);
 
    template <typename Number2 = Number>
    void write_entry(const dealii::Tensor<1, n_components, Number2> &entry,
                     const unsigned int row,
                     const unsigned int position_within_column,
                     const bool do_streaming_store = false);
 
    /* Synchronize over MPI ranks: */
 
    void update_ghost_rows_start(const unsigned int communication_channel = 0);
 
    void update_ghost_rows_finish();
 
    void update_ghost_rows();
 
  private:
    const SparsityPatternSIMD<simd_length> *sparsity;
    dealii::AlignedVector<Number> data;
    dealii::AlignedVector<Number> exchange_buffer;
    std::vector<MPI_Request> requests;
  };
 
  /*
   * Inline function  definitions:
   */
 
 
  template <int simd_length>
  DEAL_II_ALWAYS_INLINE inline unsigned int
  SparsityPatternSIMD<simd_length>::stride_of_row(const unsigned int row) const
  {
    AssertIndexRange(row, row_starts.size());
 
    if (row < n_internal_dofs)
      return simd_length;
    else
      return 1;
  }
 
 
  template <int simd_length>
  DEAL_II_ALWAYS_INLINE inline const unsigned int *
  SparsityPatternSIMD<simd_length>::columns(const unsigned int row) const
  {
    AssertIndexRange(row, row_starts.size() - 1);
 
    if (row < n_internal_dofs)
      return column_indices.data() + row_starts[row / simd_length] +
             row % simd_length;
    else
      return column_indices.data() + row_starts[row];
  }
 
 
  template <int simd_length>
  DEAL_II_ALWAYS_INLINE inline unsigned int
  SparsityPatternSIMD<simd_length>::row_length(const unsigned int row) const
  {
    AssertIndexRange(row, row_starts.size() - 1);
 
    if (row < n_internal_dofs) {
      const unsigned int simd_row = row / simd_length;
      return (row_starts[simd_row + 1] - row_starts[simd_row]) / simd_length;
    } else {
      return row_starts[row + 1] - row_starts[row];
    }
  }
 
 
  template <int simd_length>
  DEAL_II_ALWAYS_INLINE inline unsigned int
  SparsityPatternSIMD<simd_length>::n_rows() const
  {
    return row_starts.size() - 1;
  }
 
 
  template <int simd_length>
  DEAL_II_ALWAYS_INLINE inline std::size_t
  SparsityPatternSIMD<simd_length>::n_nonzero_elements() const
  {
    Assert(row_starts.size() > 0, dealii::ExcNotInitialized());
 
    return row_starts.back();
  }
 
 
  template <typename Number, int n_components, int simd_length>
  template <typename Number2>
  DEAL_II_ALWAYS_INLINE inline auto
  SparseMatrixSIMD<Number, n_components, simd_length>::get_entry(
      const unsigned int row, const unsigned int position_within_column) const
      -> EntryType<Number2>
  {
    const auto result = get_tensor<Number2>(row, position_within_column);
    if constexpr (n_components == 1)
      return result[0];
    else
      return result;
  }
 
 
  template <typename Number, int n_components, int simd_length>
  template <typename Number2>
  DEAL_II_ALWAYS_INLINE inline dealii::Tensor<1, n_components, Number2>
  SparseMatrixSIMD<Number, n_components, simd_length>::get_tensor(
      const unsigned int row, const unsigned int position_within_column) const
  {
    Assert(sparsity != nullptr, dealii::ExcNotInitialized());
    AssertIndexRange(row, sparsity->row_starts.size() - 1);
    AssertIndexRange(position_within_column, sparsity->row_length(row));
 
    dealii::Tensor<1, n_components, Number2> result;
 
    if constexpr (std::is_same<Number, Number2>::value) {
      /*
       * Non-vectorized slow access. Supports all row indices in
       * [0,n_owned)
       */
      if (row < sparsity->n_internal_dofs) {
        // go through vectorized part
        const unsigned int simd_row = row / simd_length;
        const unsigned int simd_offset = row % simd_length;
        for (unsigned int d = 0; d < n_components; ++d)
          result[d] = data[(sparsity->row_starts[simd_row] +
                            position_within_column * simd_length) *
                               n_components +
                           d * simd_length + simd_offset];
      } else {
        // go through standard part
        for (unsigned int d = 0; d < n_components; ++d)
          result[d] =
              data[(sparsity->row_starts[row] + position_within_column) *
                       n_components +
                   d];
      }
 
    } else if constexpr (std::is_same<VectorizedArray, Number2>::value) {
      /*
       * Vectorized fast access. Indices must be in the range
       * [0,n_internal), index must be divisible by simd_length
       */
 
      Assert(row < sparsity->n_internal_dofs,
             dealii::ExcMessage(
                 "Vectorized access only possible in vectorized part"));
      Assert(row % simd_length == 0,
             dealii::ExcMessage(
                 "Access only supported for rows at the SIMD granularity"));
 
      const Number *load_pos =
          data.data() + (sparsity->row_starts[row / simd_length] +
                         position_within_column * simd_length) *
                            n_components;
 
      for (unsigned int d = 0; d < n_components; ++d)
        result[d].load(load_pos + d * simd_length);
 
    } else {
      /* not implemented */
      __builtin_trap();
    }
 
    return result;
  }
 
 
  template <typename Number, int n_components, int simd_length>
  template <typename Number2>
  DEAL_II_ALWAYS_INLINE inline auto
  SparseMatrixSIMD<Number, n_components, simd_length>::get_transposed_entry(
      const unsigned int row, const unsigned int position_within_column) const
      -> EntryType<Number2>
  {
    const auto result =
        get_transposed_tensor<Number2>(row, position_within_column);
    if constexpr (n_components == 1)
      return result[0];
    else
      return result;
  }
 
 
  template <typename Number, int n_components, int simd_length>
  template <typename Number2>
  DEAL_II_ALWAYS_INLINE inline dealii::Tensor<1, n_components, Number2>
  SparseMatrixSIMD<Number, n_components, simd_length>::get_transposed_tensor(
      const unsigned int row, const unsigned int position_within_column) const
  {
    Assert(sparsity != nullptr, dealii::ExcNotInitialized());
    AssertIndexRange(row, sparsity->row_starts.size() - 1);
    AssertIndexRange(position_within_column, sparsity->row_length(row));
 
    dealii::Tensor<1, n_components, Number2> result;
 
    if constexpr (std::is_same<Number, Number2>::value) {
      /*
       * Non-vectorized slow access. Supports all row indices in
       * [0,n_owned)
       */
 
      if (row < sparsity->n_internal_dofs) {
        // go through vectorized part
        const unsigned int simd_row = row / simd_length;
        const unsigned int simd_offset = row % simd_length;
        const std::size_t index =
            sparsity->indices_transposed[sparsity->row_starts[simd_row] +
                                         simd_offset +
                                         position_within_column * simd_length];
        if (n_components > 1) {
          const unsigned int col =
              sparsity->column_indices[sparsity->row_starts[simd_row] +
                                       simd_offset +
                                       position_within_column * simd_length];
          if (col < sparsity->n_internal_dofs)
            for (unsigned int d = 0; d < n_components; ++d)
              result[d] =
                  data[index / simd_length * simd_length * n_components +
                       simd_length * d + index % simd_length];
          else
            for (unsigned int d = 0; d < n_components; ++d)
              result[d] = data[index * n_components + d];
        } else
          result[0] = data[index];
      } else {
        // go through standard part
        const std::size_t index =
            sparsity->indices_transposed[sparsity->row_starts[row] +
                                         position_within_column];
        if (n_components > 1) {
          const unsigned int col =
              sparsity->column_indices[sparsity->row_starts[row] +
                                       position_within_column];
          if (col < sparsity->n_internal_dofs)
            for (unsigned int d = 0; d < n_components; ++d)
              result[d] =
                  data[index / simd_length * simd_length * n_components +
                       simd_length * d + index % simd_length];
          else
            for (unsigned int d = 0; d < n_components; ++d)
              result[d] = data[index * n_components + d];
        } else
          result[0] = data[index];
      }
 
    } else if constexpr (std::is_same<VectorizedArray, Number2>::value &&
                         (n_components == 1)) {
      /*
       * Vectorized fast access. Indices must be in the range
       * [0,n_internal), index must be divisible by simd_length
       */
 
      Assert(row < sparsity->n_internal_dofs,
             dealii::ExcMessage(
                 "Vectorized access only possible in vectorized part"));
      Assert(row % simd_length == 0,
             dealii::ExcMessage(
                 "Access only supported for rows at the SIMD granularity"));
 
      const unsigned int offset = sparsity->row_starts[row / simd_length] +
                                  position_within_column * simd_length;
      result[0].gather(data.data(),
                       sparsity->indices_transposed.data() + offset);
 
    } else {
      /* not implemented */
      __builtin_trap();
    }
 
    return result;
  }
 
 
  template <typename Number, int n_components, int simd_length>
  template <typename Number2>
  DEAL_II_ALWAYS_INLINE inline void
  SparseMatrixSIMD<Number, n_components, simd_length>::write_entry(
      const Number2 entry,
      const unsigned int row,
      const unsigned int position_within_column,
      const bool do_streaming_store)
  {
    static_assert(
        n_components == 1,
        "Attempted to write a scalar value into a tensor-valued matrix entry");
 
    Assert(sparsity != nullptr, dealii::ExcNotInitialized());
    AssertIndexRange(row, sparsity->row_starts.size() - 1);
    AssertIndexRange(position_within_column, sparsity->row_length(row));
 
    dealii::Tensor<1, n_components, Number2> result;
    result[0] = entry;
 
    write_entry<Number2>(
        result, row, position_within_column, do_streaming_store);
  }
 
 
  template <typename Number, int n_components, int simd_length>
  template <typename Number2>
  DEAL_II_ALWAYS_INLINE inline void
  SparseMatrixSIMD<Number, n_components, simd_length>::write_entry(
      const dealii::Tensor<1, n_components, Number2> &entry,
      const unsigned int row,
      const unsigned int position_within_column,
      const bool do_streaming_store)
  {
    Assert(sparsity != nullptr, dealii::ExcNotInitialized());
    AssertIndexRange(row, sparsity->row_starts.size() - 1);
    AssertIndexRange(position_within_column, sparsity->row_length(row));
 
    if constexpr (std::is_same<Number, Number2>::value) {
      /*
       * Non-vectorized slow access. Supports all row indices in
       * [0,n_owned)
       */
 
      if (row < sparsity->n_internal_dofs) {
        // go through vectorized part
        const unsigned int simd_row = row / simd_length;
        const unsigned int simd_offset = row % simd_length;
        for (unsigned int d = 0; d < n_components; ++d)
          data[(sparsity->row_starts[simd_row] +
                position_within_column * simd_length) *
                   n_components +
               d * simd_length + simd_offset] = entry[d];
      } else {
        // go through standard part
        for (unsigned int d = 0; d < n_components; ++d)
          data[(sparsity->row_starts[row] + position_within_column) *
                   n_components +
               d] = entry[d];
      }
 
    } else if constexpr (std::is_same<VectorizedArray, Number2>::value) {
      /*
       * Vectorized fast access. Indices must be in the range
       * [0,n_internal), index must be divisible by simd_length
       */
 
      Assert(row < sparsity->n_internal_dofs,
             dealii::ExcMessage(
                 "Vectorized access only possible in vectorized part"));
      Assert(row % simd_length == 0,
             dealii::ExcMessage(
                 "Access only supported for rows at the SIMD granularity"));
 
      Number *store_pos =
          data.data() + (sparsity->row_starts[row / simd_length] +
                         position_within_column * simd_length) *
                            n_components;
      if (do_streaming_store)
        for (unsigned int d = 0; d < n_components; ++d)
          entry[d].streaming_store(store_pos + d * simd_length);
      else
        for (unsigned int d = 0; d < n_components; ++d)
          entry[d].store(store_pos + d * simd_length);
 
    } else {
      /* not implemented */
      __builtin_trap();
    }
  }
 
 
  template <typename Number, int n_components, int simd_length>
  inline void
  SparseMatrixSIMD<Number, n_components, simd_length>::update_ghost_rows_start(
      const unsigned int communication_channel)
  {
#ifdef DEAL_II_WITH_MPI
    AssertIndexRange(communication_channel, 200);
 
    const unsigned int mpi_tag =
        dealii::Utilities::MPI::internal::Tags::partitioner_export_start +
        communication_channel;
    Assert(mpi_tag <=
               dealii::Utilities::MPI::internal::Tags::partitioner_export_end,
           dealii::ExcInternalError());
 
    const std::size_t n_indices = sparsity->entries_to_be_sent.size();
    exchange_buffer.resize_fast(n_components * n_indices);
 
    requests.resize(sparsity->receive_targets.size() +
                    sparsity->send_targets.size());
 
    /*
     * Set up MPI receive requests. We will always receive data for indices
     * in the range [n_locally_owned_, n_locally_relevant_), thus the DATA
     * is stored in non-vectorized CSR format.
     */
 
    {
      const auto &targets = sparsity->receive_targets;
      for (unsigned int p = 0; p < targets.size(); ++p) {
        const int ierr = MPI_Irecv(
            data.data() +
                n_components *
                    (sparsity->row_starts[sparsity->n_locally_owned_dofs] +
                     (p == 0 ? 0 : targets[p - 1].second)),
            (targets[p].second - (p == 0 ? 0 : targets[p - 1].second)) *
                n_components * sizeof(Number),
            MPI_BYTE,
            targets[p].first,
            mpi_tag,
            sparsity->mpi_communicator,
            &requests[p]);
        AssertThrowMPI(ierr);
      }
    }
 
    /*
     * Copy all entries that we plan to send over to the exchange buffer.
     * Here, we have to be careful with indices falling into the "locally
     * internal" range that are stored in an array-of-struct-of-array type.
     */
 
    for (std::size_t c = 0; c < n_indices; ++c) {
 
      const auto &[row, position_within_column] =
          sparsity->entries_to_be_sent[c];
 
      Assert(row < sparsity->n_locally_owned_dofs, dealii::ExcInternalError());
 
      if (row < sparsity->n_internal_dofs) {
        // go through vectorized part
        const unsigned int simd_row = row / simd_length;
        const unsigned int simd_offset = row % simd_length;
        for (unsigned int d = 0; d < n_components; ++d)
          exchange_buffer[n_components * c + d] =
              data[(sparsity->row_starts[simd_row] +
                    position_within_column * simd_length) *
                       n_components +
                   d * simd_length + simd_offset];
      } else {
        // go through standard part
        for (unsigned int d = 0; d < n_components; ++d)
          exchange_buffer[n_components * c + d] =
              data[(sparsity->row_starts[row] + position_within_column) *
                       n_components +
                   d];
      }
    }
 
    /*
     * Set up MPI send requests. We have copied everything we intend to
     * send to the exchange_buffer compatible with the CSR storage format
     * of the receiving MPI rank.
     */
 
    {
      const auto &targets = sparsity->send_targets;
      for (unsigned int p = 0; p < targets.size(); ++p) {
        const int ierr = MPI_Isend(
            exchange_buffer.data() +
                n_components * (p == 0 ? 0 : targets[p - 1].second),
            (targets[p].second - (p == 0 ? 0 : targets[p - 1].second)) *
                n_components * sizeof(Number),
            MPI_BYTE,
            targets[p].first,
            mpi_tag,
            sparsity->mpi_communicator,
            &requests[p + sparsity->receive_targets.size()]);
        AssertThrowMPI(ierr);
      }
    }
#endif
  }
 
 
  template <typename Number, int n_components, int simd_length>
  inline void SparseMatrixSIMD<Number, n_components, simd_length>::
      update_ghost_rows_finish()
  {
#ifdef DEAL_II_WITH_MPI
    const int ierr =
        MPI_Waitall(requests.size(), requests.data(), MPI_STATUSES_IGNORE);
    AssertThrowMPI(ierr);
#endif
  }
 
 
  template <typename Number, int n_components, int simd_length>
  inline void
  SparseMatrixSIMD<Number, n_components, simd_length>::update_ghost_rows()
  {
    update_ghost_rows_start();
    update_ghost_rows_finish();
  }
 
} // namespace ryujin