cs_mdspan.h

Go to the documentation of this file.

#ifndef __CS_MDSPAN_H__
#define __CS_MDSPAN_H__
 
/*============================================================================
 * mdspan class and handling utilities.
 *============================================================================*/
 
/*
  This file is part of code_saturne, a general-purpose CFD tool.
 
  Copyright (C) 1998-2025 EDF S.A.
 
  This program is free software; you can redistribute it and/or modify it under
  the terms of the GNU General Public License as published by the Free Software
  Foundation; either version 2 of the License, or (at your option) any later
  version.
 
  This program 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 General Public License for more
  details.
 
  You should have received a copy of the GNU General Public License along with
  this program; if not, write to the Free Software Foundation, Inc., 51 Franklin
  Street, Fifth Floor, Boston, MA 02110-1301, USA.
*/
 
/*----------------------------------------------------------------------------*/
 
/*----------------------------------------------------------------------------
 * Standard C and C++ library headers
 *----------------------------------------------------------------------------*/
 
#include <string.h>
 
/*----------------------------------------------------------------------------
 *  Local headers
 *----------------------------------------------------------------------------*/
 
#include "base/cs_defs.h"
 
#include "base/cs_dispatch.h"
 
/*----------------------------------------------------------------------------*/
 
#if defined(__cplusplus)
 
namespace cs {
 
enum class
layout {
  right,   
  left,    
  unknown  
};
 
/*----------------------------------------------------------------------------*/
/*----------------------------------------------------------------------------*/
 
template<class T, int N, layout L = layout::right>
class mdspan {
 
/*=============================================================================
 * Public methods
 *============================================================================*/
 
public:
 
