OpenMC/include/openmc/bank_io.h
2025-12-11 22:01:10 -06:00

105 lines
2.7 KiB
C++

#ifndef OPENMC_BANK_IO_H
#define OPENMC_BANK_IO_H
#include "hdf5.h"
#include "openmc/message_passing.h"
#include "openmc/span.h"
#include "openmc/vector.h"
#include <algorithm>
#ifdef OPENMC_MPI
#include <mpi.h>
#endif
namespace openmc {
template<typename SiteType>
void write_bank_dataset(
const char* dataset_name, hid_t group_id, span<SiteType> bank,
const vector<int64_t>& bank_index, hid_t membanktype, hid_t filebanktype
#ifdef OPENMC_MPI
,
MPI_Datatype mpi_dtype
#endif
)
{
int64_t dims_size = bank_index.back();
int64_t count_size = bank_index[mpi::rank + 1] - bank_index[mpi::rank];
#ifdef PHDF5
hsize_t dims[] {static_cast<hsize_t>(dims_size)};
hid_t dspace = H5Screate_simple(1, dims, nullptr);
hid_t dset = H5Dcreate(group_id, dataset_name, filebanktype, dspace,
H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
hsize_t count[] {static_cast<hsize_t>(count_size)};
hid_t memspace = H5Screate_simple(1, count, nullptr);
hsize_t start[] {static_cast<hsize_t>(bank_index[mpi::rank])};
H5Sselect_hyperslab(dspace, H5S_SELECT_SET, start, nullptr, count, nullptr);
hid_t plist = H5Pcreate(H5P_DATASET_XFER);
H5Pset_dxpl_mpio(plist, H5FD_MPIO_COLLECTIVE);
H5Dwrite(dset, membanktype, memspace, dspace, plist, bank.data());
H5Sclose(dspace);
H5Sclose(memspace);
H5Dclose(dset);
H5Pclose(plist);
#else
if (mpi::master) {
hsize_t dims[] {static_cast<hsize_t>(dims_size)};
hid_t dspace = H5Screate_simple(1, dims, nullptr);
hid_t dset = H5Dcreate(group_id, dataset_name, filebanktype, dspace,
H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
#ifdef OPENMC_MPI
vector<SiteType> temp_bank {bank.begin(), bank.end()};
#endif
for (int i = 0; i < mpi::n_procs; ++i) {
hsize_t count[] {static_cast<hsize_t>(bank_index[i + 1] - bank_index[i])};
hid_t memspace = H5Screate_simple(1, count, nullptr);
#ifdef OPENMC_MPI
if (i > 0) {
MPI_Recv(bank.data(), count[0], mpi_dtype, i, i, mpi::intracomm,
MPI_STATUS_IGNORE);
}
#endif
hid_t dspace_rank = H5Dget_space(dset);
hsize_t start[] {static_cast<hsize_t>(bank_index[i])};
H5Sselect_hyperslab(
dspace_rank, H5S_SELECT_SET, start, nullptr, count, nullptr);
H5Dwrite(
dset, membanktype, memspace, dspace_rank, H5P_DEFAULT, bank.data());
H5Sclose(memspace);
H5Sclose(dspace_rank);
}
H5Dclose(dset);
#ifdef OPENMC_MPI
std::copy(temp_bank.begin(), temp_bank.end(), bank.begin());
#endif
}
#ifdef OPENMC_MPI
else {
if (!bank.empty()) {
MPI_Send(
bank.data(), bank.size(), mpi_dtype, 0, mpi::rank, mpi::intracomm);
}
}
#endif
#endif
}
} // namespace openmc
#endif // OPENMC_BANK_IO_H