distribution.hpp

Go to the documentation of this file.

#ifndef DISTRIBUTION_HPP
#define DISTRIBUTION_HPP

#include <Kokkos_Core.hpp>
#include "space_settings.hpp"
#include "globals.hpp"
#include "velocity_grid.hpp"

#ifndef NO_FORTRAN
extern "C" double* get_f_loc();
extern "C" double* get_f0g_loc();
extern "C" double* get_temp_ev_loc();
extern "C" double* get_df0g_tmp_loc();
extern "C" double* get_delta_n_loc();
extern "C" double* get_delta_u_loc();
extern "C" double* get_delta_t_loc();
#endif

template<class Device>
class Distribution{

    public:

    // Let f view just wrap around fortran array instead of making a copy
    // Used in collisions, but stay on the host
    Kokkos::View<double***,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f;
    Kokkos::View<double***,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> df0g_tmp;

    // Used in getf0, copied to gpu so use _h to distinguish the host versions
    // If F0_TOR_LINEAR is used, f0g has a 4th dimension of size 2 (to enable a toroidal interpolation)
#if defined(XGC1) && defined(F0_TOR_LINEAR)
    Kokkos::View<double****,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f0g_h;
#else
    Kokkos::View<double***,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f0g_h;
#endif
    Kokkos::View<double*,Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> temp_ev_h;

    // Arrays on Device:
#if defined(XGC1) && defined(F0_TOR_LINEAR)
    Kokkos::View<double****,Kokkos::LayoutRight,Device> f0g;
#else
    Kokkos::View<double***,Kokkos::LayoutRight,Device> f0g;
#endif
    Kokkos::View<double*,Kokkos::LayoutRight,Device> temp_ev;

    // These arrays are for the entire mesh on the fortran side; could be decomposed
    Kokkos::View<double*,Kokkos::LayoutRight,Device> delta_n;
    Kokkos::View<double*,Kokkos::LayoutRight,Device> delta_u;
    Kokkos::View<double*,Kokkos::LayoutRight,Device> delta_t;

    // Used in distribute_f0g
    Kokkos::View<double*,Kokkos::LayoutRight,Device> grid_vol;

    // Used in update_f0
    Kokkos::View<double*,Kokkos::LayoutRight,Device> inv_grid_vol;
    Kokkos::View<double***,Kokkos::LayoutRight,Device> n;
    Kokkos::View<double***,Kokkos::LayoutRight, Device> f_d;
    Kokkos::View<double*,Kokkos::LayoutRight,Device> temp_ev_d;

    // Initialize these Views but don't allocate any memory yet; Instead, resize when needed
    Distribution()
#if defined(XGC1) && defined(F0_TOR_LINEAR)
        : f0g("f0g",0,0,0,0),
#else
        : f0g("f0g",0,0,0),
#endif
          temp_ev("temp_ev",0),
          delta_n("delta_n",0),
          delta_u("delta_u",0),
          delta_t("delta_t",0),
          grid_vol("grid_vol",0),
          inv_grid_vol("inv_grid_vol",0),
          temp_ev_d("temp_ev_d",0),
          f_d("f_d",0,0,0),
          n("n",0,0,0)
          {}


    // Allocate memory for the device arrays by resizing them
    void resize_on_device(const VelocityGrid& vgrid, int nnode, int nnode_entire_plane){
#if defined(XGC1) && defined(F0_TOR_LINEAR)
        Kokkos::resize(f0g, 2, vgrid.nvr, nnode, vgrid.nvz);
#else
        Kokkos::resize(f0g, vgrid.nvr, nnode, vgrid.nvz);
#endif
        Kokkos::resize(temp_ev, nnode);
        Kokkos::resize(delta_n, nnode_entire_plane);
        Kokkos::resize(delta_u, nnode_entire_plane);
        Kokkos::resize(delta_t, nnode_entire_plane);
    }

    // Deallocate memory of the device arrays by resizing back to 0
    void resize_to_zero(){
#if defined(XGC1) && defined(F0_TOR_LINEAR)
        Kokkos::resize(f0g, 0,0,0,0);
#else
        Kokkos::resize(f0g, 0,0,0);
#endif
        Kokkos::resize(temp_ev, 0);
        Kokkos::resize(delta_n, 0);
        Kokkos::resize(delta_u, 0);
        Kokkos::resize(delta_t, 0);
    }

