Skip to content

Commit

Permalink
Merge pull request #391 from toxa81/develop
Browse files Browse the repository at this point in the history 
new api for cp2k
  • Loading branch information
@toxa81
toxa81 authored Jul 18, 2019
2 parents 1d01d2f + f01f2a9 commit 5d8320d
Show file tree
Hide file tree
Showing 3 changed files with 215 additions and 11 deletions.
5 changes: 5 additions & 0 deletions src/SDDK/fft3d.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1136,6 +1136,11 @@ class FFT3D : public FFT3D_grid
}
}
}

inline splindex<splindex_t::block> const& spl_z() const
{
return spl_z_;
}
};

} // namespace sddk
Expand Down
68 changes: 62 additions & 6 deletions src/generated.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2967,35 +2967,91 @@ end subroutine sirius_get_fv_eigen_values
!> @brief Set the values of the function on the regular grid.
!> @param [in] handler DFT ground state handler.
!> @param [in] label Label of the function.
!> @param [in] values Values of the function.
!> @param [in] grid_dims Dimensions of the FFT grid.
!> @param [in] local_box_origin Coordinates of the local box origin for each MPI rank
!> @param [in] local_box_size Dimensions of the local box for each MPI rank.
!> @param [in] fcomm Fortran communicator used to partition FFT grid into local boxes.
!> @param [in] values Values of the function (local buffer for each MPI rank).
!> @param [in] transform_to_pw If true, transform function to PW domain.
subroutine sirius_set_rg_values(handler,label,values,grid_dims,transform_to_pw)
subroutine sirius_set_rg_values(handler,label,grid_dims,local_box_origin,local_box_size,&
&fcomm,values,transform_to_pw)
implicit none
type(C_PTR), intent(in) :: handler
character(C_CHAR), dimension(*), intent(in) :: label
real(C_DOUBLE), intent(in) :: values
integer(C_INT), intent(in) :: grid_dims
integer(C_INT), intent(in) :: local_box_origin
integer(C_INT), intent(in) :: local_box_size
integer(C_INT), intent(in) :: fcomm
real(C_DOUBLE), intent(in) :: values
logical(C_BOOL), optional, target, intent(in) :: transform_to_pw
type(C_PTR) :: transform_to_pw_ptr
interface
subroutine sirius_set_rg_values_aux(handler,label,values,grid_dims,transform_to_pw)&
subroutine sirius_set_rg_values_aux(handler,label,grid_dims,local_box_origin,local_box_size,&
&fcomm,values,transform_to_pw)&
&bind(C, name="sirius_set_rg_values")
use, intrinsic :: ISO_C_BINDING
type(C_PTR), intent(in) :: handler
character(C_CHAR), dimension(*), intent(in) :: label
real(C_DOUBLE), intent(in) :: values
integer(C_INT), intent(in) :: grid_dims
integer(C_INT), intent(in) :: local_box_origin
integer(C_INT), intent(in) :: local_box_size
integer(C_INT), intent(in) :: fcomm
real(C_DOUBLE), intent(in) :: values
type(C_PTR), value :: transform_to_pw
end subroutine
end interface

transform_to_pw_ptr = C_NULL_PTR
if (present(transform_to_pw)) transform_to_pw_ptr = C_LOC(transform_to_pw)

call sirius_set_rg_values_aux(handler,label,values,grid_dims,transform_to_pw_ptr)
call sirius_set_rg_values_aux(handler,label,grid_dims,local_box_origin,local_box_size,&
&fcomm,values,transform_to_pw_ptr)
end subroutine sirius_set_rg_values

!> @brief Get the values of the function on the regular grid.
!> @param [in] handler DFT ground state handler.
!> @param [in] label Label of the function.
!> @param [in] grid_dims Dimensions of the FFT grid.
!> @param [in] local_box_origin Coordinates of the local box origin for each MPI rank
!> @param [in] local_box_size Dimensions of the local box for each MPI rank.
!> @param [in] fcomm Fortran communicator used to partition FFT grid into local boxes.
!> @param [out] values Values of the function (local buffer for each MPI rank).
!> @param [in] transform_to_rg If true, transform function to regular grid before fetching the values.
subroutine sirius_get_rg_values(handler,label,grid_dims,local_box_origin,local_box_size,&
&fcomm,values,transform_to_rg)
implicit none
type(C_PTR), intent(in) :: handler
character(C_CHAR), dimension(*), intent(in) :: label
integer(C_INT), intent(in) :: grid_dims
integer(C_INT), intent(in) :: local_box_origin
integer(C_INT), intent(in) :: local_box_size
integer(C_INT), intent(in) :: fcomm
real(C_DOUBLE), intent(out) :: values
logical(C_BOOL), optional, target, intent(in) :: transform_to_rg
type(C_PTR) :: transform_to_rg_ptr
interface
subroutine sirius_get_rg_values_aux(handler,label,grid_dims,local_box_origin,local_box_size,&
&fcomm,values,transform_to_rg)&
&bind(C, name="sirius_get_rg_values")
use, intrinsic :: ISO_C_BINDING
type(C_PTR), intent(in) :: handler
character(C_CHAR), dimension(*), intent(in) :: label
integer(C_INT), intent(in) :: grid_dims
integer(C_INT), intent(in) :: local_box_origin
integer(C_INT), intent(in) :: local_box_size
integer(C_INT), intent(in) :: fcomm
real(C_DOUBLE), intent(out) :: values
type(C_PTR), value :: transform_to_rg
end subroutine
end interface

