src/gromacs/mdspan/extents.h

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2018,2019,2021, 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::extents for 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  * \ingroup mdspan
  87  */
  88 #ifndef MDSPAN_EXTENTS_H
  89 #define MDSPAN_EXTENTS_H
  90
  91 #include <cstddef>
  92
  93 #include <array>
  94
  95 namespace gmx
  96 {
  97
  98 /*! \brief Define constant that signals dynamic extent.
  99  */
 100 enum : std::ptrdiff_t
 101 {
 102     dynamic_extent = -1
 103 };
 104
 105 template<std::ptrdiff_t... StaticExtents>
 106 class extents;
 107
 108 template<std::ptrdiff_t... LHS, std::ptrdiff_t... RHS>
 109 constexpr bool operator==(const extents<LHS...>& lhs, const extents<RHS...>& rhs) noexcept;
 110
 111 template<std::ptrdiff_t... LHS, std::ptrdiff_t... RHS>
 112 constexpr bool operator!=(const extents<LHS...>& lhs, const extents<RHS...>& rhs) noexcept;
 113
 114 namespace detail
 115 {
 116
 117 template<int R, std::ptrdiff_t... StaticExtents>
 118 struct extents_analyse;
 119
 120 /*! \libinternal \brief Enable querying extent of specific rank by splitting
 121  * a static extents off the variadic template arguments.
 122  *
 123  */
 124 template<int R, std::ptrdiff_t E0, std::ptrdiff_t... StaticExtents>
 125 struct extents_analyse<R, E0, StaticExtents...>
 126 {
 127
 128     //! The extent analysis of the next lower rank.
 129     using next_extents_analyse = extents_analyse<R - 1, StaticExtents...>;
 130
 131     /*! \brief Accumulate the total rank from all extents.
 132      * \returns incremented rank of the next extent
 133      */
 134     static constexpr std::size_t rank() noexcept { return next_extents_analyse::rank() + 1; }
 135     /*! \brief Accumulate the dynamic rank from all extents.
 136      * This extent is static, so hand down query to the next extent analysis.
 137      * \returns the dynamic rank of the next extent analysis.
 138      */
 139     static constexpr std::size_t rank_dynamic() noexcept
 140     {
 141         return next_extents_analyse::rank_dynamic();
 142     }
 143
 144     //! Store analysis of the next extent of next lower rank.
 145     next_extents_analyse next;
 146
 147     //! Trivial constructor.
 148     constexpr extents_analyse() : next() {}
 149
 150     /*! \brief Construction from dynamic extents hands the extents down
 151      * to the next extents analysis of lower rank.
 152      * \param[in] de dynamic extents
 153      */
 154     template<class... DynamicExtents>
 155     constexpr extents_analyse(DynamicExtents... de) : next(de...)
 156     {
 157     }
 158
 159     /*! \brief Construct from an array of dynamic extentes and rank.
 160      * Hand down the dynamic rank parameters to the next extents analysis rank
 161      * \param[in] de dynamic extents
 162      * \param[in] r rank to read from the dynamic extent
 163      */
 164     template<std::size_t Rank>
 165     constexpr extents_analyse(const std::array<std::ptrdiff_t, Rank>& de, const std::size_t r) :
 166         next(de, r)
 167     {
 168     }
 169
 170     //! Copy constructor.
 171     template<std::ptrdiff_t... OtherStaticExtents>
 172     extents_analyse(extents_analyse<R, OtherStaticExtents...> rhs) : next(rhs.next)
 173     {
 174     }
 175
 176     //! Assignment operator.
 177     template<std::ptrdiff_t... OtherStaticExtents>
 178     extents_analyse& operator=(extents_analyse<R, OtherStaticExtents...> rhs)
 179     {
 180         next = rhs.next;
 181         return *this;
 182     }
 183
 184     /*! \brief Report extent of dimension r.
 185      * \param[in] r the dimension to query
 186      * \returns the extent in dimension r.
 187      */
 188     constexpr std::ptrdiff_t extent(const std::size_t r) const noexcept
 189     {
 190         return (r == R) ? E0 : next.extent(r);
 191     }
 192     /*! \brief Report the static extent of dimension r.
 193      * \param[in] r the dimension to query
 194      * \returns the static extent in dimension r.
 195      */
 196     static constexpr std::ptrdiff_t static_extent(const std::size_t r) noexcept
 197     {
 198         return (r == R) ? E0 : next_extents_analyse::static_extent(r);
 199     }
 200
 201     //! Returns the extent with the first dimension sliced off
 202     constexpr auto sliced_extents() const noexcept { return next; }
 203 };
 204
 205 /*! \libinternal \brief Enable querying extent of specific rank by splitting
 206  * a dynamic extent off the variadic template arguments.
 207  */
 208 template<int R, std::ptrdiff_t... StaticExtents>
 209 struct extents_analyse<R, dynamic_extent, StaticExtents...>
 210 {
 211     //! The extent analysis of the next lower rank.
 212     using next_extents_analyse = extents_analyse<R - 1, StaticExtents...>;
 213     /*! \brief Accumulate the total rank from all extents.
 214      * \returns incremented rank of the next extent
 215      */
 216     static constexpr std::size_t rank() noexcept { return next_extents_analyse::rank() + 1; }
 217     /*! \brief Accumulate the dynamic rank from all extents.
 218      * \returns the dynamic rank of the next extent analysis.
 219      */
 220     static constexpr std::size_t rank_dynamic() noexcept
 221     {
 222         return next_extents_analyse::rank_dynamic() + 1;
 223     }
 224
 225     //! Store analysis of the next extent of next lower rank.
 226     next_extents_analyse next;
 227     //! The dynamic extent of this rank
 228     std::ptrdiff_t this_extent;
 229
 230     //! Trivial constructor.
 231     extents_analyse() : next(), this_extent(0) {}
 232
 233     /*! \brief Construction from dynamic extents hands the extents down
 234      * to the next extents analysis of lower rank.
 235      * \param[in] E the dynamic extent of this rank.
 236      * \param[in] de dynamic extents
 237      */
 238     template<class... DynamicExtents>
 239     extents_analyse(std::ptrdiff_t E, DynamicExtents... de) : next(de...), this_extent(E)
 240     {
 241     }
 242
 243     /*! \brief Construct from an array of dynamic extentes and rank.
 244      * Hand down the dynamic rank parameters to the next extents analysis rank
 245      * \param[in] de dynamic extents
 246      * \param[in] r rank to read from the dynamic extent
 247      */
 248     template<std::size_t Rank>
 249     extents_analyse(const std::array<std::ptrdiff_t, Rank>& de, const std::size_t r) :
 250         next(de, r + 1), this_extent(de[r])
 251     {
 252     }
 253
 254     //! Copy constructor.
 255     template<std::ptrdiff_t... OtherStaticExtents>
 256     extents_analyse(extents_analyse<R, OtherStaticExtents...> rhs) :
 257         next(rhs.next), this_extent(rhs.extent(R))
 258     {
 259     }
 260
 261     //! Assignment operator.
 262     template<std::ptrdiff_t... OtherStaticExtents>
 263     extents_analyse& operator=(extents_analyse<R, OtherStaticExtents...> rhs)
 264     {
 265         next        = rhs.next;
 266         this_extent = rhs.extent(R);
 267         return *this;
 268     }
 269
 270     /*! \brief Report extent of dimension r.
 271      * \param[in] r the dimension to query
 272      * \returns the extent in dimension r.
 273      */
 274     constexpr std::ptrdiff_t extent(const std::size_t r) const noexcept
 275     {
 276         return (r == R) ? this_extent : next.extent(r);
 277     }
 278     /*! \brief Report the static extent of dimension r.
 279      * \param[in] r the dimension to query
 280      * \returns the static extent in dimension r.
 281      */
 282     static constexpr std::ptrdiff_t static_extent(const std::size_t r) noexcept
 283     {
 284         return (r == R) ? dynamic_extent : next_extents_analyse::static_extent(r);
 285     }
 286
 287     //! Returns the extent with the first dimension sliced off
 288     constexpr auto sliced_extents() const noexcept { return next; }
 289 };
 290
 291 /*! \libinternal \brief Specialisation for rank 0 extents analysis.
 292  * Ends recursive rank analysis.
 293  */
 294 template<>
 295 struct extents_analyse<0>
 296 {
 297     /*! \brief Rank of extent of rank 0.
 298      * \returns 0
 299      */
 300     static constexpr std::size_t rank() noexcept { return 0; }
 301     /*! \brief Dynamic rank of extent of rank 0.
 302      * \returns 0
 303      */
 304     static constexpr std::size_t rank_dynamic() noexcept { return 0; }
 305
 306     //! Trivial constructor.
 307     constexpr extents_analyse() {}
 308
 309     //! Construct from array and rank, doing nothing.
 310     template<std::size_t Rank>
 311     extents_analyse(const std::array<std::ptrdiff_t, Rank>& /*de*/, const std::size_t /*r*/)
 312     {
 313     }
 314
 315     // extents_analyse & operator=(extents_analyse) = default;
 316
 317     /*! \brief Extent of rank 0 is 1, ensuring that product of extents yields required size and not zero.
 318      * NOTE changed from ORNL reference implementation in making this static constexpr instead of constexpr .. const
 319      */
 320     static constexpr std::ptrdiff_t extent(const std::size_t /*r*/) noexcept { return 1; }
 321
 322     //! Static extent of rank 0 is 1, ensuring that product of extents yields required size and not zero.
 323     static constexpr std::ptrdiff_t static_extent(const std::size_t /*r*/) noexcept { return 1; }
 324 };
 325
 326 template<std::ptrdiff_t E0, std::ptrdiff_t... StaticExtents>
 327 struct sliced_extents
 328 {
 329     using type = extents<StaticExtents...>;
 330 };
 331 } // namespace detail
 332
 333 /*! \libinternal \brief Multidimensional extents with static and dynamic dimensions.
 334  *
 335  * Describes a multidimensional index space of rank R.
 336  * This is equivalent to the Cartesian product space of integer intervals
 337  * [0, N_0) x [0, N_1) x ... x [0,N_{R-1} )
 338  *
 339  * Confer to P0009r8 of the Library Evolution Working Group and mdspan.extents
 340  *
 341  * \tparam StaticExtents rank number of extents, where the dynamic_extent
 342  * constant for static extent is used to signal a dynamic extent.
 343  */
 344 template<std::ptrdiff_t... StaticExtents>
 345 class extents
 346 {
 347 private:
 348     using extents_analyse_t = detail::extents_analyse<sizeof...(StaticExtents), StaticExtents...>;
 349
 350 public:
 351     //! Type used to index elements.
 352     using index_type = std::ptrdiff_t;
 353     //! Trivial constructor
 354     constexpr extents() noexcept {}
 355     //! Move constructor
 356     constexpr extents(extents&&) noexcept = default;
 357     //! Copy constructor.
 358     constexpr extents(const extents&) noexcept = default;
 359     /*! \brief Construct with dynamic extents.
 360      *
 361      * Allows for extents(u,v,w..) syntax when setting dynamic extents
 362      *
 363      * \tparam IndexType type of index
 364      * \param[in] dn first dynamic index
 365      * \param[in] DynamicExtents parameter pack
 366      */
 367     template<class... IndexType>
 368     constexpr extents(std::ptrdiff_t dn, IndexType... DynamicExtents) noexcept :
 369         impl(dn, DynamicExtents...)
 370     {
 371         static_assert(1 + sizeof...(DynamicExtents) == rank_dynamic(), "");
 372     }
 373
 374     /*! \brief Construct from array of dynamic extents.
 375      *
 376      * Allows for extents({{u,v,w..}}) syntax when setting dynamic extents
 377      *
 378      * \param[in] dynamic_extents array of dynamic rank size containing extents
 379      */
 380     constexpr extents(const std::array<std::ptrdiff_t, extents_analyse_t::rank_dynamic()> dynamic_extents) noexcept
 381         :
 382         impl(dynamic_extents, 0)
 383     {
 384     }
 385
 386     //! Copy constructor
 387     template<std::ptrdiff_t... OtherStaticExtents>
 388     extents(const extents<OtherStaticExtents...>& other) : impl(other.impl)
 389     {
 390     }
 391
 392     //! Default move assignment
 393     extents& operator=(extents&&) noexcept = default;
 394     //! Default copy assignment
 395     extents& operator=(const extents&) noexcept = default;
 396     //! Copy assignment
 397     template<std::ptrdiff_t... OtherStaticExtents>
 398     extents& operator=(const extents<OtherStaticExtents...>& other)
 399     {
 400         impl = other.impl;
 401         return *this;
 402     }
 403     //! Default destructor
 404     ~extents() = default;
 405
 406     // [mdspan.extents.obs]
 407     /*! \brief The rank of the extent.
 408      * \returns the rank all extents together
 409      */
 410     static constexpr std::size_t rank() noexcept { return sizeof...(StaticExtents); }
 411     /*! \brief The rank of the dynamic extents.
 412      * \returns Only the dynamic extents.
 413      */
 414     static constexpr std::size_t rank_dynamic() noexcept
 415     {
 416         return extents_analyse_t::rank_dynamic();
 417     }
 418     /*! \brief The rank of the static extents.
 419      * \returns Only the static extents.
 420      */
 421     static constexpr index_type static_extent(std::size_t k) noexcept
 422     {
 423         return extents_analyse_t::static_extent(rank() - k);
 424     }
 425     /*! \brief The extent along a specific dimension.
 426      * \param[in] k the dimension
 427      * \returns the extent along that dimension
 428      */
 429     constexpr index_type extent(std::size_t k) const noexcept { return impl.extent(rank() - k); }
 430     //! Returns the extent with the first dimension sliced off
 431     constexpr auto sliced_extents() const noexcept
 432     {
 433         return typename detail::sliced_extents<StaticExtents...>::type(impl.sliced_extents());
 434     }
 435
 436 private:
 437     extents(extents_analyse_t o) : impl(o) {}
 438     //! For copy assignment, extents are friends of extents.
 439     template<std::ptrdiff_t...>
 440     friend class extents;
 441     //! The implementation class.
 442     extents_analyse_t impl;
 443 };
 444
 445
 446 /*! \brief Comparison operator.
 447  * \returns true if extents are equal
 448  */
 449 template<std::ptrdiff_t... LHS, std::ptrdiff_t... RHS>
 450 constexpr bool operator==(const extents<LHS...>& lhs, const extents<RHS...>& rhs) noexcept
 451 {
 452     bool equal = lhs.rank() == rhs.rank();
 453     for (std::size_t r = 0; r < lhs.rank(); r++)
 454     {
 455         equal = equal && (lhs.extent(r) == rhs.extent(r));
 456     }
 457     return equal;
 458 }
 459
 460 /*! \brief Check for non-equality.
 461  * \returns true if extents are unequal
 462  */
 463 template<std::ptrdiff_t... LHS, std::ptrdiff_t... RHS>
 464 constexpr bool operator!=(const extents<LHS...>& lhs, const extents<RHS...>& rhs) noexcept
 465 {
 466     return !(lhs == rhs);
 467 }
 468
 469 } // namespace gmx
 470 #endif /* end of include guard: MDSPAN_EXTENTS_H */