kokkos_linalg_base.h

Go to the documentation of this file.

// The libMesh Finite Element Library.
// Copyright (C) 2002-2026 Benjamin S. Kirk, John W. Peterson, Roy H. Stogner
 
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
 
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
 
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
#ifndef LIBMESH_KOKKOS_LINALG_BASE_H
#define LIBMESH_KOKKOS_LINALG_BASE_H
 
// libMesh Kokkos compile-time linalg foundation.
//
// This header defines the small access/materialization layer that sits
// underneath richer vector/tensor algebra. It is intentionally limited to
// component access, storage-backed references, and conversion between
// vector-like/tensor-like objects and libMesh semantic types.
//
// The key abstractions are:
//
// - "vector-like"/"tensor-like": the types registered with the
//   is_vector_like/is_tensor_like traits, whose components are read
//   through operator()(i) / operator()(i,j).  libMesh's owning types
//   (TypeVector, TypeTensor, Point, ...) and the non-owning refs below
//   are registered, so the algorithms in kokkos_vector_ops.h and
//   kokkos_tensor_ops.h can be written once against either kind.
//
// - vector_ref/tensor_ref: lightweight non-owning handles addressing
//   one entry of a Kokkos storage view (see kokkos_storage_policy.h).
//   Reads and writes through a ref go directly to the underlying view,
//   so kernels can update storage in place without materializing
//   temporaries.
//
// - materialize_vector()/materialize_tensor(): copy any vector-like or
//   tensor-like object into an owning libMesh value type, for use in
//   register-resident arithmetic.
 
#include "libmesh/libmesh_common.h"
#include "libmesh/libmesh_device.h"
#include "libmesh/point.h"
#include "libmesh/tensor_value.h"
#include "libmesh/type_tensor.h"
#include "libmesh/type_vector.h"
 
#include <type_traits>
#include <utility>
 
namespace libMesh::Kokkos
{
 
namespace detail
{
 
template <typename T>
using remove_cvref_t =
  typename std::remove_cv<typename std::remove_reference<T>::type>::type;
 
template <typename ViewType>
using vector_view_value_t =
  remove_cvref_t<decltype(std::declval<ViewType &>()(0, 0))>;
 
template <typename ViewType>
using tensor_view_value_t =
  remove_cvref_t<decltype(std::declval<ViewType &>()(0, 0, 0))>;
 
} // namespace detail
 
template <typename T>
struct vector_traits;
 
template <typename T>
struct tensor_traits;
 
template <typename T>
struct is_vector_like : std::false_type
{
};
 
template <typename T>
struct is_tensor_like : std::false_type
{
};
 
template <typename T>
struct is_vector_ref : std::false_type
{
};
 
template <typename T>
struct is_tensor_ref : std::false_type
{
};
 
template <typename T>
inline constexpr bool is_vector_like_v = is_vector_like<detail::remove_cvref_t<T>>::value;
 
template <typename T>
inline constexpr bool is_tensor_like_v = is_tensor_like<detail::remove_cvref_t<T>>::value;
 
template <typename T>
inline constexpr bool is_vector_ref_v = is_vector_ref<detail::remove_cvref_t<T>>::value;
 
template <typename T>
inline constexpr bool is_tensor_ref_v = is_tensor_ref<detail::remove_cvref_t<T>>::value;
 
template <typename ViewType>
class vector_ref
{
public:
  // The underlying view type and the scalar type of its components.
  using view_type = ViewType;
  using value_type = detail::vector_view_value_t<ViewType>;
 
  LIBMESH_DEVICE_INLINE
  vector_ref(ViewType view, const unsigned int index) : _view(view), _index(index) {}
 
  LIBMESH_DEVICE_INLINE
  decltype(auto) operator()(const unsigned int component) const
  {
    return _view(_index, component);
  }
 
  LIBMESH_DEVICE_INLINE
  decltype(auto) operator()(const unsigned int component)
  {
    return _view(_index, component);
  }
 
  template <typename RightVector>
  LIBMESH_DEVICE_INLINE
  void assign(const RightVector & right);
 
  template <typename RightVector>
  LIBMESH_DEVICE_INLINE
  void add(const RightVector & right);
 
  template <typename RightVector>
  LIBMESH_DEVICE_INLINE
  void add_scaled(const RightVector & right, const value_type & factor);
 
  template <typename RightVector>
  LIBMESH_DEVICE_INLINE
  void subtract(const RightVector & right);
 
