celeritas/field/RZMapFieldData.covfie.hh

0001 //------------------------------- -*- C++ -*- -------------------------------//
0002 // Copyright Celeritas contributors: see top-level COPYRIGHT file for details
0003 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0004 //---------------------------------------------------------------------------//
0005 //! \file celeritas/field/RZMapFieldData.covfie.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include <memory>
0010 #include <type_traits>
0011
0012 #include "corecel/Assert.hh"
0013 #include "corecel/Macros.hh"
0014 #include "corecel/Types.hh"
0015 #include "corecel/data/DeviceVector.hh"
0016 #include "celeritas/Types.hh"
0017
0018 #include "RZMapFieldData.hh"  // primary template + const_reference/device spec
0019
0020 #include "detail/CovfieRZFieldTraits.hh"
0021
0022 namespace celeritas
0023 {
0024 //---------------------------------------------------------------------------//
0025 template<MemSpace M>
0026 struct RZMapFieldParamsDataBase
0027 {
0028     using field_t = typename detail::CovfieRZFieldTraits<M>::field_t;
0029     using view_t = typename field_t::view_t;
0030
0031     FieldDriverOptions options;
0032     real_type min_r{};
0033     real_type max_r{};
0034     real_type min_z{};
0035     real_type max_z{};
0036 };
0037
0038 // Specializations for value/host, const_reference/host, and value/device are
0039 // defined below. const_reference/device is in RZMapFieldData.hh (no covfie
0040 // dependency, safe for C++17 TUs).
0041
0042 template<>
0043 struct RZMapFieldParamsData<Ownership::value, MemSpace::host>
0044     : RZMapFieldParamsDataBase<MemSpace::host>
0045 {
0046     CELER_FUNCTION view_t get_view() const { return view_t(*field); }
0047
0048     CELER_FUNCTION explicit operator bool() const { return field.get(); }
0049
0050     std::unique_ptr<field_t> field;
0051 };
0052 template<>
0053 struct RZMapFieldParamsData<Ownership::const_reference, MemSpace::host>
0054     : RZMapFieldParamsDataBase<MemSpace::host>
0055 {
0056     CELER_FUNCTION view_t const& get_view() const { return field_view; }
0057
0058     CELER_FUNCTION explicit operator bool() const { return true; }
0059
0060     view_t field_view;
0061 };
0062
0063 template<>
0064 struct RZMapFieldParamsData<Ownership::value, MemSpace::device>
0065     : RZMapFieldParamsDataBase<MemSpace::device>
0066 {
0067     view_t const& get_view() const { return field_view.device_ref()[0]; }
0068
0069     explicit operator bool() const
0070     {
0071         return field.get() && field_view.size() == 1;
0072     }
0073
0074     RZMapFieldParamsData& operator=(
0075         RZMapFieldParamsData<Ownership::value, MemSpace::host> const& other)
0076     {
0077         using host_field_t
0078             = detail::CovfieRZFieldTraits<MemSpace::host>::field_t;
0079         if constexpr (!std::is_same_v<field_t, host_field_t>)
0080         {
0081             // Use covfie's cross-type field constructor: propagates through
0082             // all transformer layers (affine->clamp->linear->strided), with
0083             // strided's cross-type constructor performing the H2D transfer
0084             // into cuda_device_array at the bottom of the chain.
0085             field = std::make_unique<field_t>(*other.field);
0086
0087             // Store view_t in device memory; pass pointer to kernel
0088             field_view = DeviceVector<view_t>{1};
0089             field_view.copy_to_device(make_span<view_t const>({{*field}}));
0090         }
0091         else
0092         {
0093             field = std::make_unique<field_t>(*other.field);
0094         }
0095         options = other.options;
0096         min_r = other.min_r;
0097         max_r = other.max_r;
0098         min_z = other.min_z;
0099         max_z = other.max_z;
0100         return *this;
0101     }
0102
0103     std::unique_ptr<field_t> field;
0104     DeviceVector<view_t> field_view;
0105 };
0106
0107 //---------------------------------------------------------------------------//
0108 /*!
0109  * Convert shared parameter data to a covfie field view.
0110  */
0111 template<MemSpace M>
0112 CELER_FUNCTION inline auto get_rz_map_field_view(
0113     RZMapFieldParamsData<Ownership::const_reference, M> const& params) ->
0114     typename detail::CovfieRZFieldTraits<M>::field_t::view_t const&
0115 {
0116     if constexpr (M == MemSpace::host)
0117     {
0118         return params.get_view();
0119     }
0120     else
0121     {
0122         using view_t = typename detail::CovfieRZFieldTraits<M>::field_t::view_t;
0123         CELER_EXPECT(params.field_view);
0124         return *static_cast<view_t const*>(params.field_view);
0125     }
0126 }
0127
0128 //---------------------------------------------------------------------------//
0129 inline auto
0130 RZMapFieldParamsData<Ownership::const_reference, MemSpace::device>::operator=(
0131     RZMapFieldParamsData<Ownership::value, MemSpace::device> const& other)
0132     -> RZMapFieldParamsData&
0133 {
0134     field_view = static_cast<void const*>(&other.get_view());
0135     options = other.options;
0136     min_r = other.min_r;
0137     max_r = other.max_r;
0138     min_z = other.min_z;
0139     max_z = other.max_z;
0140     return *this;
0141 }
0142
0143 //---------------------------------------------------------------------------//
0144 }  // namespace celeritas