diff --git a/src/include/fields3d.hxx b/src/include/fields3d.hxx index 69e3b7d15..d015c3fa0 100644 --- a/src/include/fields3d.hxx +++ b/src/include/fields3d.hxx @@ -353,6 +353,7 @@ struct Mfields const Grid_t& grid() const { return *grid_; } auto gt() { return Base::storage().view(); } + auto gt() const { return Base::storage().view(); } template void Foreach_3d(int l, int r, FUNC&& F) const diff --git a/src/include/fields3d.inl b/src/include/fields3d.inl index 42e32514b..4c90a52c5 100644 --- a/src/include/fields3d.inl +++ b/src/include/fields3d.inl @@ -18,19 +18,13 @@ public: void put(kg::io::Engine& writer, const Mfields& mflds, const kg::io::Mode launch = kg::io::Mode::NonBlocking) { - writer.put("ib", mflds.box().ib(), launch); - writer.put("im", mflds.box().im(), launch); - - auto n_comps = mflds.n_comps(); - auto shape = makeDims(n_comps, mflds.gdims()); + const Grid_t& grid = mflds.grid(); + std::vector patchOffsets; for (int p = 0; p < mflds.n_patches(); p++) { - auto start = makeDims(0, mflds.patchOffset(p)); - auto count = makeDims(n_comps, mflds.ldims()); - auto ib = makeDims(0, -mflds.box().ib()); - auto im = makeDims(n_comps, mflds.box().im()); - writer.putVariable(const_cast(mflds)[p].storage().data(), - launch, shape, {start, count}, {ib, im}); // FIXME cast + patchOffsets.push_back(mflds.patchOffset(p)); } + + write_4d(writer, grid.ldims, grid.domain.gdims, patchOffsets, mflds.gt()); } void get(kg::io::Engine& reader, Mfields& mflds, diff --git a/src/include/io_common.h b/src/include/io_common.h index 1040f79c4..872620ea8 100644 --- a/src/include/io_common.h +++ b/src/include/io_common.h @@ -13,4 +13,71 @@ inline kg::io::Dims makeDims(int m, const Int3& dims) static_cast(dims[0])}; } +inline kg::io::Dims makeDims(const Int3& dims) +{ + return kg::io::Dims{static_cast(dims[2]), + static_cast(dims[1]), + static_cast(dims[0])}; +} + +template +inline void write_4d(kg::io::Engine& file, const Int3& ldims, const Int3& gdims, + const std::vector& patch_off, const E& h_expr) +{ + auto launch = kg::io::Mode::Blocking; + + int n_comps = h_expr.shape(3); + int n_patches = h_expr.shape(4); + Int3 im = {h_expr.shape(0), h_expr.shape(1), h_expr.shape(2)}; + Int3 ib = -(im - ldims) / 2; + std::cout << "ib " << ib[0] << " " << ib[1] << " " << ib[2] << "\n"; + std::cout << "im " << im[0] << " " << im[1] << " " << im[2] << "\n"; + file.put("ib", ib, launch); + file.put("im", im, launch); + + auto shape = makeDims(n_comps, gdims); + for (int p = 0; p < n_patches; p++) { + auto start = makeDims(0, patch_off[p]); + auto count = makeDims(n_comps, ldims); + auto _ib = makeDims(0, -ib); + auto _im = makeDims(n_comps, im); + file.putVariable(&h_expr(0, 0, 0, 0, p), launch, shape, {start, count}, + {_ib, _im}); + } + file.performPuts(); +} + +template +inline void write_3d(kg::io::Engine& file, const Int3& ldims, const Int3& gdims, + const std::vector& patch_off, const E& h_expr, + const std::string& name) +{ + // FIXME, should not pass name, but have prefixes set first, but gotta resolve + // separator inconsistency first + auto launch = kg::io::Mode::Blocking; + + int n_patches = h_expr.shape(3); + Int3 im = {h_expr.shape(0), h_expr.shape(1), h_expr.shape(2)}; + Int3 ib = -(im - ldims) / 2; + + file.prefixes_.push_back(name); + file.put("ib", ib, launch); + file.put("im", im, launch); + + auto shape = makeDims(gdims); + for (int p = 0; p < n_patches; p++) { + auto start = makeDims(patch_off[p]); + auto count = makeDims(ldims); + auto _ib = makeDims(-ib); + auto _im = makeDims(im); + file.putVariable(&h_expr(0, 0, 0, p), launch, shape, {start, count}, + {_ib, _im}); + } + file.prefixes_.pop_back(); + // FIXME, it'd be better to write within the prefix, but xarray-adios2 + // expects a "/" separator instead of "::" + file.put(std::string(name) + "/dimensions", std::string("z y x")); + file.performPuts(); +} + } // namespace diff --git a/src/include/psc_fields_cuda.inl b/src/include/psc_fields_cuda.inl index 07d77bf91..56cf48d54 100644 --- a/src/include/psc_fields_cuda.inl +++ b/src/include/psc_fields_cuda.inl @@ -23,27 +23,17 @@ public: const kg::io::Mode launch = kg::io::Mode::NonBlocking) { const auto& grid = mflds_cuda.grid(); - auto n_comps = mflds_cuda.n_comps(); writer.put("ibn", mflds_cuda.ibn(), launch); auto&& h_mflds = gt::host_mirror(mflds_cuda.storage()); gt::copy(mflds_cuda.storage(), h_mflds); - auto shape = makeDims(n_comps, grid.domain.gdims); + std::vector patchOffsets; for (int p = 0; p < mflds_cuda.n_patches(); p++) { - auto start = makeDims(0, grid.patches[p].off); - auto count = makeDims(n_comps, grid.ldims); - auto ib = makeDims(0, mflds_cuda.ibn()); - auto shp = mflds_cuda.storage().shape(); - auto im = makeDims(n_comps, {shp[0], shp[1], shp[2]}); - writer.putVariable(&h_mflds(0, 0, 0, 0, p), launch, shape, {start, count}, - {ib, im}); + patchOffsets.push_back(grid.patches[p].off); } - - // host mirror will go away as this function returns, so need - // to write - writer.performPuts(); + write_4d(writer, grid.ldims, grid.domain.gdims, patchOffsets, h_mflds); } void get(kg::io::Engine& reader, Mfields& mflds_cuda, diff --git a/src/include/writer_adios2.hxx b/src/include/writer_adios2.hxx index fec6ad15b..fd36ebfbe 100644 --- a/src/include/writer_adios2.hxx +++ b/src/include/writer_adios2.hxx @@ -71,24 +71,11 @@ public: void begin_step(const Grid_t& grid) { - int step = grid.timestep(); - double time = grid.time(); - - int len = dir_.size() + pfx_.size() + 20; - char filename[len]; - snprintf(filename, len, "%s/%s.%09d.bp", dir_.c_str(), pfx_.c_str(), step); - file_ = io_.open(filename, kg::io::Mode::Write, comm_, pfx_); - file_.beginStep(kg::io::StepMode::Append); - file_.put("step", step); - file_.put("time", time); - file_.put("length", grid.domain.length); - file_.put("corner", grid.domain.corner); - file_.performPuts(); + _begin_step(file_, grid.timestep(), grid.time(), grid.domain.length, + grid.domain.corner, grid.domain.gdims); } - void end_step() { file_.close(); } - - void set_subset(const Grid_t& grid, Int3 rn, Int3 rx) {} + void end_step() { _end_step(file_); } template void write(const E& expr, const Grid_t& grid, const std::string& name, @@ -105,31 +92,52 @@ public: auto&& evaluated = gt::eval(const_cast(expr)); auto&& h_expr = gt::host_mirror(evaluated); gt::copy(evaluated, h_expr); - Mfields> h_mflds(grid, h_expr.shape(3), {}); - h_mflds.gt() = h_expr; prof_stop(pr_eval); prof_start(pr_write); - file_.put(name, h_mflds); - file_.performPuts(); + std::vector patch_off(grid.n_patches()); + for (int p = 0; p < grid.n_patches(); p++) { + patch_off[p] = grid.patches[p].off; + } + write(file_, grid.ldims, grid.domain.gdims, patch_off, h_expr, name, + comp_names); prof_stop(pr_write); } + template + static void write(kg::io::Engine& file, const Int3& ldims, const Int3& gdims, + const std::vector& patch_off, const E& h_expr, + const std::string& name, + const std::vector& comp_names) + { +#if 0 + file.prefixes_.push_back(name); + write_4d(file, ldims, gdims, patch_off, h_expr); + file.prefixes_.pop_back(); +#else + int n_comps = h_expr.shape(3); + assert(n_comps == comp_names.size()); + + for (int m = 0; m < n_comps; m++) { + write_3d(file, ldims, gdims, patch_off, + h_expr.view(_all, _all, _all, m, _all), comp_names[m]); + } +#endif + } + template void write_step(const Grid_t& grid, const Int3& rn, const Int3& rx, const E& expr, const std::string& name, const std::vector& comp_names) { - static int pr, pr_copy, pr_wait, pr_write, pr_thread, pr_lock, pr_adios2; + static int pr, pr_copy, pr_write, pr_thread, pr_lock; if (!pr) { pr = prof_register("write_step", 1., 0, 0); pr_copy = prof_register("ws copy", 1., 0, 0); - pr_wait = prof_register("ws wait", 1., 0, 0); pr_write = prof_register("ws write", 1., 0, 0); pr_thread = prof_register("ws thread", 1., 0, 0); pr_lock = prof_register("ws lock", 1., 0, 0); - pr_adios2 = prof_register("ws adios2", 1., 0, 0); } prof_start(pr); @@ -145,7 +153,7 @@ public: #ifdef PSC_USE_IO_THREADS prof_start(pr_write); int step = grid.timestep(); - double time = step * grid.dt; + double time = grid.time(); assert(grid.n_patches() == h_expr.shape(4)); Int3 ldims = grid.ldims; @@ -165,51 +173,19 @@ public: prof_start(pr_thread); - int n_comps = h_expr.shape(3); - int n_patches = h_expr.shape(4); - Int3 im = {h_expr.shape(0), h_expr.shape(1), h_expr.shape(2)}; - Int3 ib = {-(im[0] - ldims[0]) / 2, -(im[1] - ldims[1]) / 2, - -(im[2] - ldims[2]) / 2}; - - int len = dir_.size() + pfx_.size() + 20; - char filename[len]; - snprintf(filename, len, "%s/%s.%09d.bp", dir_.c_str(), pfx_.c_str(), - step); { - auto launch = kg::io::Mode::Blocking; - // FIXME not sure how necessary this lock really is, it certainly could // spin for a long time if another thread is writing another file prof_start(pr_lock); // std::lock_guard guard(writer_mutex); prof_stop(pr_lock); - auto file = io_.open(filename, kg::io::Mode::Write, comm_, pfx_); - file.beginStep(kg::io::StepMode::Append); - file.put("step", step); - file.put("time", time); - file.put("length", length); - file.put("corner", corner); - - prof_start(pr_adios2); - file.put("ib", ib, launch); - file.put("im", im, launch); - - file.prefixes_.push_back(name); - auto shape = makeDims(n_comps, gdims); - for (int p = 0; p < n_patches; p++) { - auto start = makeDims(0, patch_off[p]); - auto count = makeDims(n_comps, ldims); - auto _ib = makeDims(0, -ib); - auto _im = makeDims(n_comps, im); - file.putVariable(&h_expr(ib[0], ib[1], ib[2], 0, p), launch, shape, - {start, count}, {_ib, _im}); // FIXME cast - } - file.prefixes_.pop_back(); - file.performPuts(); - file.endStep(); - file.close(); - prof_stop(pr_adios2); + prof_start(pr_write); + kg::io::Engine file; + _begin_step(file, step, time, length, corner, gdims); + write(file, ldims, gdims, patch_off, h_expr, name, comp_names); + _end_step(file); + prof_stop(pr_write); } prof_stop(pr_thread); @@ -222,7 +198,6 @@ public: #else prof_start(pr_write); begin_step(grid); - set_subset(grid, rn, rx); write(expr, grid, name, comp_names); end_step(); prof_stop(pr_write); @@ -231,6 +206,50 @@ public: } private: + void _begin_step(kg::io::Engine& file, int step, double time, + const Double3& length, const Double3& corner, + const Int3& gdims) + { + int len = dir_.size() + pfx_.size() + 20; + char filename[len]; + snprintf(filename, len, "%s/%s.%09d.bp", dir_.c_str(), pfx_.c_str(), step); + file = io_.open(filename, kg::io::Mode::Write, comm_, pfx_); + + file.beginStep(kg::io::StepMode::Append); + file.put("step", step); + file.put("time", time); + // additionally, write time as a variable (above writes it as an attribute) + // so that xarray-adios2 can use it as a dimension coordinate + file.file_.putVariable("time", time); + file.put("length", length); + file.put("corner", corner); + file.put("step_dimension", std::string("time")); + + if (file.mpi_rank_ == 0) { + Double3 dx = length / Double3(gdims); + for (int d = 0; d < 3; d++) { + auto crd = gt::zeros({gdims[d]}); + for (int i = 0; i < gdims[d]; i++) { + crd(i) = corner[d] + (i + .5) * dx[d]; + } + std::string name = d == 0 ? "x" : d == 1 ? "y" : "z"; + file.prefixes_.push_back(name); + file.putVariable(crd.data(), kg::io::Mode::Blocking, + kg::io::Dims({size_t(gdims[d])}), {}, {}); + file.prefixes_.pop_back(); + file.put(std::string(name) + "/dimensions", name); + } + } + + file_.performPuts(); + } + + static void _end_step(kg::io::Engine& file) + { + file.endStep(); + file.close(); + } + #ifdef PSC_USE_IO_THREADS void thread_func() { diff --git a/src/kg/include/kg/io/File.h b/src/kg/include/kg/io/File.h index b1d81b746..d5361611b 100644 --- a/src/kg/include/kg/io/File.h +++ b/src/kg/include/kg/io/File.h @@ -31,11 +31,16 @@ class File void putVariable(const std::string& name, const T* data, Mode launch, const Dims& shape, const Extents& selection, const Extents& memory_selection); + template + void putVariable(const std::string& name, const T& datum); template void getVariable(const std::string& name, T* data, Mode launch, const Extents& selection, const Extents& memory_selection); + template + void getVariable(const std::string& name, T& datum); + Dims shapeVariable(const std::string& name) const; template diff --git a/src/kg/include/kg/io/File.inl b/src/kg/include/kg/io/File.inl index 873208e8e..8f149555c 100644 --- a/src/kg/include/kg/io/File.inl +++ b/src/kg/include/kg/io/File.inl @@ -37,6 +37,14 @@ inline void File::putVariable(const std::string& name, const T* data, impl_->putVariable(name, dataVar, launch, shape, selection, memory_selection); } +template +inline void File::putVariable(const std::string& name, const T& datum) +{ + assert(impl_); + FileBase::TypeConstPointer datumVar = &datum; + impl_->putVariable(name, datumVar, kg::io::Mode::Blocking, {}, {}, {}); +} + template inline void File::getVariable(const std::string& name, T* data, Mode launch, const Extents& selection, @@ -47,6 +55,16 @@ inline void File::getVariable(const std::string& name, T* data, Mode launch, impl_->getVariable(name, dataVar, launch, selection, memory_selection); } +template +inline void File::getVariable(const std::string& name, T& datum) +{ + assert(impl_); + auto shape = impl_->shapeVariable(name); + assert(shape.size() == 0); + FileBase::TypePointer dataVar = &datum; + impl_->getVariable(name, dataVar, Mode::Blocking, {}, {}); +} + inline Dims File::shapeVariable(const std::string& name) const { assert(impl_); diff --git a/src/kg/include/kg/io/FileAdios2.inl b/src/kg/include/kg/io/FileAdios2.inl index 1926d9678..036719d70 100644 --- a/src/kg/include/kg/io/FileAdios2.inl +++ b/src/kg/include/kg/io/FileAdios2.inl @@ -71,6 +71,10 @@ inline void FileAdios2::putVariable(const std::string& name, const T* data, } if (!selection.start.empty()) { v.SetSelection({selection.start, selection.count}); + } else { + if (!shape.empty() && shape[0] != LocalValueDim) { + v.SetSelection({kg::io::Dims({0}), shape}); + } } if (!memory_selection.start.empty()) { v.SetMemorySelection({memory_selection.start, memory_selection.count}); diff --git a/src/kg/testing/io/TestIOAdios2.cxx b/src/kg/testing/io/TestIOAdios2.cxx index 8ef6c2b33..feafeeaf4 100644 --- a/src/kg/testing/io/TestIOAdios2.cxx +++ b/src/kg/testing/io/TestIOAdios2.cxx @@ -68,6 +68,32 @@ TEST(IOAdios2, FilePutGetVariable) } } +TEST(IOAdios2, FilePutGetVariableScalar) +{ + auto io = kg::io::IOAdios2{}; + { + auto file = io.openFile("test3.bp", kg::io::Mode::Write); + file.beginStep(kg::io::StepMode::Append); + double dbl = 99.; + file.putVariable("dbl_scalar", dbl); + file.performPuts(); + file.endStep(); + } + { + auto file = io.openFile("test3.bp", kg::io::Mode::Read); + file.beginStep(kg::io::StepMode::Read); + + auto shape = file.shapeVariable("dbl_scalar"); + EXPECT_EQ(shape.size(), 0); + double dbl; + file.getVariable("dbl_scalar", dbl); + file.performGets(); + file.endStep(); + + EXPECT_EQ(dbl, 99.); + } +} + TEST(IOAdios2, FilePutGetAttribute) { auto io = kg::io::IOAdios2{}; diff --git a/src/libpsc/tests/test_mfields_io.cxx b/src/libpsc/tests/test_mfields_io.cxx index cd62490c6..0dc1ae182 100644 --- a/src/libpsc/tests/test_mfields_io.cxx +++ b/src/libpsc/tests/test_mfields_io.cxx @@ -40,6 +40,56 @@ static Grid_t make_grid() return Grid_t{domain, bc, kinds, norm, dt, -1, {2, 2, 2}}; } +#ifdef PSC_HAVE_ADIOS2 +using WriterTestTypes = ::testing::Types; +#else +using WriterTestTypes = ::testing::Types; +#endif + +template +class WriterTest : public ::testing::Test +{}; + +TYPED_TEST_SUITE(WriterTest, WriterTestTypes); + +TYPED_TEST(WriterTest, WriterSingle) +{ + using Writer = TypeParam; + auto grid = make_grid(); + auto mflds = MfieldsSingle{grid, NR_FIELDS, {2, 2, 2}}; + + setupFields(mflds, [](int m, double crd[3]) { + return m + crd[0] + 100 * crd[1] + 10000 * crd[2]; + }); + + Writer writer; + writer.open("test1", "."); + writer.write_step(grid, {}, {}, + mflds.gt().view(_s(2, -2), _s(2, -2), _s(2, -2)), "mflds", + {"c0", "c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8"}); + writer.close(); +} + +TYPED_TEST(WriterTest, WriterMultiple) +// test writing multiple fields, with different numbers of components, to the +// same file +{ + using Writer = TypeParam; + auto grid = make_grid(); + auto mflds1 = MfieldsSingle{grid, 1, {2, 2, 2}}; + auto mflds2 = MfieldsSingle{grid, 2, {2, 2, 2}}; + + Writer writer; + writer.open("test2", "."); + writer.begin_step(grid); + writer.write(mflds1.gt().view(_s(2, -2), _s(2, -2), _s(2, -2)), grid, + "mflds1", {"comp"}); + writer.write(mflds2.gt().view(_s(2, -2), _s(2, -2), _s(2, -2)), grid, + "mflds2", {"comp1", "comp2"}); + writer.end_step(); + writer.close(); +} + template class MfieldsTest : public ::testing::Test {};