From 52b229e622a4327bd3f9c4fc679311adb21724eb Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Fri, 8 Mar 2024 11:30:45 -0600 Subject: [PATCH 01/10] Started adding virial contribution support to FixDLExt --- dlext/src/FixDLExt.cpp | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/dlext/src/FixDLExt.cpp b/dlext/src/FixDLExt.cpp index d57f35b..59cff16 100644 --- a/dlext/src/FixDLExt.cpp +++ b/dlext/src/FixDLExt.cpp @@ -32,6 +32,10 @@ FixDLExt::FixDLExt(LAMMPS* lmp, int narg, char** arg) if (atom->map_style != Atom::MAP_ARRAY) error->all(FLERR, "Fix dlext requires to map atoms as arrays"); + // signal that this fix contributes to the global virial + virial_global_flag = 1; + thermo_virial = 1; + kokkosable = has_kokkos_cuda_enabled(lmp); atomKK = dynamic_cast(atom); execution_space = (on_host || !kokkosable) ? kOnHost : kOnDevice; @@ -40,7 +44,17 @@ FixDLExt::FixDLExt(LAMMPS* lmp, int narg, char** arg) } int FixDLExt::setmask() { return FixConst::POST_FORCE; } -void FixDLExt::post_force(int) { callback(update->ntimestep); } +void FixDLExt::post_force(int vflag) +{ + // virial setup + v_init(vflag); + + // invoke callback + callback(update->ntimestep); + + // TODO: put the callback's virial into this fix's member variable virial[6] (see fix.h) + // callback(virial) +} void FixDLExt::set_callback(DLExtCallback& cb) { callback = cb; } void register_FixDLExt(LAMMPS* lmp) From 77781c69ec19f9c4499e21e38ad838ed01f3f27c Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Mon, 11 Mar 2024 15:34:29 -0500 Subject: [PATCH 02/10] Added callback to set virial and exposed it to the python module --- dlext/include/FixDLExt.h | 3 +++ dlext/src/FixDLExt.cpp | 6 ++++-- python/lammps_dlext.cpp | 1 + 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index 8d8f048..b27fc10 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -19,6 +19,7 @@ namespace dlext using TimeStep = bigint; // bigint depends on how LAMMPS was built using DLExtCallback = std::function; +using DLExtSetVirial = std::function; // } // Aliases @@ -39,9 +40,11 @@ class DEFAULT_VISIBILITY FixDLExt : public Fix { int setmask() override; void post_force(int) override; void set_callback(DLExtCallback& cb); + void set_virial_callback(DLExtSetVirial& cb); private: DLExtCallback callback = [](TimeStep) { }; + DLExtSetVirial setVirial = [](double*) { }; }; void register_FixDLExt(LAMMPS* lmp); diff --git a/dlext/src/FixDLExt.cpp b/dlext/src/FixDLExt.cpp index 59cff16..e83b351 100644 --- a/dlext/src/FixDLExt.cpp +++ b/dlext/src/FixDLExt.cpp @@ -52,11 +52,13 @@ void FixDLExt::post_force(int vflag) // invoke callback callback(update->ntimestep); - // TODO: put the callback's virial into this fix's member variable virial[6] (see fix.h) - // callback(virial) + // put the virial from the bias into this fix's member variable virial[6] (see fix.h) + setVirial(virial); } void FixDLExt::set_callback(DLExtCallback& cb) { callback = cb; } +void FixDLExt::set_virial_callback(DLExtSetVirial& cb) { setVirial = cb; } + void register_FixDLExt(LAMMPS* lmp) { auto fix_map = lmp->modify->fix_map; diff --git a/python/lammps_dlext.cpp b/python/lammps_dlext.cpp index 5929212..8ad36bd 100644 --- a/python/lammps_dlext.cpp +++ b/python/lammps_dlext.cpp @@ -49,6 +49,7 @@ void export_FixDLExt(py::module& m) return static_cast(fix); })) .def("set_callback", &FixDLExt::set_callback) + .def("set_virial_callback", &FixDLExt::set_virial_callback) ; } From 0689eb54c5d3e6a9b449f5a4c5b0ab3683169749 Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 11 Jul 2024 09:55:08 -0500 Subject: [PATCH 03/10] Virial contribution from FixDLExt to global virial is optional and by default is off --- dlext/include/FixDLExt.h | 1 + dlext/src/FixDLExt.cpp | 20 +++++++++++++++----- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index b27fc10..fcdadc4 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -41,6 +41,7 @@ class DEFAULT_VISIBILITY FixDLExt : public Fix { void post_force(int) override; void set_callback(DLExtCallback& cb); void set_virial_callback(DLExtSetVirial& cb); + void set_virial_global(int flag) { virial_global_flag = flag; } private: DLExtCallback callback = [](TimeStep) { }; diff --git a/dlext/src/FixDLExt.cpp b/dlext/src/FixDLExt.cpp index e83b351..f0dacb6 100644 --- a/dlext/src/FixDLExt.cpp +++ b/dlext/src/FixDLExt.cpp @@ -32,9 +32,8 @@ FixDLExt::FixDLExt(LAMMPS* lmp, int narg, char** arg) if (atom->map_style != Atom::MAP_ARRAY) error->all(FLERR, "Fix dlext requires to map atoms as arrays"); - // signal that this fix contributes to the global virial - virial_global_flag = 1; - thermo_virial = 1; + // signal that this fix contributes to the global virial or not, default no + virial_global_flag = 0; kokkosable = has_kokkos_cuda_enabled(lmp); atomKK = dynamic_cast(atom); @@ -43,20 +42,31 @@ FixDLExt::FixDLExt(LAMMPS* lmp, int narg, char** arg) datamask_modify = EMPTY_MASK; } -int FixDLExt::setmask() { return FixConst::POST_FORCE; } +int FixDLExt::setmask() +{ + return FixConst::POST_FORCE; +} + void FixDLExt::post_force(int vflag) { // virial setup + v_init(vflag); // invoke callback + callback(update->ntimestep); // put the virial from the bias into this fix's member variable virial[6] (see fix.h) - setVirial(virial); + + if (virial_global_flag) + setVirial(virial); } + +// callback from the sampling method to add the biasing forces to the atoms void FixDLExt::set_callback(DLExtCallback& cb) { callback = cb; } +// callback from the sampling method to set the virial contribution to the fix's virial void FixDLExt::set_virial_callback(DLExtSetVirial& cb) { setVirial = cb; } void register_FixDLExt(LAMMPS* lmp) From bb4e05164f5e7ff37139255f9478f055118abaa8 Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 11 Jul 2024 10:39:21 -0500 Subject: [PATCH 04/10] Wrapped fix->virial as a pointer to a DLManagedTensor struct --- dlext/include/DLExt.h | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/dlext/include/DLExt.h b/dlext/include/DLExt.h index 12bcc3e..6d61cfa 100644 --- a/dlext/include/DLExt.h +++ b/dlext/include/DLExt.h @@ -6,6 +6,7 @@ #include "LAMMPSView.h" #include "atom.h" +#include "fix.h" #ifdef LMP_KOKKOS #include "atom_kokkos.h" @@ -34,6 +35,7 @@ static struct Images { } kImages; static struct Tags { } kTags; static struct TagsMap { } kTagsMap; static struct Types { } kTypes; +static struct Virial { } kVirial; static struct SecondDim { } kSecondDim; @@ -61,6 +63,7 @@ inline void* opaque(const T* data) return const_cast(data); } +// if LAMMPS is built with KOKKOS, bind the PROPERTY struct to the corresponding ACCESSOR #ifdef LMP_KOKKOS #define DLEXT_OPAQUE_ATOM_KOKKOS(PROPERTY, ACCESSOR) \ inline void* opaque(const AtomKokkos* atom, PROPERTY) \ @@ -80,6 +83,7 @@ DLEXT_OPAQUE_ATOM_KOKKOS(Types, k_type) #undef DLEXT_OPAQUE_ATOM_KOKKOS #endif +// return the underlying pointers in LAMMPS (Property can be used as a tag, or actually bound to the KOKKOS accessor as above) inline void* opaque(const Atom* atom, Positions) { return opaque(atom->x[0]); } inline void* opaque(const Atom* atom, Velocities) { return opaque(atom->v[0]); } inline void* opaque(const Atom* atom, Masses) { return opaque(atom->mass); } @@ -91,6 +95,7 @@ inline void* opaque(const Atom* atom, TagsMap) { return opaque(const_cast(atom)->get_map_array()); } +inline void* opaque(const Fix* fix, Virial) { return opaque(fix->virial); } template inline void* opaque(const LAMMPSView& view, DLDeviceType device_type, Property p) @@ -102,11 +107,13 @@ inline void* opaque(const LAMMPSView& view, DLDeviceType device_type, Property p return opaque(view.atom_ptr(), p); } +// get the device info (id) from view and device_type, return a DLDevice struct inline DLDevice device_info(const LAMMPSView& view, DLDeviceType device_type) { return DLDevice { device_type, view.device_id() }; } +// return the DLDataType code corresonding to the actual data type of the "Tag" constexpr DLDataTypeCode dtype_code(Positions) { return kDLFloat; } constexpr DLDataTypeCode dtype_code(Velocities) { return kDLFloat; } constexpr DLDataTypeCode dtype_code(Masses) { return kDLFloat; } @@ -115,7 +122,9 @@ constexpr DLDataTypeCode dtype_code(Images) { return kDLInt; } constexpr DLDataTypeCode dtype_code(Tags) { return kDLInt; } constexpr DLDataTypeCode dtype_code(TagsMap) { return kDLInt; } constexpr DLDataTypeCode dtype_code(Types) { return kDLInt; } +constexpr DLDataTypeCode dtype_code(Virial) { return kDLFloat; } +// return the number of bits of the data type of a given PROPERTY #define DLEXT_BITS_FLOAT_ARRAY(PROPERTY, TYPE) \ inline uint8_t bits(DLDeviceType device_type, PROPERTY) \ { \ @@ -126,6 +135,7 @@ DLEXT_BITS_FLOAT_ARRAY(Positions, X_FLOAT) DLEXT_BITS_FLOAT_ARRAY(Velocities, V_FLOAT) DLEXT_BITS_FLOAT_ARRAY(Masses, LMP_FLOAT) DLEXT_BITS_FLOAT_ARRAY(Forces, F_FLOAT) +DLEXT_BITS_FLOAT_ARRAY(Virial, F_FLOAT) #undef DLEXT_BITS_FLOAT_ARRAY @@ -157,6 +167,7 @@ inline int64_t size(const LAMMPSView& view, Property) } inline int64_t size(const LAMMPSView& view, Masses) { return view.atom_ptr()->ntypes + 1; } inline int64_t size(const LAMMPSView& view, TagsMap) { return view.atom_ptr()->get_map_size(); } +inline int64_t size(const LAMMPSView& view, Virial) { return 6; } template inline int64_t size(const LAMMPSView& view, Property, SecondDim) @@ -173,25 +184,37 @@ constexpr uint64_t offset(const LAMMPSView& view, Property p) return 0; } +// a templated function for wrapping a C array given its data type and dimensions +// and returning a pointer to a DLPack tensor template DLManagedTensor* wrap(const LAMMPSView& view, Property property, ExecutionSpace exec_space) { + // get the device type of the view (host or device) + auto device_type = view.device_type(exec_space); + auto bridge = std::make_unique(); bridge->tensor.manager_ctx = bridge.get(); bridge->tensor.deleter = delete_bridge; + // acquire the actual dltensor pointer auto& dltensor = bridge->tensor.dl_tensor; - auto device_type = view.device_type(exec_space); + + // fill in the dltensor struct + // get the underlying array/accessor of the given property and assign it to data (as void*) dltensor.data = opaque(view, device_type, property); + // get the device info from view and device_type and assign it to device (as DLDevice) dltensor.device = device_info(view, device_type); + // get the data type of the underlying array (DLDataType) given the data type code and number of bits dltensor.dtype = dtype(device_type, property); + // fill in the tensor shape (dimensions), strides and byte offsets auto& shape = bridge->shape; auto size2 = size(view, property, kSecondDim); shape.push_back(size(view, property)); + // if the array is 2D if (size2 > 1) shape.push_back(size2); - + // strides between consecutive elements in each dim auto& strides = bridge->strides; strides.push_back(size2); if (size2 > 1) @@ -205,12 +228,14 @@ DLManagedTensor* wrap(const LAMMPSView& view, Property property, ExecutionSpace return &(bridge.release()->tensor); } +// macro that returns a DLManagedTensor from view for a given SELECTOR (Property) #define DLEXT_PROPERTY_FROM_VIEW(FN, SELECTOR) \ inline DLManagedTensor* FN(const LAMMPSView& view, ExecutionSpace space) \ { \ return wrap(view, SELECTOR, space); \ } +// finally, all the function instances to pack arrays into DLManagedTensor structs DLEXT_PROPERTY_FROM_VIEW(positions, kPositions) DLEXT_PROPERTY_FROM_VIEW(velocities, kVelocities) DLEXT_PROPERTY_FROM_VIEW(masses, kMasses) @@ -219,6 +244,7 @@ DLEXT_PROPERTY_FROM_VIEW(images, kImages) DLEXT_PROPERTY_FROM_VIEW(tags, kTags) DLEXT_PROPERTY_FROM_VIEW(tags_map, kTagsMap) DLEXT_PROPERTY_FROM_VIEW(types, kTypes) +DLEXT_PROPERTY_FROM_VIEW(virial, kVirial) #undef DLEXT_PROPERTY From 5b7c495c1c6fbd4608306c6521485651eb5042bb Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Wed, 14 Aug 2024 13:56:39 -0500 Subject: [PATCH 05/10] Changed from LAMMPSView to Fix for access to lmp and other per-atom pointers --- dlext/include/DLExt.h | 77 ++++++++++++++++++++-------------------- dlext/include/FixDLExt.h | 31 ++++++++++++++-- dlext/src/FixDLExt.cpp | 36 +++++++++++++++++++ 3 files changed, 102 insertions(+), 42 deletions(-) diff --git a/dlext/include/DLExt.h b/dlext/include/DLExt.h index 6d61cfa..56f4b23 100644 --- a/dlext/include/DLExt.h +++ b/dlext/include/DLExt.h @@ -4,7 +4,6 @@ #ifndef LAMMPS_DLPACK_EXTENSION_H_ #define LAMMPS_DLPACK_EXTENSION_H_ -#include "LAMMPSView.h" #include "atom.h" #include "fix.h" @@ -98,19 +97,19 @@ inline void* opaque(const Atom* atom, TagsMap) inline void* opaque(const Fix* fix, Virial) { return opaque(fix->virial); } template -inline void* opaque(const LAMMPSView& view, DLDeviceType device_type, Property p) +inline void* opaque(const Fix* fixdlext, DLDeviceType device_type, Property p) { #ifdef LMP_KOKKOS if (device_type == kDLCUDA) - return opaque(view.atom_kokkos_ptr(), p); + return opaque(fixdlext->atom_kokkos_ptr(), p); #endif - return opaque(view.atom_ptr(), p); + return opaque(fixdlext->atom_ptr(), p); } -// get the device info (id) from view and device_type, return a DLDevice struct -inline DLDevice device_info(const LAMMPSView& view, DLDeviceType device_type) +// get the device info (id) from fix and device_type, return a DLDevice struct +inline DLDevice device_info(const Fix* fixdlext, DLDeviceType device_type) { - return DLDevice { device_type, view.device_id() }; + return DLDevice { device_type, fixdlext->device_id() }; } // return the DLDataType code corresonding to the actual data type of the "Tag" @@ -161,25 +160,25 @@ inline DLDataType dtype(DLDeviceType device_type, Property p) } template -inline int64_t size(const LAMMPSView& view, Property) +inline int64_t size(const Fix* fixdlext, Property) { - return view.local_particle_number(); + return fixdlext->local_particle_number(); } -inline int64_t size(const LAMMPSView& view, Masses) { return view.atom_ptr()->ntypes + 1; } -inline int64_t size(const LAMMPSView& view, TagsMap) { return view.atom_ptr()->get_map_size(); } -inline int64_t size(const LAMMPSView& view, Virial) { return 6; } +inline int64_t size(const Fix* fixdlext, Masses) { return fixdlext->atom_ptr()->ntypes + 1; } +inline int64_t size(const Fix* fixdlext, TagsMap) { return fixdlext->atom_ptr()->get_map_size(); } +inline int64_t size(const Fix* fixdlext, Virial) { return 6; } template -inline int64_t size(const LAMMPSView& view, Property, SecondDim) +inline int64_t size(const Fix* fixdlext, Property, SecondDim) { return 1; } -inline int64_t size(const LAMMPSView& view, Positions, SecondDim) { return 3; } -inline int64_t size(const LAMMPSView& view, Velocities, SecondDim) { return 3; } -inline int64_t size(const LAMMPSView& view, Forces, SecondDim) { return 3; } +inline int64_t size(const Fix* fixdlext, Positions, SecondDim) { return 3; } +inline int64_t size(const Fix* fixdlext, Velocities, SecondDim) { return 3; } +inline int64_t size(const Fix* fixdlext, Forces, SecondDim) { return 3; } template -constexpr uint64_t offset(const LAMMPSView& view, Property p) +constexpr uint64_t offset(const Fix* fixdlext, Property p) { return 0; } @@ -187,10 +186,10 @@ constexpr uint64_t offset(const LAMMPSView& view, Property p) // a templated function for wrapping a C array given its data type and dimensions // and returning a pointer to a DLPack tensor template -DLManagedTensor* wrap(const LAMMPSView& view, Property property, ExecutionSpace exec_space) +DLManagedTensor* wrap(const Fix* fixdlext, Property property, ExecutionSpace exec_space) { - // get the device type of the view (host or device) - auto device_type = view.device_type(exec_space); + // get the device type of the fix (host or device) + auto device_type = fixdlext->device_type(exec_space); auto bridge = std::make_unique(); bridge->tensor.manager_ctx = bridge.get(); @@ -201,16 +200,16 @@ DLManagedTensor* wrap(const LAMMPSView& view, Property property, ExecutionSpace // fill in the dltensor struct // get the underlying array/accessor of the given property and assign it to data (as void*) - dltensor.data = opaque(view, device_type, property); - // get the device info from view and device_type and assign it to device (as DLDevice) - dltensor.device = device_info(view, device_type); + dltensor.data = opaque(fixdlext, device_type, property); + // get the device info from fix and device_type and assign it to device (as DLDevice) + dltensor.device = device_info(fixdlext, device_type); // get the data type of the underlying array (DLDataType) given the data type code and number of bits dltensor.dtype = dtype(device_type, property); // fill in the tensor shape (dimensions), strides and byte offsets auto& shape = bridge->shape; - auto size2 = size(view, property, kSecondDim); - shape.push_back(size(view, property)); + auto size2 = size(fixdlext, property, kSecondDim); + shape.push_back(size(fixdlext, property)); // if the array is 2D if (size2 > 1) shape.push_back(size2); @@ -223,28 +222,28 @@ DLManagedTensor* wrap(const LAMMPSView& view, Property property, ExecutionSpace dltensor.ndim = shape.size(); dltensor.shape = reinterpret_cast(shape.data()); dltensor.strides = reinterpret_cast(strides.data()); - dltensor.byte_offset = offset(view, property); + dltensor.byte_offset = offset(fixdlext, property); return &(bridge.release()->tensor); } -// macro that returns a DLManagedTensor from view for a given SELECTOR (Property) -#define DLEXT_PROPERTY_FROM_VIEW(FN, SELECTOR) \ - inline DLManagedTensor* FN(const LAMMPSView& view, ExecutionSpace space) \ +// macro that returns a DLManagedTensor from fix for a given SELECTOR (Property) +#define DLEXT_PROPERTY_FROM_FIX(FN, SELECTOR) \ + inline DLManagedTensor* FN(const Fix* fixdlext, ExecutionSpace space) \ { \ - return wrap(view, SELECTOR, space); \ + return wrap(fixdlext, SELECTOR, space); \ } // finally, all the function instances to pack arrays into DLManagedTensor structs -DLEXT_PROPERTY_FROM_VIEW(positions, kPositions) -DLEXT_PROPERTY_FROM_VIEW(velocities, kVelocities) -DLEXT_PROPERTY_FROM_VIEW(masses, kMasses) -DLEXT_PROPERTY_FROM_VIEW(forces, kForces) -DLEXT_PROPERTY_FROM_VIEW(images, kImages) -DLEXT_PROPERTY_FROM_VIEW(tags, kTags) -DLEXT_PROPERTY_FROM_VIEW(tags_map, kTagsMap) -DLEXT_PROPERTY_FROM_VIEW(types, kTypes) -DLEXT_PROPERTY_FROM_VIEW(virial, kVirial) +DLEXT_PROPERTY_FROM_FIX(positions, kPositions) +DLEXT_PROPERTY_FROM_FIX(velocities, kVelocities) +DLEXT_PROPERTY_FROM_FIX(masses, kMasses) +DLEXT_PROPERTY_FROM_FIX(forces, kForces) +DLEXT_PROPERTY_FROM_FIX(images, kImages) +DLEXT_PROPERTY_FROM_FIX(tags, kTags) +DLEXT_PROPERTY_FROM_FIX(tags_map, kTagsMap) +DLEXT_PROPERTY_FROM_FIX(types, kTypes) +DLEXT_PROPERTY_FROM_FIX(virial, kVirial) #undef DLEXT_PROPERTY diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index fcdadc4..7e6ddf6 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -4,10 +4,8 @@ #ifndef DLEXT_SAMPLER_H_ #define DLEXT_SAMPLER_H_ -#include "LAMMPSView.h" - +#include "dlpack/dlpack.h" #include "fix.h" - #include namespace LAMMPS_NS @@ -16,6 +14,8 @@ namespace dlext { // { // Aliases +const auto kOnHost = ExecutionSpace::Host; +const auto kOnDevice = ExecutionSpace::Device; using TimeStep = bigint; // bigint depends on how LAMMPS was built using DLExtCallback = std::function; @@ -43,11 +43,36 @@ class DEFAULT_VISIBILITY FixDLExt : public Fix { void set_virial_callback(DLExtSetVirial& cb); void set_virial_global(int flag) { virial_global_flag = flag; } + // Provide easy access to the atom pointers + Atom* atom_ptr() const; + AtomKokkos* atom_kokkos_ptr() const; + + //! Given an execution space, returns kDLCUDA if LAMMPS was built with KOKKOS and + //! Cuda supoprt, and it's available at runtime. Otherwise, returns kDLCPU. + DLDeviceType device_type(ExecutionSpace requested_space) const; + + //! The device id where this class instances are being executed + int device_id() const; + + // Convenience methods for retriving the number of particles + int local_particle_number() const; + bigint global_particle_number() const; + + //! If KOKKOS is available, synchronize on the particle data on the requested space + void synchronize(ExecutionSpace requested_space = kOnDevice); + private: DLExtCallback callback = [](TimeStep) { }; DLExtSetVirial setVirial = [](double*) { }; + ExecutionSpace try_pick(ExecutionSpace requested_space) const; }; +inline bool has_kokkos_cuda_enabled(LAMMPS* lmp) +{ + bool has_cuda = strcmp(LMPDeviceType::name(), "Cuda") == 0; + return has_cuda & (lmp->kokkos != nullptr); +} + void register_FixDLExt(LAMMPS* lmp); } // namespace dlext diff --git a/dlext/src/FixDLExt.cpp b/dlext/src/FixDLExt.cpp index f0dacb6..615548a 100644 --- a/dlext/src/FixDLExt.cpp +++ b/dlext/src/FixDLExt.cpp @@ -49,6 +49,17 @@ int FixDLExt::setmask() void FixDLExt::post_force(int vflag) { +#ifdef LMP_KOKKOS + if (has_kokkos_cuda_enabled(lmp)) { + // Since there's is no MASS_MASK, we need to make sure + // masses are available on the device. + atom_kokkos_ptr()->k_mass.sync_device(); + // On the other hand, MAP_MASK exists, but it's never used + // within any of the synchronizations methods. + atom_kokkos_ptr()->k_map_array.sync_device(); + } +#endif + // virial setup v_init(vflag); @@ -69,6 +80,31 @@ void FixDLExt::set_callback(DLExtCallback& cb) { callback = cb; } // callback from the sampling method to set the virial contribution to the fix's virial void FixDLExt::set_virial_callback(DLExtSetVirial& cb) { setVirial = cb; } +// return the pointer to the atom pointer +Atom* FixDLExt::atom_ptr() const { return lmp->atom; } +AtomKokkos* FixDLExt::atom_kokkos_ptr() const { return lmp->atomKK; } + +// return the device type +DLDeviceType FixDLExt::device_type(ExecutionSpace requested_space) const +{ + return (try_pick(requested_space) == kOnDevice) ? kDLCUDA : kDLCPU; +} + +// return the device id +// TODO: infer this from the LAMMPS instance +int FixDLExt::device_id() const { return 0; } + +int FixDLExt::local_particle_number() const { return atom_ptr()->nlocal; } +bigint FixDLExt::global_particle_number() const { return atom_ptr()->natoms; } + +void FixDLExt::synchronize(ExecutionSpace requested_space) +{ + if (lmp->kokkos) { + atom_kokkos_ptr()->sync(try_pick(requested_space), DLEXT_MASK); + atom_kokkos_ptr()->k_map_array.sync_device(); + } +} + void register_FixDLExt(LAMMPS* lmp) { auto fix_map = lmp->modify->fix_map; From e2a871924bd3d520d326c4262e0b39ebde86a4d4 Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Wed, 14 Aug 2024 14:30:50 -0500 Subject: [PATCH 06/10] Modified the python C++ interface for using Fix --- dlext/include/FixDLExt.h | 3 +++ dlext/src/FixDLExt.cpp | 12 ++++++++++-- python/PyDLExt.h | 4 ++-- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index 7e6ddf6..0a1cec2 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -4,6 +4,7 @@ #ifndef DLEXT_SAMPLER_H_ #define DLEXT_SAMPLER_H_ +#include "cxx11utils.h" #include "dlpack/dlpack.h" #include "fix.h" #include @@ -13,6 +14,8 @@ namespace LAMMPS_NS namespace dlext { +using namespace cxx11; + // { // Aliases const auto kOnHost = ExecutionSpace::Host; const auto kOnDevice = ExecutionSpace::Device; diff --git a/dlext/src/FixDLExt.cpp b/dlext/src/FixDLExt.cpp index 615548a..b15ff8b 100644 --- a/dlext/src/FixDLExt.cpp +++ b/dlext/src/FixDLExt.cpp @@ -91,8 +91,16 @@ DLDeviceType FixDLExt::device_type(ExecutionSpace requested_space) const } // return the device id -// TODO: infer this from the LAMMPS instance -int FixDLExt::device_id() const { return 0; } +// TODO: would be handy if this can be available from lmp->kokkos +// KokkosLMP currently initializes KOKKOS with a temporary device_id in the class constructor +// but doesn't have this variable as a pubic class member +int FixDLExt::device_id() const +{ +#ifdef LMP_KOKKOS_GPU + return 0; +#endif + return 0; +} int FixDLExt::local_particle_number() const { return atom_ptr()->nlocal; } bigint FixDLExt::global_particle_number() const { return atom_ptr()->natoms; } diff --git a/python/PyDLExt.h b/python/PyDLExt.h index 9618398..c7fc92e 100644 --- a/python/PyDLExt.h +++ b/python/PyDLExt.h @@ -22,9 +22,9 @@ const char* const kDLTensorCapsuleName = "dltensor"; // See the DLPack Documentation https://dmlc.github.io/dlpack/latest/python_spec.html template -inline PyCapsule enpycapsulate(const LAMMPSView& view, ExecutionSpace space) +inline PyCapsule enpycapsulate(const Fix* fixdlext, ExecutionSpace space) { - auto dl_managed_tensor = property(view, space); + auto dl_managed_tensor = property(fixdlext, space); return PyCapsule( dl_managed_tensor, // PyCapsule pointer kDLTensorCapsuleName, // PyCapsule name From 7a9070b2887805c99e9b52cdc797464c8629b72f Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 22 Aug 2024 13:08:01 -0500 Subject: [PATCH 07/10] Reverted some changes in the last commit, created a LAMMPSView member variable in FixDLExt --- dlext/include/DLExt.h | 48 ++++++++++++++++++++-------------------- dlext/include/FixDLExt.h | 33 ++------------------------- dlext/src/FixDLExt.cpp | 44 ------------------------------------ python/PyDLExt.h | 4 ++-- 4 files changed, 28 insertions(+), 101 deletions(-) diff --git a/dlext/include/DLExt.h b/dlext/include/DLExt.h index 56f4b23..7f58efd 100644 --- a/dlext/include/DLExt.h +++ b/dlext/include/DLExt.h @@ -97,19 +97,19 @@ inline void* opaque(const Atom* atom, TagsMap) inline void* opaque(const Fix* fix, Virial) { return opaque(fix->virial); } template -inline void* opaque(const Fix* fixdlext, DLDeviceType device_type, Property p) +inline void* opaque(const Fix* fix, DLDeviceType device_type, Property p) { #ifdef LMP_KOKKOS if (device_type == kDLCUDA) - return opaque(fixdlext->atom_kokkos_ptr(), p); + return opaque(fix->view.atom_kokkos_ptr(), p); #endif - return opaque(fixdlext->atom_ptr(), p); + return opaque(fix->view.atom_ptr(), p); } // get the device info (id) from fix and device_type, return a DLDevice struct -inline DLDevice device_info(const Fix* fixdlext, DLDeviceType device_type) +inline DLDevice device_info(const Fix* fix, DLDeviceType device_type) { - return DLDevice { device_type, fixdlext->device_id() }; + return DLDevice { device_type, fix->view.device_id() }; } // return the DLDataType code corresonding to the actual data type of the "Tag" @@ -160,25 +160,25 @@ inline DLDataType dtype(DLDeviceType device_type, Property p) } template -inline int64_t size(const Fix* fixdlext, Property) +inline int64_t size(const Fix* fix, Property) { - return fixdlext->local_particle_number(); + return fix->view.local_particle_number(); } -inline int64_t size(const Fix* fixdlext, Masses) { return fixdlext->atom_ptr()->ntypes + 1; } -inline int64_t size(const Fix* fixdlext, TagsMap) { return fixdlext->atom_ptr()->get_map_size(); } -inline int64_t size(const Fix* fixdlext, Virial) { return 6; } +inline int64_t size(const Fix* fix, Masses) { return fix->view.atom_ptr()->ntypes + 1; } +inline int64_t size(const Fix* fix, TagsMap) { return fix->view.atom_ptr()->get_map_size(); } +inline int64_t size(const Fix* fix, Virial) { return 6; } template -inline int64_t size(const Fix* fixdlext, Property, SecondDim) +inline int64_t size(const Fix* fix, Property, SecondDim) { return 1; } -inline int64_t size(const Fix* fixdlext, Positions, SecondDim) { return 3; } -inline int64_t size(const Fix* fixdlext, Velocities, SecondDim) { return 3; } -inline int64_t size(const Fix* fixdlext, Forces, SecondDim) { return 3; } +inline int64_t size(const Fix* fix, Positions, SecondDim) { return 3; } +inline int64_t size(const Fix* fix, Velocities, SecondDim) { return 3; } +inline int64_t size(const Fix* fix, Forces, SecondDim) { return 3; } template -constexpr uint64_t offset(const Fix* fixdlext, Property p) +constexpr uint64_t offset(const Fix* fix, Property p) { return 0; } @@ -186,10 +186,10 @@ constexpr uint64_t offset(const Fix* fixdlext, Property p) // a templated function for wrapping a C array given its data type and dimensions // and returning a pointer to a DLPack tensor template -DLManagedTensor* wrap(const Fix* fixdlext, Property property, ExecutionSpace exec_space) +DLManagedTensor* wrap(const Fix* fix, Property property, ExecutionSpace exec_space) { // get the device type of the fix (host or device) - auto device_type = fixdlext->device_type(exec_space); + auto device_type = fix->device_type(exec_space); auto bridge = std::make_unique(); bridge->tensor.manager_ctx = bridge.get(); @@ -200,16 +200,16 @@ DLManagedTensor* wrap(const Fix* fixdlext, Property property, ExecutionSpace exe // fill in the dltensor struct // get the underlying array/accessor of the given property and assign it to data (as void*) - dltensor.data = opaque(fixdlext, device_type, property); + dltensor.data = opaque(fix, device_type, property); // get the device info from fix and device_type and assign it to device (as DLDevice) - dltensor.device = device_info(fixdlext, device_type); + dltensor.device = device_info(fix, device_type); // get the data type of the underlying array (DLDataType) given the data type code and number of bits dltensor.dtype = dtype(device_type, property); // fill in the tensor shape (dimensions), strides and byte offsets auto& shape = bridge->shape; - auto size2 = size(fixdlext, property, kSecondDim); - shape.push_back(size(fixdlext, property)); + auto size2 = size(fix, property, kSecondDim); + shape.push_back(size(fix, property)); // if the array is 2D if (size2 > 1) shape.push_back(size2); @@ -222,16 +222,16 @@ DLManagedTensor* wrap(const Fix* fixdlext, Property property, ExecutionSpace exe dltensor.ndim = shape.size(); dltensor.shape = reinterpret_cast(shape.data()); dltensor.strides = reinterpret_cast(strides.data()); - dltensor.byte_offset = offset(fixdlext, property); + dltensor.byte_offset = offset(fix, property); return &(bridge.release()->tensor); } // macro that returns a DLManagedTensor from fix for a given SELECTOR (Property) #define DLEXT_PROPERTY_FROM_FIX(FN, SELECTOR) \ - inline DLManagedTensor* FN(const Fix* fixdlext, ExecutionSpace space) \ + inline DLManagedTensor* FN(const Fix* fix, ExecutionSpace space) \ { \ - return wrap(fixdlext, SELECTOR, space); \ + return wrap(fix, SELECTOR, space); \ } // finally, all the function instances to pack arrays into DLManagedTensor structs diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index 0a1cec2..8550829 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -4,8 +4,6 @@ #ifndef DLEXT_SAMPLER_H_ #define DLEXT_SAMPLER_H_ -#include "cxx11utils.h" -#include "dlpack/dlpack.h" #include "fix.h" #include @@ -14,11 +12,7 @@ namespace LAMMPS_NS namespace dlext { -using namespace cxx11; - // { // Aliases -const auto kOnHost = ExecutionSpace::Host; -const auto kOnDevice = ExecutionSpace::Device; using TimeStep = bigint; // bigint depends on how LAMMPS was built using DLExtCallback = std::function; @@ -46,36 +40,13 @@ class DEFAULT_VISIBILITY FixDLExt : public Fix { void set_virial_callback(DLExtSetVirial& cb); void set_virial_global(int flag) { virial_global_flag = flag; } - // Provide easy access to the atom pointers - Atom* atom_ptr() const; - AtomKokkos* atom_kokkos_ptr() const; - - //! Given an execution space, returns kDLCUDA if LAMMPS was built with KOKKOS and - //! Cuda supoprt, and it's available at runtime. Otherwise, returns kDLCPU. - DLDeviceType device_type(ExecutionSpace requested_space) const; - - //! The device id where this class instances are being executed - int device_id() const; - - // Convenience methods for retriving the number of particles - int local_particle_number() const; - bigint global_particle_number() const; +protected: + class LAMMPSView view; - //! If KOKKOS is available, synchronize on the particle data on the requested space - void synchronize(ExecutionSpace requested_space = kOnDevice); - -private: DLExtCallback callback = [](TimeStep) { }; DLExtSetVirial setVirial = [](double*) { }; - ExecutionSpace try_pick(ExecutionSpace requested_space) const; }; -inline bool has_kokkos_cuda_enabled(LAMMPS* lmp) -{ - bool has_cuda = strcmp(LMPDeviceType::name(), "Cuda") == 0; - return has_cuda & (lmp->kokkos != nullptr); -} - void register_FixDLExt(LAMMPS* lmp); } // namespace dlext diff --git a/dlext/src/FixDLExt.cpp b/dlext/src/FixDLExt.cpp index b15ff8b..f0dacb6 100644 --- a/dlext/src/FixDLExt.cpp +++ b/dlext/src/FixDLExt.cpp @@ -49,17 +49,6 @@ int FixDLExt::setmask() void FixDLExt::post_force(int vflag) { -#ifdef LMP_KOKKOS - if (has_kokkos_cuda_enabled(lmp)) { - // Since there's is no MASS_MASK, we need to make sure - // masses are available on the device. - atom_kokkos_ptr()->k_mass.sync_device(); - // On the other hand, MAP_MASK exists, but it's never used - // within any of the synchronizations methods. - atom_kokkos_ptr()->k_map_array.sync_device(); - } -#endif - // virial setup v_init(vflag); @@ -80,39 +69,6 @@ void FixDLExt::set_callback(DLExtCallback& cb) { callback = cb; } // callback from the sampling method to set the virial contribution to the fix's virial void FixDLExt::set_virial_callback(DLExtSetVirial& cb) { setVirial = cb; } -// return the pointer to the atom pointer -Atom* FixDLExt::atom_ptr() const { return lmp->atom; } -AtomKokkos* FixDLExt::atom_kokkos_ptr() const { return lmp->atomKK; } - -// return the device type -DLDeviceType FixDLExt::device_type(ExecutionSpace requested_space) const -{ - return (try_pick(requested_space) == kOnDevice) ? kDLCUDA : kDLCPU; -} - -// return the device id -// TODO: would be handy if this can be available from lmp->kokkos -// KokkosLMP currently initializes KOKKOS with a temporary device_id in the class constructor -// but doesn't have this variable as a pubic class member -int FixDLExt::device_id() const -{ -#ifdef LMP_KOKKOS_GPU - return 0; -#endif - return 0; -} - -int FixDLExt::local_particle_number() const { return atom_ptr()->nlocal; } -bigint FixDLExt::global_particle_number() const { return atom_ptr()->natoms; } - -void FixDLExt::synchronize(ExecutionSpace requested_space) -{ - if (lmp->kokkos) { - atom_kokkos_ptr()->sync(try_pick(requested_space), DLEXT_MASK); - atom_kokkos_ptr()->k_map_array.sync_device(); - } -} - void register_FixDLExt(LAMMPS* lmp) { auto fix_map = lmp->modify->fix_map; diff --git a/python/PyDLExt.h b/python/PyDLExt.h index c7fc92e..f5a278d 100644 --- a/python/PyDLExt.h +++ b/python/PyDLExt.h @@ -22,9 +22,9 @@ const char* const kDLTensorCapsuleName = "dltensor"; // See the DLPack Documentation https://dmlc.github.io/dlpack/latest/python_spec.html template -inline PyCapsule enpycapsulate(const Fix* fixdlext, ExecutionSpace space) +inline PyCapsule enpycapsulate(const LAMMPS& view, ExecutionSpace space) { - auto dl_managed_tensor = property(fixdlext, space); + auto dl_managed_tensor = property(view, space); return PyCapsule( dl_managed_tensor, // PyCapsule pointer kDLTensorCapsuleName, // PyCapsule name From 973030ef14a0bd73528870186d86cd4c9c36a3ef Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 22 Aug 2024 13:28:04 -0500 Subject: [PATCH 08/10] Exported view to python interface --- dlext/include/FixDLExt.h | 5 +++-- python/lammps_dlext.cpp | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index 8550829..22ff2cc 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -6,6 +6,7 @@ #include "fix.h" #include +#include "LAMMPSView.h" namespace LAMMPS_NS { @@ -40,9 +41,9 @@ class DEFAULT_VISIBILITY FixDLExt : public Fix { void set_virial_callback(DLExtSetVirial& cb); void set_virial_global(int flag) { virial_global_flag = flag; } -protected: - class LAMMPSView view; + LAMMPSView view; +protected: DLExtCallback callback = [](TimeStep) { }; DLExtSetVirial setVirial = [](double*) { }; }; diff --git a/python/lammps_dlext.cpp b/python/lammps_dlext.cpp index 8ad36bd..bc04214 100644 --- a/python/lammps_dlext.cpp +++ b/python/lammps_dlext.cpp @@ -50,6 +50,7 @@ void export_FixDLExt(py::module& m) })) .def("set_callback", &FixDLExt::set_callback) .def("set_virial_callback", &FixDLExt::set_virial_callback) + .def("view", &FixDLExt::view) ; } From 71c5262bbcae2f05b9ba5c6fe2fabd104cf60bf9 Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 22 Aug 2024 13:32:04 -0500 Subject: [PATCH 09/10] Added back the dlpack header file --- dlext/include/FixDLExt.h | 1 + 1 file changed, 1 insertion(+) diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index 22ff2cc..15f379b 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -4,6 +4,7 @@ #ifndef DLEXT_SAMPLER_H_ #define DLEXT_SAMPLER_H_ +#include "dlpack/dlpack.h" #include "fix.h" #include #include "LAMMPSView.h" From d89276788ae1dc89f7b58dbadf346cdda72b62d3 Mon Sep 17 00:00:00 2001 From: Trung Nguyen Date: Thu, 22 Aug 2024 15:55:13 -0500 Subject: [PATCH 10/10] Rearranged view and exposed it through pybind11 with accessors --- dlext/include/DLExt.h | 2 +- dlext/include/FixDLExt.h | 5 +++-- dlext/src/FixDLExt.cpp | 11 +++++++++++ python/lammps_dlext.cpp | 2 +- 4 files changed, 16 insertions(+), 4 deletions(-) diff --git a/dlext/include/DLExt.h b/dlext/include/DLExt.h index 7f58efd..708a200 100644 --- a/dlext/include/DLExt.h +++ b/dlext/include/DLExt.h @@ -189,7 +189,7 @@ template DLManagedTensor* wrap(const Fix* fix, Property property, ExecutionSpace exec_space) { // get the device type of the fix (host or device) - auto device_type = fix->device_type(exec_space); + auto device_type = fix->view.device_type(exec_space); auto bridge = std::make_unique(); bridge->tensor.manager_ctx = bridge.get(); diff --git a/dlext/include/FixDLExt.h b/dlext/include/FixDLExt.h index 15f379b..7af14a6 100644 --- a/dlext/include/FixDLExt.h +++ b/dlext/include/FixDLExt.h @@ -41,12 +41,13 @@ class DEFAULT_VISIBILITY FixDLExt : public Fix { void set_callback(DLExtCallback& cb); void set_virial_callback(DLExtSetVirial& cb); void set_virial_global(int flag) { virial_global_flag = flag; } - - LAMMPSView view; + void set_view(LAMMPSView _view); + LAMMPSView get_view() const; protected: DLExtCallback callback = [](TimeStep) { }; DLExtSetVirial setVirial = [](double*) { }; + LAMMPSView view; }; void register_FixDLExt(LAMMPS* lmp); diff --git a/dlext/src/FixDLExt.cpp b/dlext/src/FixDLExt.cpp index f0dacb6..83cb157 100644 --- a/dlext/src/FixDLExt.cpp +++ b/dlext/src/FixDLExt.cpp @@ -40,6 +40,11 @@ FixDLExt::FixDLExt(LAMMPS* lmp, int narg, char** arg) execution_space = (on_host || !kokkosable) ? kOnHost : kOnDevice; datamask_read = EMPTY_MASK; datamask_modify = EMPTY_MASK; + + // create an instance of LAMMPSView + // using default copy constructor and operator '=' of LAMMPSView works here + // because LAMMPSView simply encapsulates the LAMMPS pointers from lmp + view = LAMMPSView(lmp); } int FixDLExt::setmask() @@ -69,6 +74,12 @@ void FixDLExt::set_callback(DLExtCallback& cb) { callback = cb; } // callback from the sampling method to set the virial contribution to the fix's virial void FixDLExt::set_virial_callback(DLExtSetVirial& cb) { setVirial = cb; } +// set the LAMMPSView object +void FixDLExt::set_view(LAMMPSView _view) { view = _view; } + +// get the LAMMPSView object +LAMMPSView FixDLExt::get_view() const { return view; } + void register_FixDLExt(LAMMPS* lmp) { auto fix_map = lmp->modify->fix_map; diff --git a/python/lammps_dlext.cpp b/python/lammps_dlext.cpp index bc04214..fff9b41 100644 --- a/python/lammps_dlext.cpp +++ b/python/lammps_dlext.cpp @@ -50,7 +50,7 @@ void export_FixDLExt(py::module& m) })) .def("set_callback", &FixDLExt::set_callback) .def("set_virial_callback", &FixDLExt::set_virial_callback) - .def("view", &FixDLExt::view) + .def_property("view", &FixDLExt::get_view, &FixDLExt::set_view); ; }