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
{
  namespace Vectors
  {
    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 ScalarVector = dealii::LinearAlgebra::distributed::Vector<Number>;
 
      using ScalarVector::operator=;
 
      void reinit_with_scalar_partitioner(
          const std::shared_ptr<const dealii::Utilities::MPI::Partitioner>
              &scalar_partitioner);
 
      void extract_component(ScalarVector &scalar_vector,
                             unsigned int component) const;
 
      void insert_component(const ScalarVector &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);
 
      template <typename Number2 = Number,
                typename Tensor = dealii::Tensor<1, n_comp, Number2>>
      void add_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(
        ScalarVector &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 ScalarVector &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(/*add into*/ false,
                                               n_comp,
                                               &tensor[0],
                                               indices.data(),
                                               this->begin() + i * n_comp);
 
      } else {
        /* not implemented */
        __builtin_trap();
      }
    }
 
 
    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>::add_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(/*add into*/ true,
                                               n_comp,
                                               &tensor[0],
                                               indices.data(),
                                               this->begin() + i * n_comp);
 
      } else {
        /* not implemented */
        __builtin_trap();
      }
    }
#endif
  } // namespace Vectors
} // namespace ryujin