    // Send arrays needed for weight udpate
    void send_to_device(double* delta_n_array,double* delta_u_array,double* delta_t_array,int isp)
    {
        subarray_deep_copy(delta_n, delta_n_array, isp);
        subarray_deep_copy(delta_u, delta_u_array, isp);
        subarray_deep_copy(delta_t, delta_t_array, isp);

        Kokkos::deep_copy(f0g, f0g_h);
        Kokkos::deep_copy(temp_ev, temp_ev_h);
    }

    // Allocate memory for the device arrays by resizing them
    void resize_for_distribute_f0g_on_device(const VelocityGrid& vgrid, int nnode, int nnode_entire_plane){
        Kokkos::resize(grid_vol, nnode);
        Kokkos::resize(temp_ev_d, nnode);
    }

    // Deallocate memory of the device arrays by resizing back to 0
    void resize_for_distribute_f0g_to_zero(){
        Kokkos::resize(grid_vol, 0);
        Kokkos::resize(temp_ev_d, 0);
    }

    // Allocate memory for the device arrays by resizing them
    void resize_for_update_f0_on_device(const VelocityGrid& vgrid, int nnode, int nnode_entire_plane){
        Kokkos::resize(f0g, vgrid.nvr, nnode, vgrid.nvz);
        Kokkos::resize(inv_grid_vol, nnode);
        Kokkos::resize(f_d, vgrid.nvr, nnode, vgrid.nvz);
        Kokkos::resize(n, vgrid.nvr, nnode, vgrid.nvz);
        Kokkos::resize(temp_ev_d, nnode);
    }

    // Deallocate memory of the device arrays by resizing back to 0
    void resize_for_update_f0_to_zero(){
        Kokkos::resize(f0g, 0,0,0);
        Kokkos::resize(inv_grid_vol, 0);
        Kokkos::resize(f_d, 0,0,0);
        Kokkos::resize(n, 0,0,0);
        Kokkos::resize(temp_ev_d, 0);
    }

    /********************************/

    // isp here is the species index in f0, i.e. non-adiabatic species
    void update_pointers(const VelocityGrid& vgrid, int nnode, int isp, double* f0g_ptr, double* temp_ev_ptr )
    {
        // Need to update pointers to fortran arrays containing grid info
        // Since these pointers point to specific species, must add offset to get to the correct species
        // This all goes away once the distributions are resident on the C++ side
        int one_sp_size = vgrid.nvr*nnode*vgrid.nvz;
        int ptr_offset = isp*one_sp_size;

#if defined(XGC1) && defined(F0_TOR_LINEAR)
        f0g_h = Kokkos::View<double****,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
            (f0g_ptr + ptr_offset,    2, vgrid.nvr, nnode, vgrid.nvz);
#else
        f0g_h = Kokkos::View<double***,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
            (f0g_ptr + ptr_offset,    vgrid.nvr, nnode, vgrid.nvz);
#endif

        // 1D array: temp_ev
        one_sp_size = nnode;
        ptr_offset = isp*one_sp_size;
        temp_ev_h = Kokkos::View<double*,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
            (temp_ev_ptr + ptr_offset,    nnode);
    }

    // This pointer is to a temporary array allocated in fortran right before collisions, so it must be updated
    // at the beginning of collisions
    // isp here is the species index in f0, i.e. non-adiabatic species
    void update_tmp_pointer(const VelocityGrid& vgrid, int nnode, int isp, double* f_ptr, double* df0g_tmp_ptr )
    {
        // Need to update pointers to fortran arrays containing grid info
        // Since these pointers point to specific species, must add offset to get to the correct species
        // This all goes away once the distributions are resident on the C++ side
        int one_sp_size = vgrid.nvr*nnode*vgrid.nvz;
        int ptr_offset = isp*one_sp_size;
        f = Kokkos::View<double***,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
            (f_ptr + ptr_offset,    vgrid.nvr, nnode, vgrid.nvz);
        df0g_tmp = Kokkos::View<double***,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
            (df0g_tmp_ptr + ptr_offset, vgrid.nvr, nnode, vgrid.nvz);
    }

};
#endif