  template <typename RightVector>
  LIBMESH_DEVICE_INLINE
  void subtract_scaled(const RightVector & right, const value_type & factor);
 
  LIBMESH_DEVICE_INLINE
  void zero();
 
  template <typename RightVector>
  LIBMESH_DEVICE_INLINE
  auto contract(const RightVector & right) const;
 
  LIBMESH_DEVICE_INLINE
  auto norm() const;
 
  LIBMESH_DEVICE_INLINE
  auto norm_sq() const;
 
  LIBMESH_DEVICE_INLINE
  auto l1_norm() const;
 
  LIBMESH_DEVICE_INLINE
  bool is_zero() const;
 
  LIBMESH_DEVICE_INLINE
  auto unit() const;
 
  template <typename RightVector>
  LIBMESH_DEVICE_INLINE
  auto cross(const RightVector & right) const;
 
  LIBMESH_DEVICE_INLINE
  unsigned int index() const
  {
    return _index;
  }
 
private:
  ViewType _view;
  unsigned int _index;
};
 
template <typename ViewType>
class tensor_ref
{
public:
  // The underlying view type and the scalar type of its components.
  using view_type = ViewType;
  using value_type = detail::tensor_view_value_t<ViewType>;
 
  LIBMESH_DEVICE_INLINE
  tensor_ref(ViewType view, const unsigned int index) : _view(view), _index(index) {}
 
  LIBMESH_DEVICE_INLINE
  decltype(auto) operator()(const unsigned int row, const unsigned int col) const
  {
    return _view(_index, row, col);
  }
 
  LIBMESH_DEVICE_INLINE
  decltype(auto) operator()(const unsigned int row, const unsigned int col)
  {
    return _view(_index, row, col);
  }
 
  template <typename RightTensor>
  LIBMESH_DEVICE_INLINE
  void assign(const RightTensor & right);
 
  template <typename RightTensor>
  LIBMESH_DEVICE_INLINE
  void add(const RightTensor & right);
 
  template <typename RightTensor>
  LIBMESH_DEVICE_INLINE
  void add_scaled(const RightTensor & right, const value_type & factor);
 
  template <typename RightTensor>
  LIBMESH_DEVICE_INLINE
  void subtract(const RightTensor & right);
 
  template <typename RightTensor>
  LIBMESH_DEVICE_INLINE
  void subtract_scaled(const RightTensor & right, const value_type & factor);
 
  LIBMESH_DEVICE_INLINE
  void zero();
 
  template <typename RightTensor>
  LIBMESH_DEVICE_INLINE
  auto contract(const RightTensor & right) const;
 
  LIBMESH_DEVICE_INLINE
  auto norm() const;
 
  LIBMESH_DEVICE_INLINE
  auto norm_sq() const;
 
  LIBMESH_DEVICE_INLINE
  bool is_zero() const;
 
  LIBMESH_DEVICE_INLINE
  auto transpose() const;
 
  LIBMESH_DEVICE_INLINE
  auto det(const unsigned int dim = LIBMESH_DIM) const;
 
  LIBMESH_DEVICE_INLINE
  auto tr() const;
 
  LIBMESH_DEVICE_INLINE
  auto inverse(const unsigned int dim = LIBMESH_DIM) const;
 
  template <typename VectorLike, typename ResultVector>
  LIBMESH_DEVICE_INLINE
  void solve(const VectorLike & b, ResultVector & x) const;
 
  LIBMESH_DEVICE_INLINE
  auto row(const unsigned int i) const;
 
  LIBMESH_DEVICE_INLINE
  auto column(const unsigned int i) const;
 
  template <typename VectorLike>
  LIBMESH_DEVICE_INLINE
  auto left_multiply(const VectorLike & v) const;
 