  /* Make array class friend */
  template<class _T_, int _N_, layout _L_> friend class array;
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  mdspan():
    _extent{0},
    _offset{0},
    _size(0),
    _data(nullptr)
  {}
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  mdspan
  (
    T       *data,   
    Args...  indices 
  )
  {
    check_operator_args_(indices...);
    set_size_(indices...);
    _data = data;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  mdspan
  (
    T                     *data, 
    const cs_lnum_t(&dims)[N]    
  )
  {
    set_size_(dims);
    _data = data;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  mdspan
  (
    const mdspan& other
  )
  {
    set_size_(other._extent);
    _data = other._data;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST
  mdspan
  (
    mdspan&& other 
  )
  : mdspan()
  {
    swap(*this, other);
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  ~mdspan()
  {
    _data = nullptr;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST
  friend void
  swap
  (
    mdspan& first, 
    mdspan& second 
  )
  {
    using std::swap;
 
    swap(first._extent, second._extent);
    swap(first._offset, second._offset);
    swap(first._size, second._size);
    swap(first._data, second._data);
  }
 
  /*===========================================================================
   * Operators
   *==========================================================================*/
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST
  mdspan& operator=(const mdspan& other)
  {
    for (int i = 0; i < N; i++) {
      _extent[i] = other._extent[i];
      _offset[i] = other._offset[i];
    }
    _size = other._size;
    _data = other._data;
 
    return *this;
  }
 
  CS_F_HOST
  mdspan& operator=(mdspan&& other)
  {
 
    swap(*this, other);
 
    return *this;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  inline
  T& operator[]
  (
    cs_lnum_t i 
  )
  {
    return _data[i];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  inline
  T& operator[]
  (
    cs_lnum_t i 
  ) const
  {
    return _data[i];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename Id1>
  CS_F_HOST_DEVICE
  inline
  std::enable_if_t<cs::always_true<Id1>::value && N==1, T&>
  operator()
  (
    Id1 i 
  )
  {
    check_operator_args_(i);
    return _data[i];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename Id1>
  CS_F_HOST_DEVICE
  inline
  std::enable_if_t<cs::always_true<Id1>::value && N==1, T&>
  operator()
  (
    Id1 i 
  ) const
  {
    check_operator_args_(i);
    return _data[i];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename Id1, typename Id2>
  CS_F_HOST_DEVICE
  inline
  std::enable_if_t<cs::always_true<Id1, Id2>::value && N==2, T&>
  operator()
  (
    Id1 i, 
    Id2 j  
  )
  {
    check_operator_args_(i,j);
 
    if (L == layout::right)
      return _data[i*_offset[0] + j];
    else if (L == layout::left)
      return _data[i + j*_offset[1]];
    else
      return _data[i*_offset[0] + j*_offset[1]];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename Id1, typename Id2>
  CS_F_HOST_DEVICE
  inline
  std::enable_if_t<cs::always_true<Id1, Id2>::value && N==2, T&>
  operator()
  (
    Id1 i, 
    Id2 j  
  ) const
  {
    check_operator_args_(i,j);
 
    if (L == layout::right)
      return _data[i*_offset[0] + j];
    else if (L == layout::left)
      return _data[i + j*_offset[1]];
    else
      return _data[i*_offset[0] + j*_offset[1]];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  inline
  std::enable_if_t<cs::always_true<Args...>::value && (N!=2) && (N!=1), T&>
  operator()
  (
    Args... indices 
  )
  {
    check_operator_args_(indices...);
 
    return _data[data_offset_(indices...)];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  inline
  std::enable_if_t<cs::always_true<Args...>::value && (N!=2) && (N!=1), T&>
  operator()
  (
    Args... indices 
  ) const
  {
    check_operator_args_(indices...);
 
    return _data[data_offset_(indices...)];
  }
 
  /*===========================================================================
   * Getters
   *==========================================================================*/
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  inline
  cs_lnum_t
  size()
  {
    return _size;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  inline
  cs_lnum_t
  size() const
  {
    return _size;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  inline
  cs_lnum_t
  extent
  (
    int i
  )
  {
    return _extent[i];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  inline
  cs_lnum_t
  extent
  (
    int i
  ) const
  {
    return _extent[i];
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  T*
  data()
  {
    return _data;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename U>
  CS_F_HOST_DEVICE
  U*
  data()
  {
    return reinterpret_cast<U*>(_data);
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  mdspan<T, N-sizeof...(Args), L>
  sub_view
  (
    Args... indices 
  )
  {
    check_sub_function_args_(indices...);
 
    constexpr int n_idx = sizeof...(Args);
 
    cs_lnum_t dims[N - n_idx];
 
    if (L == layout::right) {
      for (int i = 0; i < N - n_idx; i++)
        dims[i] = _extent[i+n_idx];
    }
    else if (L == layout::left) {
      for (int i = 0; i < N - n_idx; i++)
        dims[i] = _extent[i];
    }
 
    return mdspan<T,N-n_idx,L>(_data + contiguous_data_offset_(indices...), dims);
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  T*
  sub_array
  (
    Args... indices 
  )
  {
    check_sub_function_args_(indices...);
 
    return _data + contiguous_data_offset_(indices...);
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST
  void set_to_val
  (
    T               val,        
    const cs_lnum_t n_vals = -1 
  )
  {
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    cs_dispatch_context ctx;
 
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = val;
    });
 
    ctx.wait();
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST
  void set_to_val
  (
    cs_dispatch_context &ctx,        
    T                    val,        
    const cs_lnum_t      n_vals = -1 
  )
  {
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    /* No wait here since context is passed as argument */
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = val;
    });
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST
  void set_to_val_on_subset
  (
    T               val,      
    const cs_lnum_t n_elts,   
    const cs_lnum_t elt_ids[] 
  )
  {
    assert(n_elts <= _size && n_elts >= 0);
 
    if (n_elts < 1)
      return;
 
    cs_dispatch_context ctx;
 
    if (elt_ids == nullptr)
      set_to_val(ctx, val, n_elts);
    else {
      // Explicit pointer, avoid passing internal member of class to the functor
      T* data_ptr = _data;
 
      ctx.parallel_for(n_elts, CS_LAMBDA (cs_lnum_t e_id) {
        data_ptr[elt_ids[e_id]] = val;
      });
    }
 
    ctx.wait();
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST
  void set_to_val_on_subset
  (
    cs_dispatch_context &ctx,      
    T                    val,      
    const cs_lnum_t      n_elts,   
    const cs_lnum_t      elt_ids[] 
  )
  {
    assert(n_elts <= _size && n_elts >= 0);
 
    if (n_elts < 1)
      return;
 
    /* No wait here since context is passed as argument */
    if (elt_ids == nullptr)
      set_to_val(ctx, val, n_elts);
    else {
      // Explicit pointer, avoid passing internal member of class to the functor
      T* data_ptr = _data;
      ctx.parallel_for(n_elts, CS_LAMBDA (cs_lnum_t e_id) {
        data_ptr[elt_ids[e_id]] = val;
      });
    }
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  zero()
  {
    cs_dispatch_context ctx;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    ctx.parallel_for(_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = static_cast<T>(0);
    });
 
    ctx.wait();
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  zero
  (
    cs_dispatch_context &ctx 
  )
  {
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    ctx.parallel_for(_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = static_cast<T>(0);
    });
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  copy_data
  (
    T               *data,       
    const cs_lnum_t  n_vals = -1 
  )
  {
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    cs_dispatch_context ctx;
 
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
        data_ptr[e_id] = data[e_id];
    });
 
    ctx.wait();
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  copy_data
  (
    const T         *data,       
    const cs_lnum_t  n_vals = -1 
  )
  {
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    cs_dispatch_context ctx;
 
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
        data_ptr[e_id] = data[e_id];
    });
 
    ctx.wait();
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  copy_data
  (
    mdspan&         other,      
    const cs_lnum_t n_vals = -1 
  )
  {
    assert(other._size == _size);
 
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
    T* o_data_ptr = other._data;
 
    cs_dispatch_context ctx;
 
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = o_data_ptr[e_id];
    });
 
    ctx.wait();
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  copy_data
  (
    cs_dispatch_context &ctx,        
    T                   *data,       
    const cs_lnum_t      n_vals = -1 
  )
  {
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = data[e_id];
    });
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  copy_data
  (
    cs_dispatch_context &ctx,        
    const T             *data,       
    const cs_lnum_t      n_vals = -1 
  )
  {
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
 
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = data[e_id];
    });
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  copy_data
  (
    cs_dispatch_context &ctx,        
    mdspan              &other,      
    const cs_lnum_t      n_vals = -1 
  )
  {
    assert(other.size() == _size);
 
    assert(n_vals <= _size);
 
    const cs_lnum_t loop_size = (n_vals == -1) ? _size : n_vals;
 
    // Explicit pointer, avoid passing internal member of class to the functor
    T* data_ptr = _data;
    T* o_data_ptr = other._data;
 
    ctx.parallel_for(loop_size, CS_LAMBDA (cs_lnum_t e_id) {
      data_ptr[e_id] = o_data_ptr[e_id];
    });
  }
 
private:
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  static inline
  void
  check_operator_args_
  (
    Args... 
  )
  {
    static_assert(sizeof...(Args) == N, "Wrong number of arguments");
    static_assert(sizeof...(Args) != 0, "No input arguments provided...");
    static_assert(cs::are_integral<Args...>::value, "Non integral input arguments.");
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  static inline
  void
  check_sub_function_args_
  (
    Args... 
  )
  {
    static_assert(sizeof...(Args) < N, "Too many input arguments.");
    static_assert(sizeof...(Args) != 0, "No input arguments provided...");
    static_assert(cs::are_integral<Args...>::value, "Non integral indices provided");
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  inline
  cs_lnum_t
  data_offset_
  (
    Args... indices 
  ) const
  {
    static_assert(sizeof...(Args) <= N && sizeof...(Args) > 0,
                  "Number of indices is out of bounds");
 
    constexpr int n_idx = sizeof...(Args);
 
    cs_lnum_t _indices[n_idx] = {indices...};
 
    cs_lnum_t retval = 0;
    for (int i = 0; i < n_idx; i++)
      retval +=_indices[i] * _offset[i];
 
    return retval;
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  inline
  cs_lnum_t
  contiguous_data_offset_
  (
    Args... indices 
  ) const
  {
    static_assert(sizeof...(Args) <= N && sizeof...(Args) > 0,
                  "Number of indices is out of bounds");
 
    constexpr int n_idx = sizeof...(Args);
 
    cs_lnum_t _indices[n_idx] = {indices...};
 
    cs_lnum_t retval = 0;
    if (L == layout::right) {
      for (int i = 0; i < n_idx; i++)
        retval +=_indices[i] * _offset[i];
    }
    else if (L == layout::left) {
      for (int i = 0; i < n_idx; i++)
        retval +=_indices[i] * _offset[N-1-i];
    }
 
    return retval;
  }
 
  /*===========================================================================
   * Private methods
   *==========================================================================*/
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  CS_F_HOST_DEVICE
  void
  set_size_
  (
    const cs_lnum_t(&dims)[N] 
  )
  {
    _size = (N > 0) ? 1 : 0;
    for (int i = 0; i < N; i++) {
      _extent[i] = dims[i];
      _size *= dims[i];
      _offset[i] = 1;
    }
 
    /* Compute offset values for getters */
 
    if (L == layout::right) {
      /* Version for Layout right */
      for (int i = 0; i < N-1; i++) {
        for (int j = i + 1; j < N; j++)
          _offset[i] *= dims[j];
      }
    }
    else if (L == layout::left) {
      for (int i = N-1; i >= 1; i--) {
        for (int j = i - 1; j >= 0; j--)
          _offset[i] *= dims[j];
      }
    }
  }
 
  /*--------------------------------------------------------------------------*/
  /*--------------------------------------------------------------------------*/
 
  template<typename... Args>
  CS_F_HOST_DEVICE
  void
  set_size_
  (
    Args... dims 
  )
  {
    check_operator_args_(dims...);
 
    cs_lnum_t loc_dims[N] = {dims...};
 
    set_size_(loc_dims);
  }
 
  /*===========================================================================
   * Private members
   *==========================================================================*/
 
  cs_lnum_t  _extent[N];
  cs_lnum_t  _offset[N];
  cs_lnum_t  _size;
  T*         _data;
};
 
} /* namespace cs */
 
template<class T, cs::layout L = cs::layout::right>
using cs_span = cs::mdspan<T, 1, L>;
 
template<class T, cs::layout L = cs::layout::right>
using cs_span_2d = cs::mdspan<T, 2, L>;
 
template<class T, cs::layout L = cs::layout::right>
using cs_span_3d = cs::mdspan<T, 3, L>;
 
template<class T, cs::layout L = cs::layout::right>
using cs_span_4d = cs::mdspan<T, 4, L>;
 
template<class T, int N>
using cs_mdspan_r = cs::mdspan<T, N, cs::layout::right>;
 
template<class T, int N>
using cs_mdspan_l = cs::mdspan<T, N, cs::layout::left>;
 
#endif /* __cplusplus */
 
/*--------------------------------------------------------------------------*/
 
#endif /* __CS_MDSPAN_H__ */