multicomponent_vector.h

Go to the documentation of this file.

//
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// Copyright (C) 2020 - 2023 by the ryujin authors
//
 
#pragma once
 
#include "simd.h"
 
#include <deal.II/base/mpi.h>
#include <deal.II/base/partitioner.h>
#include <deal.II/base/vectorization.h>
#include <deal.II/lac/la_parallel_vector.h>
 
namespace ryujin
{
  std::shared_ptr<const dealii::Utilities::MPI::Partitioner>
  create_vector_partitioner(
      const std::shared_ptr<const dealii::Utilities::MPI::Partitioner>
          &scalar_partitioner,
      const unsigned int n_components);
 
 
  template <typename Number,
            int n_comp,
            int simd_length = dealii::VectorizedArray<Number>::size()>
  class MultiComponentVector
      : public dealii::LinearAlgebra::distributed::Vector<Number>
  {
  public:
    using VectorizedArray = dealii::VectorizedArray<Number, simd_length>;
 
    using scalar_type = dealii::LinearAlgebra::distributed::Vector<Number>;
 
    using scalar_type::operator=;
 
    void reinit_with_scalar_partitioner(
        const std::shared_ptr<const dealii::Utilities::MPI::Partitioner>
            &scalar_partitioner);
 
    void extract_component(scalar_type &scalar_vector,
                           unsigned int component) const;
 
    void insert_component(const scalar_type &scalar_vector,
                          unsigned int component);
 
    template <typename Number2 = Number,
              typename Tensor = dealii::Tensor<1, n_comp, Number2>>
    Tensor get_tensor(const unsigned int i) const;
 
    template <typename Number2 = Number,
              typename Tensor = dealii::Tensor<1, n_comp, Number2>>
    Tensor get_tensor(const unsigned int *js) const;
 
    template <typename Number2 = Number,
              typename Tensor = dealii::Tensor<1, n_comp, Number2>>
    void write_tensor(const Tensor &tensor, const unsigned int i);
  };
 
 
#ifndef DOXYGEN
  /* Template definitions: */
 
  template <typename Number, int n_comp, int simd_length>
  void MultiComponentVector<Number, n_comp, simd_length>::
      reinit_with_scalar_partitioner(
          const std::shared_ptr<const dealii::Utilities::MPI::Partitioner>
              &scalar_partitioner)
  {
    /* Special case of a zero component vector */
    if (n_comp == 0)
      return;
 
    auto vector_partitioner =
        create_vector_partitioner(scalar_partitioner, n_comp);
 
    dealii::LinearAlgebra::distributed::Vector<Number>::reinit(
        vector_partitioner);
  }
 
 
  template <typename Number, int n_comp, int simd_length>
  void MultiComponentVector<Number, n_comp, simd_length>::extract_component(
      scalar_type &scalar_vector, unsigned int component) const
  {
    Assert(n_comp > 0,
           dealii::ExcMessage(
               "Cannot extract from a vector with zero components."));
 
    Assert(n_comp * scalar_vector.get_partitioner()->locally_owned_size() ==
               this->get_partitioner()->locally_owned_size(),
           dealii::ExcMessage("Called with a scalar_vector argument that has "
                              "incompatible local range."));
    const auto local_size =
        scalar_vector.get_partitioner()->locally_owned_size();
    for (unsigned int i = 0; i < local_size; ++i)
      scalar_vector.local_element(i) =
          this->local_element(i * n_comp + component);
    scalar_vector.update_ghost_values();
  }
 
 
  template <typename Number, int n_comp, int simd_length>
  void MultiComponentVector<Number, n_comp, simd_length>::insert_component(
      const scalar_type &scalar_vector, unsigned int component)
  {
    Assert(n_comp > 0,
           dealii::ExcMessage(
               "Cannot insert into a vector with zero components."));
 
    Assert(n_comp * scalar_vector.get_partitioner()->locally_owned_size() ==
               this->get_partitioner()->locally_owned_size(),
           dealii::ExcMessage("Called with a scalar_vector argument that has "
                              "incompatible local range."));
    const auto local_size =
        scalar_vector.get_partitioner()->locally_owned_size();
    for (unsigned int i = 0; i < local_size; ++i)
      this->local_element(i * n_comp + component) =
          scalar_vector.local_element(i);
  }
 
