backend/transformer/morton.hpp

0001 /*
0002  * SPDX-PackageName: "covfie, a part of the ACTS project"
0003  * SPDX-FileCopyrightText: 2022 CERN
0004  * SPDX-License-Identifier: MPL-2.0
0005  */
0006
0007 #pragma once
0008
0009 #include <algorithm>
0010 #include <climits>
0011 #include <iostream>
0012 #include <memory>
0013 #include <numeric>
0014 #include <type_traits>
0015
0016 #if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__)) &&        \
0017     defined(__BMI2__)
0018 #define HAVE_BMI2
0019 #include <x86intrin.h>
0020 #endif
0021
0022 #include <covfie/core/backend/primitive/array.hpp>
0023 #include <covfie/core/concepts.hpp>
0024 #include <covfie/core/definitions.hpp>
0025 #include <covfie/core/parameter_pack.hpp>
0026 #include <covfie/core/qualifiers.hpp>
0027 #include <covfie/core/utility/binary_io.hpp>
0028 #include <covfie/core/utility/nd_map.hpp>
0029 #include <covfie/core/utility/nd_size.hpp>
0030 #include <covfie/core/utility/numeric.hpp>
0031 #include <covfie/core/vector.hpp>
0032
0033 namespace covfie::backend {
0034 #ifdef HAVE_BMI2
0035 template <typename Ix, typename Ox, std::size_t N>
0036 struct morton_pdep_mask {
0037     template <std::size_t I>
0038     struct get_mask {
0039         template <typename>
0040         struct get_mask_helper {
0041         };
0042
0043         template <std::size_t... Js>
0044         struct get_mask_helper<std::index_sequence<Js...>> {
0045             template <Ox S>
0046             struct shiftl {
0047                 static constexpr Ox value = static_cast<Ox>(1) << S;
0048             };
0049
0050             static constexpr Ox value =
0051                 ((std::conditional_t<
0052                      Js % N == 0,
0053                      std::integral_constant<Ox, shiftl<Js>::value>,
0054                      std::integral_constant<Ox, 0>>::value) |
0055                  ...);
0056         };
0057
0058         static constexpr Ox value =
0059             get_mask_helper<
0060                 std::make_index_sequence<CHAR_BIT * sizeof(Ox)>>::value
0061             << I;
0062     };
0063
0064     template <typename C, std::size_t... Idxs>
0065     static constexpr Ox compute(C c, std::index_sequence<Idxs...>)
0066     {
0067         return (_pdep_u64(c[Idxs], get_mask<Idxs>::value) | ...);
0068     }
0069
0070     template <typename C, typename Ids = std::make_index_sequence<N>>
0071     static constexpr Ox compute(C c)
0072     {
0073         return compute(std::forward<C>(c), Ids{});
0074     }
0075 };
0076 #endif
0077
0078 template <
0079     concepts::vector_descriptor _input_vector_t,
0080     concepts::field_backend _storage_t,
0081     bool use_bmi2 = true>
0082 struct morton {
0083     using this_t = morton<_input_vector_t, _storage_t>;
0084     static constexpr bool is_initial = false;
0085
0086     using backend_t = _storage_t;
0087
0088     using contravariant_input_t =
0089         covfie::vector::array_vector_d<_input_vector_t>;
0090     using contravariant_output_t = typename backend_t::contravariant_input_t;
0091     using covariant_input_t = typename backend_t::covariant_output_t;
0092     using covariant_output_t = covariant_input_t;
0093
0094     using array_t = backend_t;
0095
0096     using configuration_t = utility::nd_size<contravariant_input_t::dimensions>;
0097
0098     static constexpr uint32_t IO_MAGIC_HEADER = 0xAB020006;
0099
0100     COVFIE_HOST_DEVICE static std::size_t
0101     calculate_index(typename contravariant_input_t::vector_t c)
0102     {
0103 #ifdef HAVE_BMI2
0104         if constexpr (use_bmi2) {
0105             return morton_pdep_mask<
0106                 typename contravariant_input_t::scalar_t,
0107                 typename contravariant_output_t::scalar_t,
0108                 contravariant_input_t::dimensions>::compute(c);
0109         } else {
0110             std::size_t idx = 0;
0111             for (std::size_t i = 0;
0112                  i < ((CHAR_BIT *
0113                        sizeof(typename contravariant_output_t::scalar_t)) /
0114                       contravariant_input_t::dimensions);
0115                  ++i)
0116             {
0117                 for (std::size_t j = 0; j < contravariant_input_t::dimensions;
0118                      ++j)
0119                 {
0120                     idx |= (c[j] & (1UL << i))
0121                            << (i * (contravariant_input_t::dimensions - 1) + j);
0122                 }
0123             }
0124             return idx;
0125         }
0126 #else
0127         std::size_t idx = 0;
0128         for (std::size_t i = 0;
0129              i <
0130              ((CHAR_BIT * sizeof(typename contravariant_output_t::scalar_t)) /
0131               contravariant_input_t::dimensions);
0132              ++i)
0133         {
0134             for (std::size_t j = 0; j < contravariant_input_t::dimensions; ++j)
0135             {
0136                 idx |= (c[j] & (static_cast<std::size_t>(1) << i))
0137                        << (i * (contravariant_input_t::dimensions - 1) + j);
0138             }
0139         }
0140         return idx;
0141 #endif
0142     }
0143
0144     template <typename T>
0145     static std::unique_ptr<
0146         std::decay_t<typename backend_t::covariant_output_t::vector_t>[]>
0147     make_morton_copy(const T & other)
0148     {
0149         configuration_t sizes = other.get_configuration();
0150         std::unique_ptr<
0151             std::decay_t<typename backend_t::covariant_output_t::vector_t>[]>
0152             res = std::make_unique<std::decay_t<
0153                 typename backend_t::covariant_output_t::vector_t>[]>(
0154                 utility::ipow(
0155                     utility::round_pow2(
0156                         *std::max_element(sizes.begin(), sizes.end())
0157                     ),
0158                     contravariant_input_t::dimensions
0159                 )
0160             );
0161         typename T::parent_t::non_owning_data_t nother(other);
0162
0163         utility::nd_map<decltype(sizes)>(
0164             [&nother, &res](decltype(sizes) t) {
0165                 typename contravariant_input_t::vector_t c;
0166
0167                 for (std::size_t i = 0; i < contravariant_input_t::dimensions;
0168                      ++i) {
0169                     c[i] = static_cast<
0170                         typename contravariant_input_t::vector_t::value_type>(
0171                         t[i]
0172                     );
0173                 }
0174
0175                 std::size_t idx = calculate_index(c);
0176
0177                 for (std::size_t i = 0; i < covariant_output_t::dimensions; ++i)
0178                 {
0179                     res[idx][i] = nother.at(c)[i];
0180                 }
0181             },
0182             sizes
0183         );
0184
0185         return res;
0186     }
0187
0188     struct owning_data_t {
0189         using parent_t = this_t;
0190
0191         owning_data_t() = default;
0192
0193         template <typename T>
0194         requires(std::same_as<
0195                  typename T::parent_t::configuration_t,
0196                  configuration_t> &&
0197                      std::constructible_from<
0198                          typename backend_t::owning_data_t,
0199                          std::size_t,
0200                          std::add_rvalue_reference_t<std::unique_ptr<std::decay_t<
0201                              typename backend_t::covariant_output_t::
0202                                  vector_t>[]>>>) explicit owning_data_t(const T &
0203                                                                             o)
0204             : m_sizes(o.get_configuration())
0205             , m_storage(
0206                   utility::ipow(
0207                       utility::round_pow2(
0208                           *std::max_element(m_sizes.begin(), m_sizes.end())
0209                       ),
0210                       contravariant_input_t::dimensions
0211                   ),
0212                   make_morton_copy(o)
0213               )
0214         {
0215         }
0216
0217         template <typename... Args>
0218         requires(sizeof...(Args) > 0) explicit owning_data_t(
0219             parameter_pack<configuration_t, Args...> && args
0220         )
0221             : m_sizes(args.x)
0222             , m_storage(std::move(args.xs))
0223         {
0224         }
0225
0226         template <typename T>
0227         requires(std::constructible_from<
0228                  owning_data_t,
0229                  T>) explicit owning_data_t(parameter_pack<T> && args)
0230             : owning_data_t(args.x)
0231         {
0232         }
0233
0234         explicit owning_data_t(
0235             const configuration_t & c, typename backend_t::owning_data_t && b
0236         )
0237             : m_sizes(c)
0238             , m_storage(std::forward<typename backend_t::owning_data_t>(b))
0239         {
0240         }
0241
0242         typename backend_t::owning_data_t & get_backend(void)
0243         {
0244             return m_storage;
0245         }
0246
0247         const typename backend_t::owning_data_t & get_backend(void) const
0248         {
0249             return m_storage;
0250         }
0251
0252         configuration_t get_configuration(void) const
0253         {
0254             return m_sizes;
0255         }
0256
0257         static owning_data_t read_binary(std::istream & fs)
0258         {
0259             utility::read_io_header(fs, IO_MAGIC_HEADER);
0260
0261             auto sizes = utility::read_binary<decltype(m_sizes)>(fs);
0262             auto be = backend_t::owning_data_t::read_binary(fs);
0263
0264             utility::read_io_footer(fs, IO_MAGIC_HEADER);
0265
0266             return owning_data_t(sizes, std::move(be));
0267         }
0268
0269         static void write_binary(std::ostream & fs, const owning_data_t & o)
0270         {
0271             utility::write_io_header(fs, IO_MAGIC_HEADER);
0272
0273             fs.write(
0274                 reinterpret_cast<const char *>(&o.m_sizes),
0275                 sizeof(decltype(o.m_sizes))
0276             );
0277
0278             backend_t::owning_data_t::write_binary(fs, o.m_storage);
0279
0280             utility::write_io_footer(fs, IO_MAGIC_HEADER);
0281         }
0282
0283         configuration_t m_sizes;
0284         typename backend_t::owning_data_t m_storage;
0285     };
0286
0287     struct non_owning_data_t {
0288         using parent_t = this_t;
0289
0290         non_owning_data_t(const owning_data_t & o)
0291             : m_sizes(o.m_sizes)
0292             , m_storage(o.m_storage)
0293         {
0294         }
0295
0296         COVFIE_HOST_DEVICE typename covariant_output_t::vector_t
0297         at(typename contravariant_input_t::vector_t c) const
0298         {
0299 #ifndef NDEBUG
0300             for (std::size_t i = 0; i < contravariant_input_t::dimensions; ++i)
0301             {
0302                 assert(c[i] < m_sizes[i]);
0303             }
0304 #endif
0305
0306             return m_storage.at(calculate_index(c));
0307         }
0308
0309         typename backend_t::non_owning_data_t & get_backend(void)
0310         {
0311             return m_storage;
0312         }
0313
0314         const typename backend_t::non_owning_data_t & get_backend(void) const
0315         {
0316             return m_storage;
0317         }
0318
0319         configuration_t m_sizes;
0320         typename backend_t::non_owning_data_t m_storage;
0321     };
0322 };
0323 }