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"

// Functions to fetch fortran array addresses
extern "C" double* get_delta_n_loc();
extern "C" double* get_delta_u_loc();
extern "C" double* get_delta_t_loc();

template<class Device>
struct Distribution{

    // Used in getf0, copied to gpu so use _h to distinguish the host versions
    Kokkos::View<double***,Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f0g_h;
    Kokkos::View<double*,Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> temp_ev_h;

    View<double*,CLayout, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> den_global_h;
    View<double*,CLayout, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> temp_global_h;

    // Arrays on Device:
    Kokkos::View<double***,Kokkos::LayoutRight,Device> f0g;
    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()
        : f0g("f0g",0,0,0),
          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){
        Kokkos::resize(f0g, vgrid.nvr, nnode, vgrid.nvz);
        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(){
        Kokkos::resize(f0g, 0,0,0);
        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){
        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){
        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);
    }
};
#endif