src/gromacs/mdspan/mdspan.h

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2018,2019, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /*
  36  * This file is a modified version of original work of Sandia Corporation.
  37  * In the spirit of the original code, this particular file can be distributed
  38  * on the terms of Sandia Corporation.
  39  */
  40 /*
  41  *                          Kokkos v. 2.0
  42  *               Copyright (2014) Sandia Corporation
  43  *
  44  * Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
  45  * the U.S. Government retains certain rights in this software.
  46  *
  47  * Kokkos is licensed under 3-clause BSD terms of use:
  48  *
  49  * Redistribution and use in source and binary forms, with or without
  50  * modification, are permitted provided that the following conditions are
  51  * met:
  52  *
  53  * 1. Redistributions of source code must retain the above copyright
  54  * notice, this list of conditions and the following disclaimer.
  55  *
  56  * 2. Redistributions in binary form must reproduce the above copyright
  57  * notice, this list of conditions and the following disclaimer in the
  58  * documentation and/or other materials provided with the distribution.
  59  *
  60  * 3. Neither the name of the Corporation nor the names of the
  61  * contributors may be used to endorse or promote products derived from
  62  * this software without specific prior written permission.
  63  *
  64  * THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
  65  * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  66  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  67  * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
  68  * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
  69  * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  70  * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
  71  * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  72  * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
  73  * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  74  * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  75  *
  76  * Questions? Contact Christian R. Trott (crtrott@sandia.gov)
  77  */
  78 /*! \libinternal \file
  79  * \brief Declares gmx::mdspan
  80  *
  81  * \author Christian Trott <crtrott@sandia.gov>
  82  * \author Ronan Keryell <ronan.keryell@xilinx.com>
  83  * \author Carter Edwards <hedwards@nvidia.com>
  84  * \author David Hollman <dshollm@sandia.gov>
  85  * \author Christian Blau <cblau@gwdg.de>
  86  * \inlibraryapi
  87  * \ingroup mdspan
  88  */
  89 #ifndef MDSPAN_MDSPAN_H
  90 #define MDSPAN_MDSPAN_H
  91
  92 #include <array>
  93 #include <type_traits>
  94
  95 #include "accessor_policy.h"
  96 #include "extents.h"
  97 #include "layouts.h"
  98
  99 namespace gmx
 100 {
 101
 102 /*! \libinternal \brief Multidimensional array indexing and memory access with flexible mapping and access model.
 103  *
 104  * \tparam ElementType Type of elemnt to be viewed
 105  * \tparam Extents The dimensions of the multidimenisonal array to view.
 106  * \tparam LayoutPolicy Describes is the memory layout of the multidimensional array; right by default.
 107  * \tparam AccessorPolicy Describes memory access model.
 108  */
 109 template<class ElementType, class Extents, class LayoutPolicy = layout_right, class AccessorPolicy = accessor_basic<ElementType> >
 110 class basic_mdspan
 111 {
 112     public:
 113         //! Expose type used to define the extents of the data.
 114         using extents_type     = Extents;
 115         //! Expose type used to define the layout of the data.
 116         using layout_type      = LayoutPolicy;
 117         //! Expose type used to define the memory access model of the data.
 118         using accessor_type    = AccessorPolicy;
 119         //! Expose type used to map multidimensional indices to one-dimensioal indices.
 120         using mapping_type     = typename layout_type::template mapping<extents_type>;
 121         //! Exposes the type of stored element.
 122         using element_type     = typename accessor_type::element_type;
 123         //! Expose the underlying type of the stored elements.
 124         using value_type       = std::remove_cv_t<element_type>;
 125         //! Expose the type used for indexing.
 126         using index_type       = ptrdiff_t;
 127         //! Expose type for index differences.
 128         using difference_type  = ptrdiff_t;
 129         //! Expose underlying pointer to data type.
 130         using pointer          = typename accessor_type::pointer;
 131         //! Expose reference to data type.
 132         using reference        = typename accessor_type::reference;
 133
 134         //! Trivial constructor
 135         constexpr basic_mdspan() noexcept : acc_(), map_(), ptr_() {}
 136         //! Move constructor
 137         constexpr basic_mdspan(basic_mdspan &&other) noexcept = default;
 138         //! copy constructor
 139         constexpr basic_mdspan(const basic_mdspan &other) noexcept = default;
 140         //! Copy assignment
 141         basic_mdspan &operator=(const basic_mdspan &other) noexcept = default;
 142         //! Move assignment
 143         basic_mdspan &operator=(basic_mdspan &&other) noexcept = default;
 144
 145         //! Copy constructor
 146         template<class OtherElementType,
 147                  class OtherExtents,
 148                  class OtherLayoutPolicy,
 149                  class OtherAccessor>
 150         constexpr basic_mdspan(
 151             const basic_mdspan<OtherElementType,
 152                                OtherExtents,
 153                                OtherLayoutPolicy,
 154                                OtherAccessor> &rhs ) noexcept
 155             : acc_( rhs.acc_ ),
 156               map_( rhs.map_ ),
 157               ptr_( rhs.ptr_ )
 158         {}
 159         //! Copy assignment constructor
 160         template<class OtherElementType,
 161                  class OtherExtents,
 162                  class OtherLayoutPolicy,
 163                  class OtherAccessor>
 164         basic_mdspan &operator= (
 165             const basic_mdspan<OtherElementType,
 166                                OtherExtents,
 167                                OtherLayoutPolicy,
 168                                OtherAccessor> &rhs ) noexcept
 169         { acc_ = rhs.acc_; map_ = rhs.map_; ptr_ = rhs.ptr_; return *this; }
 170
 171         /*!\brief Construct mdspan by setting the dynamic extents and pointer to data.
 172          * \param[in] ptr Pointer to data to be accessed by this span
 173          * \param[in] DynamicExtents
 174          * \tparam IndexType index type to describe dynamic extents
 175          */
 176         template<class ... IndexType >
 177         explicit constexpr basic_mdspan
 178             ( pointer ptr, IndexType ... DynamicExtents ) noexcept
 179             : acc_(accessor_type()), map_( extents_type(DynamicExtents ...) ), ptr_(ptr) {}
 180         /*! \brief Construct from array describing dynamic extents.
 181          * \param[in] ptr Pointer to data to be accessed by this span
 182          * \param[in] dynamic_extents Array the size of dynamic extents.
 183          */
 184         constexpr basic_mdspan( pointer ptr, const std::array<ptrdiff_t, extents_type::rank_dynamic()> &dynamic_extents)
 185             : acc_(accessor_type()), map_( extents_type(dynamic_extents)), ptr_(ptr) {}
 186         /*! \brief Construct from pointer and mapping.
 187          * \param[in] ptr Pointer to data to be accessed by this span
 188          * \param[in] m Mapping from multidimenisonal indices to one-dimensional offset.
 189          */
 190         constexpr basic_mdspan( pointer ptr, const mapping_type &m ) noexcept
 191             : acc_(accessor_type()), map_( m ), ptr_(ptr) {}
 192         /*! \brief Construct with pointer, mapping and accessor.
 193          * \param[in] ptr Pointer to data to be accessed by this span
 194          * \param[in] m Mapping from multidimenisonal indices to one-dimensional offset.
 195          * \param[in] a Accessor implementing memory access model.
 196          */
 197         constexpr basic_mdspan( pointer ptr, const mapping_type &m, const accessor_type &a ) noexcept
 198             : acc_(a), map_( m ), ptr_(ptr) {}
 199         /*! \brief Construct mdspan from multidimensional arrays implemented with mdspan
 200          *
 201          * Requires the container to have a view_type describing the mdspan, which is
 202          * accessible through an asView() call
 203          *
 204          *  This allows functions to declare mdspans as arguments, but take e.g. multidimensional
 205          *  arrays implicitly during the function call
 206          * \tparam U container type
 207          * \param[in] other mdspan-implementing container
 208          */
 209         template<typename U,
 210                  typename = std::enable_if_t<
 211                          std::is_same<typename std::remove_reference_t<U>::view_type::element_type,
 212                                       ElementType>::value> >
 213         constexpr basic_mdspan(U &&other) : basic_mdspan(other.asView()) {}
 214         /*! \brief Construct mdspan of const Elements from multidimensional arrays implemented with mdspan
 215          *
 216          * Requires the container to have a const_view_type describing the mdspan, which is
 217          * accessible through an asConstView() call
 218          *
 219          *  This allows functions to declare mdspans as arguments, but take e.g. multidimensional
 220          *  arrays implicitly during the function call
 221          * \tparam U container type
 222          * \param[in] other mdspan-implementing container
 223          */
 224         template<typename U,
 225                  typename = std::enable_if_t<
 226                          std::is_same<typename std::remove_reference_t<U>::const_view_type::element_type,
 227                                       ElementType>::value> >
 228         constexpr basic_mdspan(const U &other) : basic_mdspan(other.asConstView()) {}
 229         /*! \brief Brace operator to access multidimensional array element.
 230          * \param[in] indices The multidimensional indices of the object.
 231          * Requires rank() == sizeof...(IndexType). Slicing is implemented via sub_span.
 232          * \returns reference to element at indices.
 233          */
 234         template<class... IndexType >
 235         constexpr std::enable_if_t<sizeof ... (IndexType) == extents_type::rank(), reference>
 236         operator()( IndexType... indices) const noexcept
 237         { return acc_.access( ptr_, map_( indices ... ) ); }
 238         /*! \brief Canonical bracket operator for one-dimensional arrays.
 239          * Allows mdspan to act like array in one-dimension.
 240          * Enabled only when rank==1.
 241          * \param[in] i one-dimensional index
 242          * \returns reference to element stored at position i
 243          */
 244         template<class IndexType>
 245         constexpr std::enable_if_t<std::is_integral<IndexType>::value &&
 246                                    extents_type::rank() == 1, reference>
 247         operator[]( const IndexType &i ) const noexcept
 248         { return acc_.access( ptr_, map_(i) ); }
 249         /*! \brief Bracket operator for multi-dimensional arrays.
 250          *
 251          * \note Prefer operator() for better compile-time and run-time performance
 252          *
 253          * Slices two- and higher-dimensional arrays along a given slice by
 254          * returning a new basic_mdspan that drops the first extent and indexes
 255          * the remaining extents
 256          *
 257          * \note Currently only implemented for layout_right
 258          * \note For layout_right this implementation has significant
 259          *       performance benefits over implementing a more general slicing
 260          *       operator with a strided layout
 261          * \note Enabled only when rank() > 1
 262          *
 263          * \tparam IndexType integral tyoe for the index that enables indexing
 264          *                   with, e.g., int or size_t
 265          * \param[in] index  one-dimensional index of the slice to be indexed
 266          *
 267          * \returns basic_mdspan that is sliced at the given index
 268          */
 269         template <class IndexType,
 270                   typename sliced_mdspan_type =
 271                       basic_mdspan<element_type,
 272                                    decltype(extents_type().sliced_extents()),
 273                                    LayoutPolicy,
 274                                    AccessorPolicy> >
 275         constexpr std::enable_if_t<std::is_integral<IndexType>::value &&
 276                                    (extents_type::rank() > 1) &&
 277                                    std::is_same<LayoutPolicy, layout_right>::value,
 278                                    sliced_mdspan_type> operator[](const IndexType index) const noexcept
 279         {
 280             return sliced_mdspan_type(ptr_ + index * stride(0),
 281                                       extents().sliced_extents());
 282         }
 283         //! Report the rank.
 284         static constexpr int rank() noexcept
 285         { return extents_type::rank(); }
 286         //! Report the dynamic rank.
 287         static constexpr int rank_dynamic() noexcept
 288         { return extents_type::rank_dynamic(); }
 289         /*! \brief Return the static extent.
 290          * \param[in] k dimension to query for static extent
 291          * \returns static extent along specified dimension
 292          */
 293         constexpr index_type static_extent( size_t k ) const noexcept
 294         { return map_.extents().static_extent( k ); }
 295
 296         /*! \brief Return the extent.
 297          * \param[in] k dimension to query for extent
 298          * \returns extent along specified dimension
 299          */
 300         constexpr index_type extent( int k ) const noexcept
 301         { return map_.extents().extent( k ); }
 302
 303         //! Return all extents
 304         constexpr const extents_type &extents() const noexcept
 305         { return map_.extents(); }
 306         //! Report if mappings for this basic_span is always unique.
 307         static constexpr bool is_always_unique()     noexcept { return mapping_type::is_always_unique(); }
 308         //! Report if mapping for this basic_span is always strided
 309         static constexpr bool is_always_strided()    noexcept { return mapping_type::is_always_strided(); }
 310         //! Report if mapping for this basic_span is always is_contiguous
 311         static constexpr bool is_always_contiguous() noexcept { return mapping_type::is_always_contiguous(); }
 312         //! Report if the currently applied map is unique
 313         constexpr bool is_unique() const noexcept  { return map_.is_unique(); }
 314         //! Report if the currently applied map is strided
 315         constexpr bool is_strided() const noexcept { return map_.is_strided(); }
 316         //! Report if the currently applied map is contiguous
 317         constexpr bool is_contiguous() const noexcept {return map_.is_contiguous(); }
 318         //! Report stride along a specific rank.
 319         constexpr index_type stride( size_t r ) const noexcept
 320         { return map_.stride(r); }
 321         //! Return the currently applied mapping.
 322         constexpr mapping_type mapping() const noexcept { return map_; }
 323         //! Return the memory access model.
 324         constexpr accessor_type accessor() const noexcept { return acc_; }
 325         //! Return pointer to underlying data
 326         constexpr pointer data() const noexcept { return ptr_; }
 327     private:
 328         //! The memory access model
 329         accessor_type acc_;
 330         //! The transformation from multidimenisonal index to memory offset.
 331         mapping_type  map_;
 332         //! Memory location handle
 333         pointer       ptr_;
 334 };
 335
 336 //! basic_mdspan with wrapped indices, basic_accessor policiy and right-aligned  memory layout.
 337 template<class T, ptrdiff_t... Indices>
 338 using mdspan = basic_mdspan<T, extents<Indices...>, layout_right, accessor_basic<T> >;
 339
 340 }      // namespace gmx
 341
 342 #endif /* end of include guard: MDSPAN_MDSPAN_H */