  LIBMESH_DEVICE_INLINE
  unsigned int index() const
  {
    return _index;
  }
 
private:
  ViewType _view;
  unsigned int _index;
};
 
// Owning vector types are their own semantic type; a vector_ref
// materializes to the TypeVector of its component type.
 
template <typename T>
struct vector_traits<libMesh::TypeVector<T>>
{
  using semantic_type = libMesh::TypeVector<T>;
};
 
template <typename T>
struct vector_traits<libMesh::VectorValue<T>>
{
  using semantic_type = libMesh::VectorValue<T>;
};
 
template <>
struct vector_traits<libMesh::Point>
{
  using semantic_type = libMesh::Point;
};
 
template <typename ViewType>
struct vector_traits<vector_ref<ViewType>>
{
  using value_type = typename vector_ref<ViewType>::value_type;
  using semantic_type = libMesh::TypeVector<value_type>;
};
 
// The vector-like family: libMesh owning vector types plus vector_ref.
 
template <typename T>
struct is_vector_like<libMesh::TypeVector<T>> : std::true_type
{
};
 
template <typename T>
struct is_vector_like<libMesh::VectorValue<T>> : std::true_type
{
};
 
template <>
struct is_vector_like<libMesh::Point> : std::true_type
{
};
 
template <typename ViewType>
struct is_vector_like<vector_ref<ViewType>> : std::true_type
{
};
 
template <typename ViewType>
struct is_vector_ref<vector_ref<ViewType>> : std::true_type
{
};
 
// Owning tensor types are their own semantic type; a tensor_ref
// materializes to the TypeTensor of its component type.
 
template <typename T>
struct tensor_traits<libMesh::TypeTensor<T>>
{
  using semantic_type = libMesh::TypeTensor<T>;
};
 
template <typename T>
struct tensor_traits<libMesh::TensorValue<T>>
{
  using semantic_type = libMesh::TensorValue<T>;
};
 
template <typename ViewType>
struct tensor_traits<tensor_ref<ViewType>>
{
  using value_type = typename tensor_ref<ViewType>::value_type;
  using semantic_type = libMesh::TypeTensor<value_type>;
};
 
// The tensor-like family: libMesh owning tensor types plus tensor_ref.
 
template <typename T>
struct is_tensor_like<libMesh::TypeTensor<T>> : std::true_type
{
};
 
template <typename T>
struct is_tensor_like<libMesh::TensorValue<T>> : std::true_type
{
};
 
template <typename ViewType>
struct is_tensor_like<tensor_ref<ViewType>> : std::true_type
{
};
 
template <typename ViewType>
struct is_tensor_ref<tensor_ref<ViewType>> : std::true_type
{
};
 
template <typename T>
using vector_semantic_type_t = typename vector_traits<detail::remove_cvref_t<T>>::semantic_type;
 
template <typename T>
using tensor_semantic_type_t = typename tensor_traits<detail::remove_cvref_t<T>>::semantic_type;
 
template <typename ViewType>
LIBMESH_DEVICE_INLINE
vector_ref<typename std::remove_reference<ViewType>::type>
make_vector_ref(ViewType && view, const unsigned int index)
{
  return vector_ref<typename std::remove_reference<ViewType>::type>(std::forward<ViewType>(view),
                                                                    index);
}
 
template <typename ViewType>
LIBMESH_DEVICE_INLINE
tensor_ref<typename std::remove_reference<ViewType>::type>
make_tensor_ref(ViewType && view, const unsigned int index)
{
  return tensor_ref<typename std::remove_reference<ViewType>::type>(std::forward<ViewType>(view),
                                                                    index);
}
 
template <typename OutputVector, typename VectorLike>
LIBMESH_DEVICE_INLINE
OutputVector materialize_vector(const VectorLike & v)
{
  static_assert(is_vector_like_v<OutputVector>,
                "materialize_vector() requires a vector-like output type");
  static_assert(is_vector_like_v<VectorLike>,
                "materialize_vector() requires a vector-like input type");
 
  OutputVector out;
  out.zero();
 
  for (unsigned int component = 0; component < LIBMESH_DIM; ++component)
    out(component) = v(component);
 
  return out;
}
 
template <typename OutputTensor, typename TensorLike>
LIBMESH_DEVICE_INLINE
OutputTensor materialize_tensor(const TensorLike & T_in)
{
  static_assert(is_tensor_like_v<OutputTensor>,
                "materialize_tensor() requires a tensor-like output type");
  static_assert(is_tensor_like_v<TensorLike>,
                "materialize_tensor() requires a tensor-like input type");
 
  OutputTensor out;
  out.zero();
 
  for (unsigned int row = 0; row < LIBMESH_DIM; ++row)
    for (unsigned int col = 0; col < LIBMESH_DIM; ++col)
      out(row, col) = T_in(row, col);
 
  return out;
}
 
} // namespace libMesh::Kokkos
 
#endif // LIBMESH_KOKKOS_LINALG_BASE_H