Merge branch 'master' into mdspan

fem/CMakeLists.txt

@@ -101,6 +101,7 @@ set(SRCS
   lor/lor_ads.cpp
   lor/lor_ams.cpp
   lor/lor_batched.cpp
+  mdgridfunc.hpp
   multigrid.cpp
   nonlinearform.cpp
   nonlinearform_ext.cpp

fem/mdgridfunc.hpp

@@ -0,0 +1,154 @@
+// Copyright (c) 2010-2023, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#ifndef MFEM_MDGRIDFUNC
+#define MFEM_MDGRIDFUNC
+
+#include "../config/config.hpp"
+
+#include "fem/gridfunc.hpp"
+#include "general/mdspan.hpp"
+
+namespace mfem
+{
+
+template<int N, class Layout = MDLayoutLeft<N>>
+class MDGridFunction : public MDSpan<GridFunction, N, Layout>
+{
+   using base_t = MDSpan<GridFunction, N, Layout>;
+   using base_t::Nd;
+   using base_t::Sd;
+   using GridFunction::data;
+
+public:
+
+   /**
+    * @brief MDGridFunction default constructor (recursion)
+    */
+   MDGridFunction(): base_t() { }
+
+   /**
+    * @brief MDGridFunction recursion constructor
+    * @param[in] fes Finite element space to use
+    * @param[in] args Rest of dimension indices
+    */
+   template <typename... Ts>
+   MDGridFunction(FiniteElementSpace *fes, Ts... args): MDGridFunction(args...)
+   {
+      SetSpace(fes);
+      MFEM_VERIFY(fes->GetVDim() == 1,
+                  "Only FiniteElementSpace with vdim of 1 are supported");
+      base_t::Setup(fes->GetNDofs(), args...);
+   }
+
+   /**
+    * @brief MDGridFunction recursion constructor
+    * @param[in] dim Dimension indice
+    * @param[in] args Rest of dimension indices or finite element space to use
+    */
+   template <typename... Ts>
+   MDGridFunction(int dim, Ts... args): MDGridFunction(args...)
+   {
+      base_t::Setup(dim, args...);
+   }
+
+   /// Move constructor not supported
+   MDGridFunction(MDGridFunction&&) = delete;
+
+   /// Copy constructor not supported
+   MDGridFunction(const MDGridFunction&) = delete;
+
+   /// Move assignment not supported
+   MDGridFunction& operator=(MDGridFunction&&) = delete;
+
+   /// Copy assignment not supported
+   MDGridFunction& operator=(const MDGridFunction&) = delete;
+
+   /**
+    * @brief Returns the specific GridFunction from dimension indices
+    * @param[out] gf Returned GridFunction
+    * @param[in] args Rest of dimension indices
+    */
+   template <int n = 1, typename... Ts>
+   void GetScalarGridFunction(GridFunction &gf, Ts... args) const
+   {
+      FiniteElementSpace *fes = GridFunction::fes;
+      MFEM_VERIFY(fes->GetNDofs() == Nd[n-1], "Error in dofs size!");
+      gf.SetSpace(fes);
+      for (int s = 0; s < Nd[n-1]; s++)
+      {
+         gf[s] = data[get_vdofs_offset +
+                      MDOffset<n,N,int,Ts...>::offset(Sd, s, args...)];
+      }
+      get_vdofs_offset = 0; // re-init for next calls
+   }
+
+   /**
+    * @brief Returns the specific GridFunction from dimension indices
+    * @param[in] dim Dimension indice
+    * @param args Rest of dimension indices or GridFunction to be returned
+    */
+   template <int n = 1, typename... Ts>
+   void GetScalarGridFunction(int dim, Ts&&... args) const
+   {
+      get_vdofs_offset += dim * Sd[n-1];
+      MDGridFunction::GetScalarGridFunction<n+1>(std::forward<Ts>(args)...);
+   }
+
+   /**
+    * @brief Sets the given GridFunction at the specific dimension indices
+    * @param[in] gf GridFunction to set
+    * @param[in] args Rest of dimension indices
+    */
+   template <int n = 1, typename... Ts>
+   void SetScalarGridFunction(const GridFunction &gf, Ts... args)
+   {
+      MFEM_VERIFY(GridFunction::fes->GetNDofs() == Nd[n-1], "Error in dofs size!");
+      for (int s = 0; s < Nd[n-1]; s++)
+      {
+         data[get_vdofs_offset +
+              MDOffset<n,N,int,Ts...>::offset(Sd, s, args...)] = gf[s];
+      }
+      get_vdofs_offset = 0; // re-init for next calls
+   }
+
+   /**
+    * @brief Sets the given GridFunction at the specific dimension indices
+    * @param[in] dim Dimension indice
+    * @param args Rest of dimension indices or given GridFunction to be used
+    */
+   template <int n = 1, typename... Ts>
+   void SetScalarGridFunction(int dim, Ts... args)
+   {
+      get_vdofs_offset += dim * Sd[n-1];
+      MDGridFunction::SetScalarGridFunction<n+1>(args...);
+   }
+
+   using GridFunction::Read;
+   using GridFunction::Write;
+   using GridFunction::ReadWrite;
+   using GridFunction::HostRead;
+   using GridFunction::HostWrite;
+   using GridFunction::HostReadWrite;
+
+   using GridFunction::GetData;
+   using GridFunction::SetData;
+   using GridFunction::SetSpace;
+
+   using Vector::operator=;
+
+private:
+   mutable int get_vdofs_offset = 0;
+};
+
+} // namespace mfem
+
+#endif // MFEM_MDGRIDFUNC

general/CMakeLists.txt

@@ -47,6 +47,8 @@ list(APPEND HDRS
   zstr.hpp
   hash.hpp
   isockstream.hpp
+  mdarray.hpp
+  mdspan.hpp
   kdtree.hpp
   mem_alloc.hpp
   mem_manager.hpp

general/array.hpp

@@ -45,6 +45,8 @@ template <class T>
 class Array
 {
 protected:
+   template<typename mfem_type, int N, typename L> friend class MDSpan;
+
    /// Pointer to data
    Memory<T> data;
    /// Size of the array

general/mdarray.hpp

@@ -0,0 +1,70 @@
+// Copyright (c) 2010-2023, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#ifndef MFEM_MDARRAY
+#define MFEM_MDARRAY
+
+#include "../config/config.hpp"
+
+#include "array.hpp"
+#include "mdspan.hpp"
+
+namespace mfem
+{
+
+template<typename T, int N, typename Layout = MDLayoutLeft<N>>
+struct MDArray : public MDSpan<Array<T>, N, Layout>
+{
+   using base_t = MDSpan<Array<T>, N, Layout>;
+
+   /**
+    * @brief MDArray default constructor (recursion)
+    */
+   MDArray(): base_t() { }
+
+   /**
+    * @brief MDArray recursion constructor
+    * @param[in] n Dimension indice
+    * @param[in] args Rest of dimension indices
+    */
+   template <typename... Ts>
+   MDArray(int n, Ts... args): MDArray(args...) { base_t::Setup(n, args...); }
+
+   /// Move constructor not supported
+   MDArray(MDArray&&) = delete;
+
+   /// Copy constructor not supported
+   MDArray(const MDArray&) = delete;
+
+   /// Move assignment not supported
+   MDArray& operator=(MDArray&&) = delete;
+
+   /// Copy assignment not supported
+   MDArray& operator=(const MDArray&) = delete;
+
+   using Array<T>::Read;
+   using Array<T>::Write;
+   using Array<T>::ReadWrite;
+   using Array<T>::HostRead;
+   using Array<T>::HostWrite;
+   using Array<T>::HostReadWrite;
+
+   using Array<T>::Assign;
+   using Array<T>::Print;
+
+   using Array<T>::GetData;
+
+   using Array<T>::operator=;
+};
+
+} // namespace mfem
+
+#endif // MFEM_MDARRAY

general/mdspan.hpp

@@ -0,0 +1,417 @@
+// Copyright (c) 2010-2023, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#ifndef MFEM_MDSPAN_HPP
+#define MFEM_MDSPAN_HPP
+
+#include <list>
+#include <array>
+#include <vector>
+#include <utility>
+#include <type_traits>
+
+#include "device.hpp"
+#include "backends.hpp"
+
+namespace mfem
+{
+
+namespace internal // experimental helper functions for mfem::MDLayout
+{
+
+// md_sequence represents a compile-time sequence of integers
+template <typename T, T... args> struct md_sequence { };
+
+template <typename T, int N, bool left> struct make_md_sequence;
+
+// make_sequence, specialized for left (default) and right layout
+template <typename T, int N, bool left = true>
+using make_sequence = typename make_md_sequence<T, N, left>::type;
+
+} // namespace internal
+
+/// @brief The MDOffset class computes the multi-dimensional offsets
+template <int n, int N, typename T, typename... Ts>
+struct MDOffset
+{
+   static MFEM_HOST_DEVICE inline
+   T offset(const int (&Sd)[N], T nd, Ts... args)
+   { return nd * Sd[n-1] + MDOffset<n+1, N, Ts...>::offset(Sd, args...); }
+};
+
+template <int N, typename T, typename... Ts>
+struct MDOffset<N, N, T, Ts...>
+{
+   static MFEM_HOST_DEVICE inline
+   T offset(const int (&Sd)[N], T nd) { return nd * Sd[N-1]; }
+};
+
+/// @brief The MDTensor class holds the pointer and strides for each dimension
+template<int N, typename T> class MDTensor
+{
+   T *ptr;
+   int Sd[N];
+
+public:
+   /// Default constructor
+   MDTensor() = delete;
+
+   /// Copy constructor (default)
+   MDTensor(const MDTensor&) = default;
+
+   /// Copy assignment (default)
+   MDTensor& operator=(const MDTensor&) = default;
+
+   /// Constructor to initialize a tensor from a pointer and strides
+   template <typename... Args> MFEM_HOST_DEVICE
+   MDTensor(T *ptr, const int (&sd)[N]): ptr(ptr)
+   { for (int i = 0; i < N; ++i) { Sd[i] = sd[i]; } }
+
+   /// Accessor for the data
+   template <typename... Ts> MFEM_HOST_DEVICE inline
+   T& operator()(Ts... args) { return ptr[Offset(args...)]; }
+
+   /// Const accessor for the data
+   template <typename... Ts> MFEM_HOST_DEVICE inline
+   T& operator()(Ts... args) const { return ptr[Offset(args...)]; }
+
+   /// Offset computation
+   template <typename... Ts> MFEM_HOST_DEVICE inline
+   int Offset(Ts... args) const
+   {
+      static_assert(sizeof...(args) == N, "Wrong number of dimensions");
+      return MDOffset<1, N, Ts...>::offset(Sd, args...);
+   }
+};
+
+/// \brief The MDLayout class, defaulted to a column-major (left) ordering
+template<int N, bool left = true> struct MDLayout
+{
+   /// Create a layout with the internal::md_sequence
+   template <int... args>
+   static constexpr auto Make(internal::md_sequence<int, args...>)
+   -> std::array<int, sizeof...(args)> { return {(static_cast<int>(args))...}; }
+
+   /// Array holding the layout permutation
+   using perm_type = std::array<int, N>;
+   perm_type perm = Make(internal::make_sequence<int, N, left> {});
+
+   /// Default constructor
+   MDLayout() = default;
+
+   /// Copy constructor (default)
+   MDLayout(const MDLayout&) = default;
+
+   /// Copy assignment (default)
+   MDLayout& operator=(const MDLayout&) = default;
+
+   /// Constructor to initialize a layout from an array of indices
+   template <typename... Ts>
+   MDLayout(int n, Ts... args) noexcept: MDLayout(args...)
+   {
+      constexpr int k = N - sizeof...(args) - 1;
+      static_assert(0 <= k && k < N, "Index out of bounds!");
+      perm[k] = n;
+   }
+
+   /// Access layout entries using operator()
+   inline int operator()(int i) const
+   { return Assert(i), perm[static_cast<typename perm_type::size_type>(i)]; }
+
+   /// Access layout entries using operator[]
+   inline int operator[](int i) const
+   { return Assert(i), perm[static_cast<typename perm_type::size_type>(i)]; }
+
+   /// Asserts the given index is valid (only in MFEM_DEBUG)
+   inline void Assert(const int k) const
+   {
+      MFEM_CONTRACT_VAR(k);
+      MFEM_ASSERT(0 <= k && k < N, "Index should be in [0," << (N-1) << "]");
+   }
+};
+
+/// Left (Column-major (Fortran)) and Right (Row-major (C/C++)) layouts
+template<int N> using MDLayoutLeft = MDLayout<N, true>;
+template<int N> using MDLayoutRight = MDLayout<N, false>;
+
+/// \brief The MDSpan base class is a generic non-owning mfem_type's view
+/// that reinterprets it as a multidimensional type.
+template<typename mfem_type, int N, class layout_type = MDLayoutLeft<N>>
+class MDSpan : protected mfem_type
+{
+protected:
+   using T = typename std::remove_pointer<decltype(mfem_type::data.h_ptr)>::type;
+
+   int Nd[N], Sd[N]; // dimension sizes and strides, once the layout is set
+   layout_type layout; // stored layout, useful for reshapes
+
+   /// Set the dimensions (Nd) and strides (Sd) during contruction.
+   /// When all the arguments have been processed, SetSize is called on the
+   /// mfem_type with Device::GetMemoryType() as memory type and SetLayout is
+   /// called using the layout.
+   template <typename... Ts> void Setup(int dim, Ts... args)
+   {
+      constexpr int k = N - sizeof...(args) - 1;
+      Sd[k] = Nd[k] = dim;
+      if (k > 0) { return; }
+      int psize = 1;
+      for (int i = 0; i < N; i++) { psize *= Nd[i]; }
+      mfem_type::SetSize(static_cast<int>(psize), Device::GetMemoryType());
+      SetLayout(layout);
+   }
+
+public:
+
+   /// Default constructor (recursion)
+   MDSpan() noexcept: mfem_type() { }
+
+   /// Recursion constructor
+   template <typename... Ts>
+   MDSpan(int n, Ts... args): MDSpan(args...) { Setup(n, args...); }
+
+   /// Move constructor (delete)
+   MDSpan(MDSpan&&) = delete;
+
+   /// Copy constructor (delete)
+   MDSpan(const MDSpan&) = delete;
+
+   /// Move assignment (delete)
+   MDSpan& operator=(MDSpan&&) = delete;
+
+   /// Copy assignment (delete)
+   MDSpan& operator=(const MDSpan&) = delete;
+
+   /// Return the ith dimension
+   int Extent(int i) const { return Nd[i]; }
+
+   /// Return the size of the span.
+   int Size() const { return mfem_type::Size(); }
+
+   /// Store and use the given layout to update the strides
+   template<typename Layout> void SetLayout(const Layout &l)
+   {
+      layout = l;
+      Sd[l[0]] = 1;
+      for (int i = 1; i < N; i++) { Sd[l[i]] = Nd[l[i-1]] * Sd[l[i-1]]; }
+   }
+
+   /// Variadic resize the mfem_type
+   template <typename... Ts> inline void SetSize(int size, Ts... args)
+   {
+      constexpr int k = N - sizeof...(args) - 1;
+      Sd[k] = Nd[k] = size;
+      const int msize = mfem_type::Size();
+      MFEM_VERIFY(size > 0, "Size should be positive!");
+      mfem_type::SetSize(msize > 0 ? msize*size : size, Device::GetMemoryType());
+      MDSpan::SetSize(args...);
+   }
+
+   /// Variadic terminal case of the mfem_type resize
+   inline void SetSize(int size)
+   {
+      Sd[N-1] = Nd[N-1] = size;
+      const int msize = mfem_type::Size();
+      MFEM_VERIFY(size > 0, "Size should be positive!");
+      mfem_type::SetSize(msize > 0 ? msize*size : size, Device::GetMemoryType());
+      SetLayout(layout);
+   }
+
+   /// Access mfem_type data entries using operator()
+   template <typename... Ts> inline
+   T& operator()(Ts... args) { return mfem_type::data[Offset(args...)]; }
+
+   /// Const access mfem_type data entries using operator()
+   template <typename... Ts> inline const T& operator()(Ts... args) const
+   {
+      return mfem_type::data[Offset(args...)];
+   }
+
+   /// Offset computation
+   template <typename... Ts> inline int Offset(Ts... args) const
+   {
+      static_assert(sizeof...(args) == N, "Wrong number of dimensions");
+      return MDOffset<1,N,Ts...>::offset(Sd, args...);
+   }
+
+   /// Shortcut for mfem::Read(mfem_type::data, mfem_type::size, on_dev)
+   /// and return an MDTensor with the MDSpan's pointer and strides
+   const MDTensor<N,const T> MDRead(bool on_dev = true) const
+   {
+      const T *ptr = mfem::Read(mfem_type::data, mfem_type::size, on_dev);
+      return MDTensor<N,const T>(ptr, Sd);
+   }
+
+   /// Shortcut for mfem::Read(mfem_type::data, mfem_type::size, false)
+   /// and return an MDTensor with the MDSpan's pointer and strides
+   const MDTensor<N,const T> MDHostRead() const
+   {
+      const T *ptr = mfem::Read(mfem_type::data, mfem_type::size, false);
+      return MDTensor<N,const T>(ptr, Sd);
+   }
+
+   /// Shortcut for mfem::Write(mfem_type::data, mfem_type::size, on_dev)
+   /// and return an MDTensor with the MDSpan's pointer and strides
+   MDTensor<N,T> MDWrite(bool on_dev = true)
+   {
+      T *ptr = mfem::Write(mfem_type::data, mfem_type::size, on_dev);
+      return MDTensor<N,T>(ptr, Sd);
+   }
+
+   /// Shortcut for mfem::Write(mfem_type::data, mfem_type::size, false)
+   /// and return an MDTensor with the MDSpan's pointer and strides
+   MDTensor<N,T> MDHostWrite()
+   {
+      T *ptr = mfem::Write(mfem_type::data, mfem_type::size, false);
+      return MDTensor<N,T>(ptr, Sd);
+   }
+
+   /// Shortcut for mfem::ReadWrite(mfem_type::data, mfem_type::size, on_dev)
+   /// and return an MDTensor with the MDSpan's pointer and strides
+   MDTensor<N,T> MDReadWrite(bool on_dev = true)
+   {
+      T *ptr = mfem::ReadWrite(mfem_type::data, mfem_type::size, on_dev);
+      return MDTensor<N,T>(ptr, Sd);
+   }
+
+   /// Shortcut for mfem::ReadWrite(mfem_type::data, mfem_type::size, false)
+   /// and return an MDTensor with the MDSpan's pointer and strides
+   MDTensor<N,T> MDHostReadWrite()
+   {
+      T *ptr = mfem::ReadWrite(mfem_type::data, mfem_type::size, false);
+      return MDTensor<N,T>(ptr, Sd);
+   }
+
+   /// The MDReshape function allows to reshape the multi-dimentional view
+   /// into a new multi-dimentional one, by the use of std::array blocks.
+   /// For example, if 'this' has three dimensions {N1, N2, N3}, it could handle
+   /// this->MDReshape<4>(ptr, N1, std::array<int,2> {2, N2/2}, N3);
+
+   // Parameter R could be omitted with c++14 standard's deduced return types
+
+   // first method with given data pointer and rest of arguments
+   template <int R, int m = 0, int M = 0, typename... Ts>
+   inline auto MDReshape(T *ptr, Ts&&... args) -> MDTensor<R,T>
+   {
+      rNd.clear();
+      reshape_ptr = ptr;
+      reshape_offset = 1, reshape_shifts[0] = reshape_shifts[1] = 0;
+      return MDReshape<R,m,M>(std::forward<Ts>(args)...);
+   }
+
+   // variadic method, where a new block of reshape is given in argument
+   template <int R, int m = 0, int M = 0, size_t P, typename... Ts>
+   inline auto MDReshape(std::array<int,P> list, Ts&&... args) -> MDTensor<R,T>
+   {
+      reshape_shifts[0] = layout.perm[m]; // store layout shift begin
+      int shifted_layout = reshape_shifts[1] + layout.perm[m];
+      for (int dim: list)
+      {
+         rNd.push_back(dim);
+         rLt[m].push_back(sub_layout_pair{shifted_layout,-1});
+         shifted_layout += 1; // default left layout
+      }
+      reshape_shifts[1] += P-1; // update end
+      return MDReshape<R,m+1,M+P>(std::forward<Ts>(args)...);
+   }
+
+   // variadic method, where a new dimension of reshape is given
+   template <int R, int m = 0, int M = 0, typename... Ts>
+   inline auto MDReshape(int dim, Ts&&... args) -> MDTensor<R,T>
+   {
+      rNd.push_back(dim);
+      const int shift =
+         reshape_shifts[0] < layout.perm[m] ? reshape_shifts[1] : 0;
+      rLt[m].push_back(sub_layout_pair{layout.perm[m] + shift,-1});
+      return MDReshape<R,m+1,M+1>(std::forward<Ts>(args)...);
+   }
+
+   // terminal case which returns the resulting MDTensor
+   template <int R, int m = 0, int M = 0>
+   inline MDTensor<R,T> MDReshape()
+   {
+      int k = 0, rLt_idx[M], rSd[M];
+      // initialize sub_layout_pair's second
+      for (sub_layout_type &sub: rLt)
+      {
+         for (sub_layout_pair &p: sub) { p.second = k++; }
+      }
+      // scan with the previous layout (N) order the reshaped layout (M)
+      for (int i = 0, j = 0; i < N; i++)
+      {
+         for (sub_layout_pair &p: rLt[layout[i]])
+         {
+            rLt_idx[j++] = p.second;
+         }
+      }
+      // apply the reshaped layout (M)
+      rSd[rLt_idx[0]] = 1;
+      for (int i = 1; i < M; i++)
+      {
+         rSd[rLt_idx[i]] = rNd[rLt_idx[i-1]] * rSd[rLt_idx[i-1]];
+      }
+      // construct the MDTensor with the given pointer and reshaped sizes
+      static_assert(R == M, "R != M");
+      return MDTensor<R,T>(reshape_ptr, rSd);
+   }
+
+private:
+   T *reshape_ptr;
+   std::vector<int> rNd; // reshape sizes
+   int reshape_offset, reshape_shifts[2];// shift begin & end
+   using sub_layout_pair = std::pair<int,int>;
+   using sub_layout_type = std::list<sub_layout_pair>;
+   std::array<sub_layout_type,N> rLt; // layout
+};
+
+// md_sequence, md_extend and make_md_sequence implementation
+namespace internal
+{
+
+template <typename T, int N, int mod, bool left> struct md_extend;
+
+template <typename T, T... args, int N>
+struct md_extend<md_sequence<T, args...>, N, 0, true>
+{
+   using type = md_sequence<T, args..., (args + N)...>;
+};
+
+template <typename T, T... args, int N>
+struct md_extend<md_sequence<T, args...>, N, 1, true>
+{
+   using type = md_sequence<T, args..., (args + N)..., 2*N>;
+};
+
+template <typename T, T... args, int N>
+struct md_extend<md_sequence<T, args...>, N, 0, false>
+{
+   using type = md_sequence<T, (args + N)..., args...>;
+};
+
+template <typename T, T... args, int N>
+struct md_extend<md_sequence<T, args...>, N, 1, false>
+{
+   using type = md_sequence<T, 2*N, (args + N)..., args...>;
+};
+
+template <typename T, int N, bool L> struct make_md_sequence
+{
+   using sequence_type = typename make_md_sequence<T,N/2,L>::type;
+   using type = typename md_extend<sequence_type, N/2, N%2, L>::type;
+};
+
+template <typename T, bool L>
+struct make_md_sequence<T,0,L> { using type = md_sequence<T>; };
+
+} // namespace internal
+
+} // namespace mfem
+
+#endif // MFEM_MDSPAN_HPP

general/mem_manager.hpp

@@ -169,6 +169,7 @@ class Memory
 protected:
    friend class MemoryManager;
    friend void MemoryPrintFlags(unsigned flags);
+   template<typename mfem_type, int N, typename L> friend class MDSpan;
 
    enum FlagMask: unsigned
    {

linalg/CMakeLists.txt

@@ -57,6 +57,7 @@ list(APPEND HDRS
   lapack.hpp
   linalg.hpp
   matrix.hpp
+  mdvector.hpp
   ode.hpp
   operator.hpp
   solvers.hpp

linalg/mdvector.hpp

@@ -0,0 +1,68 @@
+// Copyright (c) 2010-2023, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#ifndef MFEM_MDVECTOR
+#define MFEM_MDVECTOR
+
+#include "../config/config.hpp"
+
+#include "vector.hpp"
+#include "general/mdspan.hpp"
+
+namespace mfem
+{
+
+template<int N, typename Layout = MDLayoutLeft<N>>
+struct MDVector : public MDSpan<Vector, N, Layout>
+{
+   using base_t = MDSpan<Vector, N, Layout>;
+
+   /**
+    * @brief MDVector default constructor (recursion)
+    */
+   MDVector(): base_t() { }
+
+   /**
+    * @brief MDVector recursion constructor
+    * @param[in] n Dimension indice
+    * @param[in] args Rest of dimension indices
+    */
+   template <typename... Ts>
+   MDVector(int n, Ts... args): MDVector(args...) { base_t::Setup(n, args...); }
+
+   /// Move constructor not supported
+   MDVector(MDVector&&) = delete;
+
+   /// Copy constructor not supported
+   MDVector(const MDVector&) = delete;
+
+   /// Move assignment not supported
+   MDVector& operator=(MDVector&&) = delete;
+
+   /// Copy assignment not supported
+   MDVector& operator=(const MDVector&) = delete;
+
+   using Vector::Read;
+   using Vector::Write;
+   using Vector::ReadWrite;
+   using Vector::HostRead;
+   using Vector::HostWrite;
+   using Vector::HostReadWrite;
+
+   using Vector::GetData;
+   using Vector::SetData;
+
+   using Vector::operator=;
+};
+
+} // namespace mfem
+
+#endif // MFEM_MDVECTOR

linalg/vector.hpp

@@ -79,6 +79,7 @@ inline real_t rand_real()
 class Vector
 {
 protected:
+   template<typename mfem_type, int N, typename L> friend class MDSpan;
 
    Memory<real_t> data;
    int size;

tests/unit/CMakeLists.txt

@@ -18,6 +18,7 @@ include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR})
 #    for d in general linalg mesh fem enzyme; do ls -1 $d/*.cpp; done
 set(UNIT_TESTS_SRCS
   general/test_array.cpp
+  general/test_mdspan.cpp
   general/test_arrays_by_name.cpp
   general/test_error.cpp
   general/test_mem.cpp

tests/unit/general/test_mdspan.cpp

@@ -0,0 +1,406 @@
+// Copyright (c) 2010-2023, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+#ifdef _WIN32
+#define _USE_MATH_DEFINES
+#include <cmath>
+#endif
+
+#include <list>
+#include <type_traits>
+
+#include "mfem.hpp"
+#include "unit_tests.hpp"
+
+#include "general/mdspan.hpp"
+#include "general/forall.hpp"
+
+#include "fem/mdgridfunc.hpp"
+#include "general/mdarray.hpp"
+#include "linalg/mdvector.hpp"
+
+using namespace mfem;
+
+static bool is_equal(const Vector &a, const Vector &b);
+
+TEST_CASE("MDArray", "[MDSpan][MDArray]")
+{
+   SECTION("Types")
+   {
+      MDArray<int,3> mda;
+      REQUIRE(mda.Size() == 0);
+      REQUIRE(std::is_same<decltype(mda.HostRead()), int const*>());
+      REQUIRE(std::is_same<decltype(mda.MDHostRead()), MDTensor<3, int const> const>());
+   }
+
+   SECTION("SetSize")
+   {
+      constexpr int NA = 11, NB = 22, NC = 33;
+      {
+
+         const int A = 7;
+         MDArray<int,3> abc;
+         abc.SetSize(NA, NB, NC);
+         abc = 7;
+         REQUIRE(abc.Size() == NA*NB*NC);
+         REQUIRE(abc.Read());
+         REQUIRE(abc.Write());
+         REQUIRE(abc.HostRead());
+         REQUIRE(abc.HostWrite());
+         REQUIRE(abc.MDRead()(0,0,0) == A);
+         REQUIRE(abc.MDWrite()(0,0,0) == A);
+         REQUIRE(abc.MDHostRead()(0,0,0) == A);
+         REQUIRE(abc.MDHostWrite()(0,0,0) == A);
+      }
+      {
+         MDArray<int,3> abc(NA, NB, NC);
+         REQUIRE(abc.Size() == NA*NB*NC);
+      }
+      {
+         const int A[6] = {0, 1, 2, 3, 4, 7};
+         MDArray<int,3,MDLayoutLeft<3>> abc_l(1,2,3);
+         MDArray<int,3,MDLayoutRight<3>> abc_r(1,2,3);
+
+         abc_l.Assign(A);
+         REQUIRE(abc_l.MDRead()(0,0,0) == 0);
+         REQUIRE(abc_l.MDRead()(0,1,2) == 7); // = 0 + 1( 1 + 2( 2)) = 5
+
+         abc_r.Assign(A);
+         REQUIRE(abc_r.MDRead()(0,0,0) == 0);
+         REQUIRE(abc_r.MDRead()(0,1,2) == 7); // = ((0)*2 + 1) * 3 + 2 = 5
+      }
+   }
+
+   SECTION("Offset")
+   {
+      constexpr int NA = 11, NB = 22, NC = 33;
+      constexpr int na =  0, nb =  1, nc =  2;
+      MDLayout<3> layout_012({na,nb,nc});
+
+      MDArray<int,3> abc(NA, NB, NC);
+      MDArray<int,3, MDLayout<3>> abc_ini(NA, NB, NC);
+      MDArray<int,3, MDLayout<3>> abc_set(NA, NB, NC);
+      abc_set.SetLayout(layout_012);
+
+      REQUIRE(abc_set.Offset(na,nb,nc) == abc.Offset(na,nb,nc));
+      REQUIRE(abc_set.Offset(na,nb,nc) == abc_ini.Offset(na,nb,nc));
+   }
+
+   SECTION("SetLayout")
+   {
+      constexpr int NA = 18, NB = 2, NC = 36;
+      // Fortran col major: (18, 2, 36)
+      //                    ( 0, 1,  2)
+      // = 0 + 18( 1 + 2( 2)) = 90
+      MDArray<int,3> left(NA,NB,NC); // default layout is LayoutLeft
+      REQUIRE(left.Offset(0,1,2) == 90);
+
+      // C/C++ row major: (18, 2, 36)
+      //                  ( 0, 1,  2)
+      // = 32( 2 + 36( 1 + 2(0))) = 38
+      // = ((0)*2 + 1) * 36 + 2
+      MDArray<int,3> right(NA, NB, NC);
+      right.SetLayout(MDLayout<3>({2,1,0}));
+      REQUIRE(right.Offset(0,1,2) == 38);
+
+      MDArray<int,3,MDLayoutRight<3>> right4(NA, NB, NC);
+      right4.SetLayout(MDLayoutRight<3>({2,1,0}));
+      REQUIRE(right4.Offset(0,1,2) == 38);
+   }
+}
+
+TEST_CASE("MDVector", "[MDSpan][MDVector]")
+{
+   SECTION("Types")
+   {
+      MDVector<3> mdv;
+      REQUIRE(mdv.Size() == 0);
+      REQUIRE(std::is_same<decltype(mdv.HostRead()), double const*>());
+      REQUIRE(std::is_same<decltype(mdv.MDHostRead()), MDTensor<3, double const> const>());
+   }
+
+   SECTION("SetSize")
+   {
+      constexpr int NA = 11, NB = 22, NC = 33;
+      {
+         MDVector<3> abc;
+         abc.SetSize(NA, NB, NC);
+         REQUIRE(abc.Size() == NA*NB*NC);
+         abc.HostRead();
+         abc.MDHostRead();
+      }
+      {
+         MDVector<3> abc(NA, NB, NC);
+         REQUIRE(abc.Size() == NA*NB*NC);
+      }
+   }
+
+   SECTION("Offset")
+   {
+      constexpr int NA = 11, NB = 22, NC = 33;
+      constexpr int na =  0, nb =  1, nc =  2;
+      MDLayout<3> layout_012({na,nb,nc});
+
+      MDVector<3> abc(NA, NB, NC);
+      MDVector<3, MDLayout<3>> abc_ini(NA, NB, NC);
+      MDVector<3, MDLayout<3>> abc_set(NA, NB, NC);
+      abc_set.SetLayout(layout_012);
+
+      REQUIRE(abc_set.Offset(na,nb,nc) == abc.Offset(na,nb,nc));
+      REQUIRE(abc_set.Offset(na,nb,nc) == abc_ini.Offset(na,nb,nc));
+   }
+
+   SECTION("SetLayout")
+   {
+      constexpr int NA = 18, NB = 2, NC = 36, ND = 32;
+      // Fortran col major: (N1:18, 2, 36, Nd:32)
+      //                    (    0, 1,  2,     3)
+      // = 0 + 18( 1 + 2( 2 + 36( 3))) = 3978
+      MDVector<4> left(NA,NB,NC,ND); // default layout is LayoutLeft
+      REQUIRE(left.Offset(0,1,2,3) == 3978);
+
+      // C/C++ row major: (N1:18, 2, 36, Nd:32)
+      //                  (    0, 1,  2,     3)
+      // = 3 + 32( 2 + 36( 1 + 2(0))) = 1219
+      // = (((0)*2 + 1) * 36 + 2) * 32 + 3
+      MDVector<4> right(NA, NB, NC, ND);
+      right.SetLayout(MDLayout<4>({3,2,1,0}));
+      REQUIRE(right.Offset(0,1,2,3) == 1219);
+
+      MDVector<4,MDLayoutRight<4>> right4(NA, NB, NC, ND);
+      right4.SetLayout(MDLayoutRight<4>({3,2,1,0}));
+      REQUIRE(right4.Offset(0,1,2,3) == 1219);
+   }
+}
+
+TEST_CASE("MDGridFunction layouts", "[MDSpan][MDGridFunction]")
+{
+   constexpr int NE = 7, NG = 3, NA = 5;
+
+   const bool all = launch_all_non_regression_tests;
+   auto p = all ? GENERATE(1,2) : 3;
+   auto nx = all ? GENERATE(3,5) : 2;
+   auto dim = all ? GENERATE(1,2,3) : 2;
+   CAPTURE(p, nx, dim);
+
+   auto MakeCartesian = [](int dim, int nx)
+   {
+      return dim == 2 ? Mesh::MakeCartesian2D(nx, nx, Element::QUADRILATERAL):
+             dim == 3 ? Mesh::MakeCartesian3D(nx, nx, nx, Element::HEXAHEDRON):
+             Mesh::MakeCartesian1D(nx);
+   };
+   Mesh mesh = MakeCartesian(dim, nx);
+
+   H1_FECollection fec(p, dim);
+   FiniteElementSpace fes(&mesh, &fec);
+   const int ND = fes.GetNDofs();
+
+   SECTION("Types")
+   {
+      MDGridFunction<4> mdgf(NE, NG, &fes, NA);
+      REQUIRE(mdgf.Size() == (NE * NG * fes.GetVSize() * NA));
+      REQUIRE(std::is_same<decltype(mdgf.HostRead()), double const*>());
+      REQUIRE(std::is_same<decltype(mdgf.MDHostRead()), MDTensor<4, double const> const>());
+   }
+
+   SECTION("LeftOffset")
+   {
+      MDGridFunction<4> gsa(NE, NG, &fes, NA);
+      const int gsa_0123 = gsa.Offset(0, 1, 2, 3);
+      REQUIRE(gsa_0123 == 0 + 1*(NE) + 2*(NE*NG) + 3*(NE*NG*ND));
+   }
+
+   SECTION("RightOffset")
+   {
+      MDGridFunction<4, MDLayoutRight<4>> gsa(NE, NG, &fes, NA);
+      const int gsa_0123 = gsa.Offset(0,1,2,3);
+      REQUIRE(gsa_0123 == 0*(NG*ND*NA) + 1*(ND*NA) + 2*(NA) + 3);
+   }
+
+   SECTION("Set/Get ScalarGridFunction")
+   {
+      MDGridFunction<3> egda(NG, &fes, NA);
+
+      GridFunction gf, rho(&fes);
+
+      BilinearForm M_ho(&fes);
+      M_ho.AddDomainIntegrator(new MassIntegrator);
+      M_ho.Assemble();
+      M_ho.Finalize();
+
+      auto compute_mass = [](GridFunction &gf)
+      {
+         FiniteElementSpace *fes = gf.FESpace();
+         ConstantCoefficient one(1.0);
+         BilinearForm ML2(fes);
+         ML2.AddDomainIntegrator(new MassIntegrator(one));
+         ML2.Assemble();
+         GridFunction ones(fes);
+         ones = 1.0;
+         return ML2.InnerProduct(gf, ones);
+      };
+
+      FunctionCoefficient rho_cft([](const Vector &x)
+      {
+         return x(1) + 0.25*cos(2*M_PI*x.Norml2());
+      });
+      rho.ProjectCoefficient(rho_cft);
+      const double rho_mass = compute_mass(rho);
+
+      const std::list<MDLayout<3>> layouts =
+      { {0,1,2}, {0,2,1}, {1,0,2}, {1,2,0}, {2,1,0}, {2,0,1} };
+
+      for (auto &layout: layouts)
+      {
+         egda = M_PI;
+         egda.SetLayout(layout);
+         for (int na = 0; na < NA; na++)
+         {
+            for (int ng = 0; ng < NG; ng++)
+            {
+               egda.GetScalarGridFunction(ng, gf, na);
+               REQUIRE(gf.Size() == fes.GetVSize());
+               REQUIRE(gf[0] == M_PI);
+               gf = rho;
+               egda.SetScalarGridFunction(ng, gf, na);
+               gf = 0.0;
+               egda.GetScalarGridFunction(ng, gf, na);
+               REQUIRE(is_equal((Vector&)gf, (Vector&)rho));
+               REQUIRE(compute_mass(gf) == MFEM_Approx(rho_mass));
+            }
+         }
+      }
+   }
+}
+
+TEST_CASE("MDGridFunction reshapes", "[MDSpan][MDReshapes]")
+{
+   SECTION("MDReshapes")
+   {
+      constexpr int p = 2;
+      constexpr int dim = 3;
+      constexpr int nx = 5, ny = 3, nz = 2;
+      Mesh mesh = Mesh::MakeCartesian3D(nx, ny, nz, Element::HEXAHEDRON);
+
+      H1_FECollection fec_mesh(p, dim);
+      FiniteElementSpace fes_mesh(&mesh, &fec_mesh, dim);
+      mesh.SetNodalFESpace(&fes_mesh);
+
+      L2_FECollection fec(p, dim);
+      FiniteElementSpace fes(&mesh, &fec);
+
+      const std::list<MDLayout<3>> layouts =
+      { {0,1,2}, {0,2,1}, {1,0,2}, {1,2,0}, {2,1,0}, {2,0,1} };
+
+      for (auto &layout: layouts)
+      {
+         constexpr int numGroups = 4, numAngles = 7;
+
+         MDGridFunction<3> psi(&fes, numGroups, numAngles);
+         psi.SetLayout(layout);
+
+         const GridFunction *nodes = mesh.GetNodes();
+         const FiniteElementSpace *mfes = mesh.GetNodalFESpace();
+         const int ng = numGroups, na = numAngles, ne = mfes->GetNE();
+         const ElementDofOrdering e_ordering = ElementDofOrdering::LEXICOGRAPHIC;
+         const Operator *R = mfes->GetElementRestriction(e_ordering);
+         REQUIRE(R);
+         const FiniteElement *mfe = mfes->GetFE(0);
+         const int nd = mfe->GetDof(), vdim = mfes->GetVDim();
+         Vector nodes_e(vdim*nd*ne); nodes_e.UseDevice(true);
+         constexpr int D1D = p + 1;
+         REQUIRE(fes.GetVSize() == D1D*D1D*D1D*ne);
+         nodes_e.Read();
+         REQUIRE(nodes);
+         R->Mult(*nodes, nodes_e);
+         const auto X = Reshape(nodes_e.Read(), D1D, D1D, D1D, vdim, ne);
+         auto dY = psi.MDWrite();
+
+         MDGridFunction<3> rY1(&fes, numGroups, numAngles);
+         rY1.SetLayout(MDLayout<3>(layout));
+         REQUIRE(rY1.Size() == psi.Size());
+         auto drY1 = rY1.MDWrite();
+
+         MDGridFunction<3> rY2(&fes, numGroups, numAngles);
+         rY2.SetLayout(MDLayout<3>(layout));
+         REQUIRE(ng%2 == 0);
+         auto drY2 = rY2.MDReshape<4>(rY2.Write(),
+                                      D1D*D1D*D1D*ne,
+                                      std::array<int,2> {2, ng/2},
+                                      na);
+
+         MDGridFunction<3> rY3(&fes, numGroups, numAngles);
+         rY3.SetLayout(MDLayout<3>(layout));
+         auto drY3 = rY3.MDReshape<6>(rY3.Write(),
+                                      std::array<int,4> {D1D, D1D, D1D, ne},
+                                      ng, na);
+
+         MDGridFunction<3> rY4(&fes, numGroups, numAngles);
+         rY4.SetLayout(MDLayout<3>(layout));
+         auto drY4 = rY4.MDReshape<7>(rY4.Write(),
+                                      std::array<int,4> {D1D, D1D, D1D, ne},
+                                      std::array<int,2> {1, ng},
+                                      na);
+
+         MDGridFunction<3> rY5(&fes, numGroups, numAngles);
+         rY5.SetLayout(MDLayout<3>(layout));
+         auto drY5 = rY5.MDReshape<7>(rY5.Write(),
+                                      std::array<int,4> {D1D, D1D, D1D, ne},
+                                      std::array<int,2> {2, ng/2},
+                                      na);
+
+         const double exp_m08 = exp(-0.8);
+
+         mfem::forall_3D(ne*ng*na, D1D,D1D,D1D, [=] MFEM_HOST_DEVICE(int ega)
+         {
+            const int e = ega/(ng*na), ga = ega%(ng*na), g = ga/na, a = ga%na;
+            MFEM_FOREACH_THREAD(dz,z,D1D)
+            {
+               MFEM_FOREACH_THREAD(dy,y,D1D)
+               {
+                  MFEM_FOREACH_THREAD(dx,x,D1D)
+                  {
+                     const int xyze = dx + D1D*(dy + D1D*(dz + D1D*(e)));
+
+                     const double p0 = X(dx,dy,dz,0,e), p1 = X(dx,dy,dz,1,e);
+                     const double value = 1.0 - exp_m08*cos(M_PI*p0)*cos(M_PI*p1);
+
+                     dY(xyze,g,a) = value;
+                     drY1(xyze,g,a) = value;
+                     drY2(xyze,g%2,g/2,a) = value;
+                     drY3(dx,dy,dz,e, g, a) = value;
+                     drY4(dx,dy,dz,e, 0,g, a) = value;
+                     drY5(dx,dy,dz,e, g%2,g/2, a) = value;
+                  }
+               }
+            }
+         });
+         psi.MDHostRead(); rY1.HostRead();
+         REQUIRE(is_equal((Vector&)rY1, (Vector&)psi));
+         REQUIRE(is_equal((Vector&)rY2, (Vector&)psi));
+         REQUIRE(is_equal((Vector&)rY3, (Vector&)psi));
+         REQUIRE(is_equal((Vector&)rY4, (Vector&)psi));
+         REQUIRE(is_equal((Vector&)rY5, (Vector&)psi));
+      }
+   }
+}
+
+static bool is_equal(const Vector &a, const Vector &b)
+{
+   REQUIRE(a.Size() == b.Size());
+   for (int i = 0; i < a.Size(); i++)
+   {
+      const double va = a.GetData()[i], vb = b.GetData()[i];
+      REQUIRE(va == MFEM_Approx(vb));
+   };
+   return true;
+};
+