  /* Inline function  definitions: */
 
  template <typename Number, int n_comp, int simd_length>
  template <typename Number2, typename Tensor>
  DEAL_II_ALWAYS_INLINE inline Tensor
  MultiComponentVector<Number, n_comp, simd_length>::get_tensor(
      const unsigned int i) const
  {
    static_assert(std::is_same<Number2, typename Tensor::value_type>::value,
                  "dummy type mismatch");
    Tensor tensor;
 
    /* Special case of a zero component vector */
    if constexpr (n_comp == 0)
      return tensor;
 
    if constexpr (std::is_same<Number, Number2>::value) {
      /* Non-vectorized sequential access. */
 
      for (unsigned int d = 0; d < n_comp; ++d)
        tensor[d] = this->local_element(i * n_comp + d);
 
    } else if constexpr (std::is_same<VectorizedArray, Number2>::value) {
 
      /* Vectorized fast access. index must be divisible by simd_length */
      std::array<unsigned int, VectorizedArray::size()> indices;
      for (unsigned int k = 0; k < VectorizedArray::size(); ++k)
        indices[k] = k * n_comp;
 
      dealii::vectorized_load_and_transpose(
          n_comp, this->begin() + i * n_comp, indices.data(), &tensor[0]);
 
    } else {
      /* not implemented */
      __builtin_trap();
    }
 
    return tensor;
  }
 
 
  template <typename Number, int n_comp, int simd_length>
  template <typename Number2, typename Tensor>
  DEAL_II_ALWAYS_INLINE inline Tensor
  MultiComponentVector<Number, n_comp, simd_length>::get_tensor(
      const unsigned int *js) const
  {
    static_assert(std::is_same<Number2, typename Tensor::value_type>::value,
                  "dummy type mismatch");
    Tensor tensor;
 
    /* Special case of a zero component vector */
    if constexpr (n_comp == 0)
      return tensor;
 
    if constexpr (std::is_same<Number, Number2>::value) {
      /* Non-vectorized sequential access. */
 
      for (unsigned int d = 0; d < n_comp; ++d)
        tensor[d] = this->local_element(js[0] * n_comp + d);
 
    } else if constexpr (std::is_same<VectorizedArray, Number2>::value) {
      /* Vectorized fast access. index must be divisible by simd_length */
 
      std::array<unsigned int, VectorizedArray::size()> indices;
      for (unsigned int k = 0; k < VectorizedArray::size(); ++k)
        indices[k] = js[k] * n_comp;
 
      dealii::vectorized_load_and_transpose(
          n_comp, this->begin(), indices.data(), &tensor[0]);
 
    } else {
      /* not implemented */
      __builtin_trap();
    }
 
    return tensor;
  }
 
 
  template <typename Number, int n_comp, int simd_length>
  template <typename Number2, typename Tensor>
  DEAL_II_ALWAYS_INLINE inline void
  MultiComponentVector<Number, n_comp, simd_length>::write_tensor(
      const Tensor &tensor, const unsigned int i)
  {
    static_assert(std::is_same<Number2, typename Tensor::value_type>::value,
                  "dummy type mismatch");
 
    /* Special case of a zero component vector */
    if constexpr (n_comp == 0)
      return;
 
    if constexpr (std::is_same<Number, Number2>::value) {
      /* Non-vectorized sequential access. */
 
      for (unsigned int d = 0; d < n_comp; ++d)
        this->local_element(i * n_comp + d) = tensor[d];
 
    } else if constexpr (std::is_same<VectorizedArray, Number2>::value) {
      /* Vectorized fast access. index must be divisible by simd_length */
 
      std::array<unsigned int, VectorizedArray::size()> indices;
      for (unsigned int k = 0; k < VectorizedArray::size(); ++k)
        indices[k] = k * n_comp;
 
      dealii::vectorized_transpose_and_store(false,
                                             n_comp,
                                             &tensor[0],
                                             indices.data(),
                                             this->begin() + i * n_comp);
 
    } else {
      /* not implemented */
      __builtin_trap();
    }
  }
#endif
 
} // namespace ryujin