diff --git a/include/bout/adios_object.hxx b/include/bout/adios_object.hxx index 63a2d4f85f..ee7e82e45a 100755 --- a/include/bout/adios_object.hxx +++ b/include/bout/adios_object.hxx @@ -16,11 +16,23 @@ #if BOUT_HAS_ADIOS2 +#include "bout/assert.hxx" +#include "bout/boutcomm.hxx" #include "bout/boutexception.hxx" +#include "bout/field2d.hxx" +#include "bout/field3d.hxx" +#include "bout/fieldperp.hxx" +#include "bout/globals.hxx" +#include "bout/mesh.hxx" +#include "bout/utils.hxx" #include #include #include +#include +#include +#include +#include namespace bout { @@ -29,21 +41,17 @@ void ADIOSInit(const std::string configFile, MPI_Comm comm); void ADIOSFinalize(); using ADIOSPtr = std::shared_ptr; -using EnginePtr = std::shared_ptr; -using IOPtr = std::shared_ptr; ADIOSPtr GetADIOSPtr(); -IOPtr GetIOPtr(const std::string IOName); class ADIOSStream { public: adios2::IO io; - adios2::Variable vTime; - adios2::Variable vStep; int adiosStep = 0; /** create or return the ADIOSStream based on the target file name */ - static ADIOSStream& ADIOSGetStream(const std::string& fname, adios2::Mode mode); + static ADIOSStream& ADIOSGetStream(const std::string& fname, adios2::Mode mode, + const std::string& engineType = "BP5"); ~ADIOSStream(); @@ -56,6 +64,91 @@ public: return v; } + template + void Get(const std::string& varname, T& value, adios2::Mode mode = adios2::Mode::Sync) { + auto variable = io.InquireVariable(varname); + ASSERT1(variable); + ASSERT1(variable.ShapeID() == adios2::ShapeID::GlobalValue); + engine().Get(variable, &value, mode); + } + + template + void Get(const std::string& varname, Array& value, + adios2::Mode mode = adios2::Mode::Sync) { + GetArrayLike(varname, value, mode); + } + + template + void Get(const std::string& varname, Matrix& value, + adios2::Mode mode = adios2::Mode::Sync) { + GetArrayLike(varname, value, mode); + } + + template + void Get(const std::string& varname, Tensor& value, + adios2::Mode mode = adios2::Mode::Sync) { + GetArrayLike(varname, value, mode); + } + + void Get(const std::string& varname, Field2D& value, + adios2::Mode mode = adios2::Mode::Sync) { + value.allocate(); + GetField(varname, {"x", "y"}, *value.getMesh(), &value(0, 0), mode); + } + + void Get(const std::string& varname, Field3D& value, + adios2::Mode mode = adios2::Mode::Sync) { + value.allocate(); + GetField(varname, {"x", "y", "z"}, *value.getMesh(), &value(0, 0, 0), mode); + } + + void Get(const std::string& varname, FieldPerp& value, + adios2::Mode mode = adios2::Mode::Sync) { + value.allocate(); + GetField(varname, {"x", "z"}, *value.getMesh(), &value(0, 0), mode); + } + + template + void Put(const std::string& varname, T value, adios2::Mode mode = adios2::Mode::Sync) { + if (BoutComm::rank() != 0) { + return; + } + engine().Put(GetValueVariable(varname), value, mode); + } + + template + void Put(const std::string& varname, const Array& value, + adios2::Mode mode = adios2::Mode::Sync) { + PutArrayLike(varname, value, mode); + } + + template + void Put(const std::string& varname, const Matrix& value, + adios2::Mode mode = adios2::Mode::Sync) { + PutArrayLike(varname, value, mode); + } + + template + void Put(const std::string& varname, const Tensor& value, + adios2::Mode mode = adios2::Mode::Sync) { + PutArrayLike(varname, value, mode); + } + + void Put(const std::string& varname, const Field2D& value, + adios2::Mode mode = adios2::Mode::Sync) { + PutField(varname, {"x", "y"}, *value.getMesh(), &value(0, 0), mode); + } + + void Put(const std::string& varname, const Field3D& value, + adios2::Mode mode = adios2::Mode::Sync) { + PutField(varname, {"x", "y", "z"}, *value.getMesh(), &value(0, 0, 0), mode); + } + + void Put(const std::string& varname, const FieldPerp& value, + adios2::Mode mode = adios2::Mode::Sync) { + PutField(varname, {"x", "z"}, *value.getMesh(), &value(0, 0), mode); + } + template adios2::Variable GetArrayVariable(const std::string& varname, const adios2::Dims& shape, @@ -78,7 +171,7 @@ public: if (not engine_) { engine_ = io.Open(fname, file_mode); if (not engine_) { - throw BoutException("Could not open ADIOS file '{:s}' for writing", fname); + throw BoutException("Could not open ADIOS file '{:s}'", fname); } } return engine_; @@ -101,25 +194,184 @@ public: void finish() { if (engine_) { - engine().EndStep(); + if (isInStep) { + engine().EndStep(); + isInStep = false; + } engine().Close(); + engine_ = adios2::Engine(); } } private: - ADIOSStream(const std::string& fname, adios2::Mode mode) - : fname(fname), file_mode(mode) { - - ADIOSPtr adiosp = GetADIOSPtr(); - std::string ioname = "write_" + fname; - try { - io = adiosp->AtIO(ioname); - } catch (const std::invalid_argument& e) { - io = adiosp->DeclareIO(ioname); - io.SetEngine("BP5"); + ADIOSStream(const std::string& fname, adios2::Mode mode, const std::string& engineType); + + struct FieldSelection { + adios2::Dims shape; + adios2::Dims start; + adios2::Dims count; + adios2::Dims mem_start; + adios2::Dims mem_count; + + auto selection() const { return adios2::Box{start, count}; } + auto memorySelection() const { + return adios2::Box{mem_start, mem_count}; } }; + void GetField(const std::string& varname, const std::vector& dim_names, + const Mesh& mesh, BoutReal* data, adios2::Mode mode) { + auto variable = io.InquireVariable(varname); + ASSERT1(variable); + ASSERT1(variable.ShapeID() == adios2::ShapeID::GlobalArray); + + auto selection = makeFieldSelection(dim_names, mesh); + variable.SetSelection(selection.selection()); + variable.SetMemorySelection(selection.memorySelection()); + engine().Get(variable, data, mode); + } + + void PutField(const std::string& varname, const std::vector& dim_names, + const Mesh& mesh, const BoutReal* data, adios2::Mode mode) { + auto selection = makeFieldSelection(dim_names, mesh); + auto variable = + GetArrayVariable(varname, selection.shape, dim_names, BoutComm::rank()); + variable.SetSelection(selection.selection()); + variable.SetMemorySelection(selection.memorySelection()); + engine().Put(variable, data, mode); + } + + FieldSelection makeFieldSelection(const std::vector& dim_names, + const Mesh& mesh) const { + ASSERT1(!dim_names.empty()); + ASSERT1(dim_names.size() <= 3); + ASSERT1(dim_names[0] == "x"); + + FieldSelection selection; + selection.shape.push_back(static_cast(mesh.GlobalNx)); + selection.start.push_back(static_cast(mesh.MapGlobalX)); + selection.count.push_back(static_cast(mesh.MapCountX)); + selection.mem_start.push_back(static_cast(mesh.MapLocalX)); + selection.mem_count.push_back(static_cast(mesh.LocalNx)); + + if (dim_names.size() > 1) { + if (dim_names[1] == "y") { + selection.shape.push_back(static_cast(mesh.GlobalNy)); + selection.start.push_back(static_cast(mesh.MapGlobalY)); + selection.count.push_back(static_cast(mesh.MapCountY)); + selection.mem_start.push_back(static_cast(mesh.MapLocalY)); + selection.mem_count.push_back(static_cast(mesh.LocalNy)); + } else if (dim_names[1] == "z") { + selection.shape.push_back(static_cast(mesh.GlobalNz)); + selection.start.push_back(static_cast(mesh.MapGlobalZ)); + selection.count.push_back(static_cast(mesh.MapCountZ)); + selection.mem_start.push_back(static_cast(mesh.MapLocalZ)); + selection.mem_count.push_back(static_cast(mesh.LocalNz)); + } else { + ASSERT1(false); + } + } + + if (dim_names.size() > 2) { + ASSERT1(dim_names[1] == "y"); + ASSERT1(dim_names[2] == "z"); + selection.shape.push_back(static_cast(mesh.GlobalNz)); + selection.start.push_back(static_cast(mesh.MapGlobalZ)); + selection.count.push_back(static_cast(mesh.MapCountZ)); + selection.mem_start.push_back(static_cast(mesh.MapLocalZ)); + selection.mem_count.push_back(static_cast(mesh.LocalNz)); + } + + return selection; + } + + template + void GetArrayLike(const std::string& varname, Container& value, adios2::Mode mode) { + using T = typename Container::data_type; + auto variable = io.InquireVariable(varname); + ASSERT1(variable); + ASSERT1(variable.ShapeID() == adios2::ShapeID::GlobalArray); + + const auto shape = variable.Shape(); + auto dims_attr = io.InquireAttribute(varname + "/__xarray_dimensions__"); + auto read_shape = shape; + + if (dims_attr) { + const auto dim_names = dims_attr.Data(); + if (!dim_names.empty() && dim_names[0] == "rank") { + ASSERT1(!shape.empty()); + ASSERT1(static_cast(BoutComm::rank()) < shape[0]); + + adios2::Dims start{static_cast(BoutComm::rank())}; + adios2::Dims count{1}; + for (std::size_t i = 1; i < shape.size(); i++) { + start.push_back(0); + count.push_back(shape[i]); + } + + variable.SetSelection(adios2::Box{start, count}); + variable.SetMemorySelection(adios2::Box{start, count}); + read_shape = adios2::Dims(shape.begin() + 1, shape.end()); + } else if (!dim_names.empty() && dim_names[0] == "x") { + ASSERT1(globals::mesh); + auto selection = makeFieldSelection(dim_names, *globals::mesh); + variable.SetSelection(selection.selection()); + variable.SetMemorySelection(selection.memorySelection()); + read_shape = selection.mem_count; + } + } + + constexpr auto ndims = std::tuple_size_v; + ASSERT1(read_shape.size() == ndims); + + resizeForShape(value, read_shape, std::make_index_sequence{}); + engine().Get(variable, value.begin(), mode); + } + + template + void PutArrayLike(const std::string& varname, const Container& value, + adios2::Mode mode) { + using T = typename Container::data_type; + auto var = GetArrayVariable(varname, makeShape(BoutComm::size(), value), + makeDimNames(value), BoutComm::rank()); + var.SetSelection(adios2::Box{makeStart(value), makeShape(1, value)}); + engine().Put(var, value.begin(), mode); + } + + template + adios2::Dims makeShape(std::size_t first, const Container& value) const { + return std::apply( + [first](auto... sizes) { + return adios2::Dims{first, static_cast(sizes)...}; + }, + value.shape()); + } + + template + adios2::Dims makeStart(const Container& value) const { + constexpr auto ndims = std::tuple_size_v; + adios2::Dims start(ndims + 1, 0); + start[0] = static_cast(BoutComm::rank()); + return start; + } + + template + std::vector makeDimNames(const Container& value) const { + constexpr auto ndims = std::tuple_size_v; + std::vector dim_names{"rank"}; + dim_names.reserve(ndims + 1); + for (std::size_t i = 0; i < ndims; i++) { + dim_names.push_back("dim_" + std::to_string(i)); + } + return dim_names; + } + + template + void resizeForShape(Container& value, const adios2::Dims& shape, + std::index_sequence /*indices*/) { + value.reallocate(static_cast(shape[I])...); + } + std::string fname; adios2::Mode file_mode; adios2::Engine engine_; @@ -128,10 +380,6 @@ private: bool isInStep = false; }; -/** Set user parameters for an IO group */ -void ADIOSSetParameters(const std::string& input, char delimKeyValue, char delimItem, - adios2::IO& io); - } // namespace bout #endif //BOUT_HAS_ADIOS2 diff --git a/src/sys/adios_object.cxx b/src/sys/adios_object.cxx index 1b3e0a12e2..596be66b93 100644 --- a/src/sys/adios_object.cxx +++ b/src/sys/adios_object.cxx @@ -7,6 +7,7 @@ #include +#include #include namespace bout { @@ -14,6 +15,30 @@ namespace bout { static ADIOSPtr adios = nullptr; static std::unordered_map adiosStreams; +namespace { +std::string GetIOName(const std::string& fname, adios2::Mode mode) { + switch (mode) { + case adios2::Mode::Read: + return "read_" + fname; + case adios2::Mode::ReadRandomAccess: + return "read_random_access_" + fname; + case adios2::Mode::Append: + return "append_" + fname; + case adios2::Mode::Write: + return "write_" + fname; + default: + return fname; + } +} + +adios2::IO GetIO(const std::string& fname, adios2::Mode mode, + const std::string& engineType) { + auto io = GetADIOSPtr()->DeclareIO(GetIOName(fname, mode)); + io.SetEngine(engineType); + return io; +} +} // namespace + void ADIOSInit(MPI_Comm comm) { adios = std::make_shared(comm); } void ADIOSInit(const std::string configFile, MPI_Comm comm) { @@ -37,16 +62,6 @@ ADIOSPtr GetADIOSPtr() { return adios; } -IOPtr GetIOPtr(const std::string IOName) { - auto adios = GetADIOSPtr(); - IOPtr io = nullptr; - try { - io = std::make_shared(adios->AtIO(IOName)); - } catch (std::invalid_argument& e) { - } - return io; -} - ADIOSStream::~ADIOSStream() { if (engine_) { if (isInStep) { @@ -57,42 +72,19 @@ ADIOSStream::~ADIOSStream() { } } -ADIOSStream& ADIOSStream::ADIOSGetStream(const std::string& fname, adios2::Mode mode) { - auto it = adiosStreams.find(fname); +ADIOSStream::ADIOSStream(const std::string& fname, adios2::Mode mode, + const std::string& engineType) + : io(GetIO(fname, mode, engineType)), fname(fname), file_mode(mode) {} + +ADIOSStream& ADIOSStream::ADIOSGetStream(const std::string& fname, adios2::Mode mode, + const std::string& engineType) { + const auto key = GetIOName(fname, mode); + auto it = adiosStreams.find(key); if (it == adiosStreams.end()) { - it = adiosStreams.emplace(fname, ADIOSStream(fname, mode)).first; + it = adiosStreams.emplace(key, ADIOSStream(fname, mode, engineType)).first; } return it->second; } -void ADIOSSetParameters(const std::string& input, char delimKeyValue, char delimItem, - adios2::IO& io) { - auto lf_Trim = [](std::string& input) { - input.erase(0, input.find_first_not_of(" \n\r\t")); // prefixing spaces - input.erase(input.find_last_not_of(" \n\r\t") + 1); // suffixing spaces - }; - - std::istringstream inputSS(input); - std::string parameter; - while (std::getline(inputSS, parameter, delimItem)) { - const size_t position = parameter.find(delimKeyValue); - if (position == std::string::npos) { - throw BoutException("ADIOSSetParameters(): wrong format for IO parameter " - + parameter + ", format must be key" + delimKeyValue - + "value for each entry"); - } - - std::string key = parameter.substr(0, position); - lf_Trim(key); - std::string value = parameter.substr(position + 1); - lf_Trim(value); - if (value.length() == 0) { - throw BoutException("ADIOS2SetParameters: empty value in IO parameter " + parameter - + ", format must be key" + delimKeyValue + "value"); - } - io.SetParameter(key, value); - } -} - } // namespace bout #endif //BOUT_HAS_ADIOS2 diff --git a/src/sys/options/options_adios.cxx b/src/sys/options/options_adios.cxx index 844284a400..4abf55bcd5 100644 --- a/src/sys/options/options_adios.cxx +++ b/src/sys/options/options_adios.cxx @@ -10,157 +10,41 @@ #include "bout/bout_types.hxx" #include "bout/boutexception.hxx" #include "bout/field2d.hxx" -#include "bout/globals.hxx" -#include "bout/mesh.hxx" #include "bout/options.hxx" #include "bout/options_io.hxx" #include "bout/output.hxx" #include "bout/sys/timer.hxx" #include "bout/sys/variant.hxx" -#include "bout/traits.hxx" #include "bout/utils.hxx" #include // IWYU pragma: keep #include #include -#include #include #include #include -#include #include #include -#include -#include namespace { -auto to_int_dims(const adios2::Dims& dims) { - std::vector int_dims; - int_dims.reserve(dims.size()); - std::transform(dims.begin(), dims.end(), std::back_inserter(int_dims), - [](auto dim) { return static_cast(dim); }); - return int_dims; -} - -// Helper class to construct Adios hyperslices -struct Selection { - // Offset of this processor's data into the global array - adios2::Dims start; - // The size of the mapped region - adios2::Dims count; - // Where the actual data starts in data pointer (to exclude ghost cells) - adios2::Dims mem_start; - // The actual size of data pointer in memory (including ghost cells) - adios2::Dims mem_count; - // Global shape, including boundaries but not guard cells - adios2::Dims shape; - // Shape of the local variable to read into - std::vector dims; - - // Distributed Field/Array/Matrix/Tensor - bool should_set_selection{false}; - - Selection(const std::vector& dim_names, const std::vector& dim_sizes, - const Mesh& mesh) { - const auto ndims = dim_names.size(); - const bool dim0_is_rank = ndims > 0 ? (dim_names[0] == "rank") : false; - const bool dim0_is_x = ndims > 0 ? (dim_names[0] == "x") : false; - const bool dim1_is_y = ndims > 1 ? (dim_names[1] == "y") : false; - const bool dim1_is_z = ndims > 1 ? (dim_names[1] == "z") : false; - const bool dim2_is_z = ndims > 2 ? (dim_names[2] == "z") : false; - - should_set_selection = dim0_is_rank or (ndims == 2 and dim0_is_x and dim1_is_y) - or (ndims == 2 and dim0_is_x and dim1_is_z) - or (ndims == 3 and dim0_is_x and dim1_is_y and dim2_is_z); - - if (dim0_is_rank) { - const auto ndim_sizes = dim_sizes.size(); - ASSERT3(ndim_sizes > 1); - - // This is a distributed array, so the local variable is going - // to be shape (dim_sizes[1]...) (that is, drop the rank) - dims.push_back(dim_sizes[1]); - // but we tell adios to read our rank's bit with the full ndims - start = {static_cast(BoutComm::rank()), 0}; - count = {std::size_t{1}, static_cast(dim_sizes[0])}; - - if (ndim_sizes > 2) { - dims.push_back(dim_sizes[2]); - start.push_back(0); - count.push_back(dim_sizes[1]); - } - if (ndim_sizes > 3) { - dims.push_back(dim_sizes[3]); - start.push_back(0); - count.push_back(dim_sizes[2]); - } - mem_count = count; - mem_start = start; - return; - } - - if (ndims > 0) { - dims.push_back(dim0_is_x ? mesh.LocalNx : dim_sizes[0]); - } - if (ndims > 1) { - if (dim1_is_y) { - dims.push_back(mesh.LocalNy); - } else if (dim1_is_z) { - dims.push_back(mesh.LocalNz); - } else { - dims.push_back(dim_sizes[1]); - } - } - if (ndims > 2) { - dims.push_back(dim2_is_z ? mesh.LocalNz : dim_sizes[2]); - } - - shape.push_back(static_cast(mesh.GlobalNx)); - start.push_back(static_cast(mesh.MapGlobalX)); - count.push_back(static_cast(mesh.MapCountX)); - mem_start.push_back(static_cast(mesh.MapLocalX)); - mem_count.push_back(static_cast(mesh.LocalNx)); - - if (dim1_is_y) { - shape.push_back(static_cast(mesh.GlobalNy)); - start.push_back(static_cast(mesh.MapGlobalY)); - count.push_back(static_cast(mesh.MapCountY)); - mem_start.push_back(static_cast(mesh.MapLocalY)); - mem_count.push_back(static_cast(mesh.LocalNy)); - } else if (dim1_is_z) { - shape.push_back(static_cast(mesh.GlobalNz)); - start.push_back(static_cast(mesh.MapGlobalZ)); - count.push_back(static_cast(mesh.MapCountZ)); - mem_start.push_back(static_cast(mesh.MapLocalZ)); - mem_count.push_back(static_cast(mesh.LocalNz)); - } - - if (dim2_is_z) { - shape.push_back(static_cast(mesh.GlobalNz)); - start.push_back(static_cast(mesh.MapGlobalZ)); - count.push_back(static_cast(mesh.MapCountZ)); - mem_start.push_back(static_cast(mesh.MapLocalZ)); - mem_count.push_back(static_cast(mesh.LocalNz)); +std::size_t getLocalNDims(adios2::IO& io, const std::string& name, + const adios2::Dims& shape) { + auto dims_attr = io.InquireAttribute(name + "/__xarray_dimensions__"); + if (dims_attr) { + const auto dim_names = dims_attr.Data(); + if (!dim_names.empty() && dim_names[0] == "rank") { + ASSERT1(!shape.empty()); + return shape.size() - 1; } } + return shape.size(); +} - auto selection() const { return adios2::Box{start, count}; } - auto memorySelection() const { return adios2::Box{mem_start, mem_count}; } -}; - -template