transform_to_rg_ptr = C_NULL_PTR
if (present(transform_to_rg)) transform_to_rg_ptr = C_LOC(transform_to_rg)

call sirius_get_rg_values_aux(handler,label,grid_dims,local_box_origin,local_box_size,&
&fcomm,values,transform_to_rg_ptr)
end subroutine sirius_get_rg_values

!> @brief Get the total magnetization of the system.
!> @param [in] handler DFT ground state handler.
!> @param [out] mag 3D magnetization vector (x,y,z components).
Expand Down
153 changes: 148 additions & 5 deletions src/sirius_api.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2923,16 +2923,24 @@ void sirius_get_fv_eigen_values(void* const* handler__,
/* @fortran begin function void sirius_set_rg_values Set the values of the function on the regular grid.
@fortran argument in required void* handler DFT ground state handler.
@fortran argument in required string label Label of the function.
@fortran argument in required double values Values of the function.
@fortran argument in required int grid_dims Dimensions of the FFT grid.
@fortran argument in required int local_box_origin Coordinates of the local box origin for each MPI rank
@fortran argument in required int local_box_size Dimensions of the local box for each MPI rank.
@fortran argument in required int fcomm Fortran communicator used to partition FFT grid into local boxes.
@fortran argument in required double values Values of the function (local buffer for each MPI rank).
@fortran argument in optional bool transform_to_pw If true, transform function to PW domain.
@fortran end */
void sirius_set_rg_values(void* const* handler__,
char const* label__,
double const* values__,
int const* grid_dims__,
int const* local_box_origin__,
int const* local_box_size__,
int const* fcomm__,
double const* values__,
bool const* transform_to_pw__)
{
PROFILE("sirius_api::sirius_set_rg_values");

GET_GS(handler__);

std::string label(label__);
Expand All @@ -2957,9 +2965,46 @@ void sirius_set_rg_values(void* const* handler__,

try {
auto& f = func.at(label);
auto offset = gs.ctx().fft().offset_z() * gs.ctx().fft().size(0) * gs.ctx().fft().size(1);
/* copy local part of the buffer */
std::copy(values__ + offset, values__ + offset + gs.ctx().fft().local_size(), &f->f_rg(0));

auto& comm = Communicator::map_fcomm(*fcomm__);

mdarray<int, 2> local_box_size(const_cast<int*>(local_box_size__), 3, comm.size());
mdarray<int, 2> local_box_origin(const_cast<int*>(local_box_origin__), 3, comm.size());

for (int rank = 0; rank < comm.size(); rank++) {
/* dimensions of this rank's local box */
int nx = local_box_size(0, rank);
int ny = local_box_size(1, rank);
int nz = local_box_size(2, rank);

mdarray<double, 3> buf(nx, ny, nz);
/* if this is that rank's turn to broadcast */
if (comm.rank() == rank) {
/* copy values to buf */
std::copy(values__, values__ + nx * ny * nz, &buf[0]);
}
/* send a copy of local box to all ranks */
comm.bcast(&buf[0], nx * ny * nz, rank);

for (int iz = 0; iz < nz; iz++) {
/* global z coordinate inside FFT box */
int z = local_box_origin(2, rank) + iz - 1; /* Fortran counts from 1 */
/* each rank on SIRIUS side, for which this condition is fulfilled copies data from the local box */
if (z >= gs.ctx().fft().offset_z() && z < gs.ctx().fft().offset_z() + gs.ctx().fft().local_size_z()) {
/* make z local for SIRIUS FFT partitioning */
z -= gs.ctx().fft().offset_z();
for (int iy = 0; iy < ny; iy++) {
/* global y coordinate inside FFT box */
int y = local_box_origin(1, rank) + iy - 1; /* Fortran counts from 1 */
for (int ix = 0; ix < nx; ix++) {
/* global x coordinate inside FFT box */
int x = local_box_origin(0, rank) + ix - 1; /* Fortran counts from 1 */
f->f_rg(gs.ctx().fft().index_by_coord(x, y, z)) = buf(ix, iy, iz);
}
}
}
}
} /* loop over ranks */
if (transform_to_pw__ && *transform_to_pw__) {
f->fft_transform(-1);
}
Expand All @@ -2968,6 +3013,104 @@ void sirius_set_rg_values(void* const* handler__,
}
}

/* @fortran begin function void sirius_get_rg_values Get the values of the function on the regular grid.
@fortran argument in required void* handler DFT ground state handler.
@fortran argument in required string label Label of the function.
@fortran argument in required int grid_dims Dimensions of the FFT grid.
@fortran argument in required int local_box_origin Coordinates of the local box origin for each MPI rank
@fortran argument in required int local_box_size Dimensions of the local box for each MPI rank.
@fortran argument in required int fcomm Fortran communicator used to partition FFT grid into local boxes.
@fortran argument out required double values Values of the function (local buffer for each MPI rank).
@fortran argument in optional bool transform_to_rg If true, transform function to regular grid before fetching the values.
@fortran end */
void sirius_get_rg_values(void* const* handler__,
char const* label__,
int const* grid_dims__,
int const* local_box_origin__,
int const* local_box_size__,
int const* fcomm__,
double* values__,
bool const* transform_to_rg__)
{
PROFILE("sirius_api::sirius_get_rg_values");

GET_GS(handler__);

std::string label(label__);

for (int x: {0, 1, 2}) {
if (grid_dims__[x] != gs.ctx().fft().size(x)) {
TERMINATE("wrong FFT grid size");
}
}

std::map<std::string, sirius::Smooth_periodic_function<double>*> func = {
{"rho", &gs.density().rho()},
{"magz", &gs.density().magnetization(0)},
{"magx", &gs.density().magnetization(1)},
{"magy", &gs.density().magnetization(2)},
{"veff", &gs.potential().effective_potential()},
{"bz", &gs.potential().effective_magnetic_field(0)},
{"bx", &gs.potential().effective_magnetic_field(1)},
{"by", &gs.potential().effective_magnetic_field(2)},
{"vxc", &gs.potential().xc_potential()},
};

try {
auto& f = func.at(label);

auto& comm = Communicator::map_fcomm(*fcomm__);

if (transform_to_rg__ && *transform_to_rg__) {
f->fft_transform(1);
}

auto& spl_z = f->fft().spl_z();

mdarray<int, 2> local_box_size(const_cast<int*>(local_box_size__), 3, comm.size());
mdarray<int, 2> local_box_origin(const_cast<int*>(local_box_origin__), 3, comm.size());

for (int rank = 0; rank < f->fft().comm().size(); rank++) {
/* slab of FFT grid for a given rank */
mdarray<double, 3> buf(f->fft().size(0), f->fft().size(1), spl_z.local_size(rank));
if (rank == f->fft().comm().rank()) {
std::copy(&f->f_rg(0), &f->f_rg(0) + f->fft().local_size(), &buf[0]);
}
f->fft().comm().bcast(&buf[0], static_cast<int>(buf.size()), rank);

/* ranks on the F90 side */
int r = comm.rank();

/* dimensions of this rank's local box */
int nx = local_box_size(0, r);
int ny = local_box_size(1, r);
int nz = local_box_size(2, r);
mdarray<double, 3> values(values__, nx, ny, nz);

for (int iz = 0; iz < nz; iz++) {
/* global z coordinate inside FFT box */
int z = local_box_origin(2, r) + iz - 1; /* Fortran counts from 1 */
if (z >= spl_z.global_offset(rank) && z < spl_z.global_offset(rank) + spl_z.local_size(rank)) {
/* make z local for SIRIUS FFT partitioning */
z -= spl_z.global_offset(rank);
for (int iy = 0; iy < ny; iy++) {
/* global y coordinate inside FFT box */
int y = local_box_origin(1, r) + iy - 1; /* Fortran counts from 1 */
for (int ix = 0; ix < nx; ix++) {
/* global x coordinate inside FFT box */
int x = local_box_origin(0, r) + ix - 1; /* Fortran counts from 1 */
values(ix, iy, iz) = buf(x, y, z);
}
}
}
}
} /* loop over ranks */
} catch(...) {
TERMINATE("wrong label");
}
}


/* @fortran begin function void sirius_get_total_magnetization Get the total magnetization of the system.
@fortran argument in required void* handler DFT ground state handler.
@fortran argument out required double mag 3D magnetization vector (x,y,z components).
Expand Down

0 comments on commit 5d8320d

Please sign in to comment.