diff --git a/source/api_cc/include/DeepPotPT.h b/source/api_cc/include/DeepPotPT.h index 62035d9010..efe453ae8e 100644 --- a/source/api_cc/include/DeepPotPT.h +++ b/source/api_cc/include/DeepPotPT.h @@ -1,9 +1,10 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once +#include + #include "DeepPot.h" #include "commonPT.h" -#include namespace deepmd { /** @@ -33,8 +34,7 @@ class DeepPotPT : public DeepPotBase { * @param[in] file_content The content of the model file. If it is not empty, *DP will read from the string instead of the file. **/ - void init(const std::string& model, - const int& gpu_rank = 0); + void init(const std::string& model, const int& gpu_rank = 0); private: /** @@ -63,13 +63,13 @@ class DeepPotPT : public DeepPotBase { void compute(ENERGYVTYPE& ener, std::vector& force, std::vector& virial, - // std::vector& atom_energy, - // std::vector& atom_virial, + // std::vector& atom_energy, + // std::vector& atom_virial, const std::vector& coord, const std::vector& atype, const std::vector& box); - // const std::vector& fparam = std::vector(), - // const std::vector& aparam = std::vector()); + // const std::vector& fparam = std::vector(), + // const std::vector& aparam = std::vector()); /** * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial *by using this DP. @@ -99,16 +99,16 @@ class DeepPotPT : public DeepPotBase { void compute(ENERGYVTYPE& ener, std::vector& force, std::vector& virial, - // std::vector& atom_energy, - // std::vector& atom_virial, + // std::vector& atom_energy, + // std::vector& atom_virial, const std::vector& coord, const std::vector& atype, const std::vector& box, - // const int nghost, + // const int nghost, const InputNlist& lmp_list, const int& ago); - // const std::vector& fparam = std::vector(), - // const std::vector& aparam = std::vector()); + // const std::vector& fparam = std::vector(), + // const std::vector& aparam = std::vector()); /** * @brief Evaluate the energy, force, and virial with the mixed type *by using this DP. @@ -310,7 +310,6 @@ class DeepPotPT : public DeepPotBase { const std::vector& aparam = std::vector()); private: - bool inited; int ntypes; int ntypes_spin; @@ -320,10 +319,10 @@ class DeepPotPT : public DeepPotBase { torch::jit::script::Module module; double rcut; NeighborListData nlist_data; - //InputNlist nlist; + // InputNlist nlist; int max_num_neighbors; int gpu_id; - at::Tensor firstneigh_tensor; + at::Tensor firstneigh_tensor; }; } // namespace deepmd diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 6e7b957aeb..d2d2e106c7 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -1,26 +1,29 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later #include #ifndef COMMON_H #define COMMON_H -#include -#include #include +#include +#include #include -#include "neighbor_list.h" +#include "neighbor_list.h" struct NeighborListData { /// Array stores the core region atom's index std::vector ilist; /// Array stores the core region atom's neighbor index - //std::vector> jlist; - int *jlist; + // std::vector> jlist; + int* jlist; /// Array stores the number of neighbors of core region atoms std::vector numneigh; /// Array stores the the location of the first neighbor of core region atoms std::vector firstneigh; public: - void copy_from_nlist(const InputNlist& inlist, int& max_num_neighbors,int nnei); - //void make_inlist(InputNlist& inlist); + void copy_from_nlist(const InputNlist& inlist, + int& max_num_neighbors, + int nnei); + // void make_inlist(InputNlist& inlist); }; -#endif \ No newline at end of file +#endif diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 15030b8841..e3640517cf 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -38,7 +38,7 @@ void DeepPot::init(const std::string& model, // TODO: throw errors if TF backend is not built, without mentioning TF dp = std::make_shared(model, gpu_rank, file_content); } else if (deepmd::DPBackend::PyTorch == backend) { - //throw deepmd::deepmd_exception("PyTorch backend is not supported yet"); + // throw deepmd::deepmd_exception("PyTorch backend is not supported yet"); dp = std::make_shared(model, gpu_rank, file_content); } else if (deepmd::DPBackend::Paddle == backend) { throw deepmd::deepmd_exception("PaddlePaddle backend is not supported yet"); diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 956d6084bd..9f3cb998bf 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -1,12 +1,12 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later #include "DeepPotPT.h" using namespace deepmd; -DeepPotPT::DeepPotPT() -: inited(false) { } +DeepPotPT::DeepPotPT() : inited(false) {} DeepPotPT::DeepPotPT(const std::string& model, const int& gpu_rank, const std::string& file_content) - : inited(false){ + : inited(false) { try { init(model, gpu_rank); } catch (...) { @@ -15,146 +15,156 @@ DeepPotPT::DeepPotPT(const std::string& model, } } void DeepPotPT::init(const std::string& model, const int& gpu_rank) { - if (inited) { + if (inited) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " "nothing at the second call of initializer" << std::endl; return; - } - std::cout << "load model from: " <(rcut_); - inited = true; + } + std::cout << "load model from: " << model << " to gpu " << gpu_rank + << std::endl; + gpu_id = gpu_rank; + torch::Device device(torch::kCUDA, gpu_rank); + + module = torch::jit::load(model, device); + torch::jit::FusionStrategy strategy; + strategy = {{torch::jit::FusionBehavior::DYNAMIC, 10}}; + torch::jit::setFusionStrategy(strategy); + + // at::globalContext().setAllowTF32CuBLAS(true); + // at::globalContext().setAllowTF32CuDNN(true); + auto rcut_ = module.run_method("get_rcut").toDouble(); + rcut = static_cast(rcut_); + inited = true; } -DeepPotPT::~DeepPotPT() { } - - +DeepPotPT::~DeepPotPT() {} template void DeepPotPT::compute(ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const InputNlist& lmp_list, - const int& ago) -{ - torch::Device device(torch::kCUDA, gpu_id); - std::vector coord_wrapped = coord; - int natoms = atype.size(); - auto options = torch::TensorOptions().dtype(torch::kFloat64); - auto int_options = torch::TensorOptions().dtype(torch::kInt64); - auto int32_options = torch::TensorOptions().dtype(torch::kInt32); - std::vector inputs; - at::Tensor coord_wrapped_Tensor = torch::from_blob(coord_wrapped.data(), {1,natoms, 3}, options).to(device); - inputs.push_back(coord_wrapped_Tensor); - std::vector atype_64(atype.begin(), atype.end()); - at::Tensor atype_Tensor = torch::from_blob(atype_64.data(), {1,natoms}, int_options).to(device); - inputs.push_back(atype_Tensor); - if(ago == 0) - { - int64_t nnei = module.run_method("get_nnei").toInt(); - nlist_data.copy_from_nlist(lmp_list,max_num_neighbors,nnei); - if(max_num_neighbors > nnei) - { - at::Tensor firstneigh = torch::from_blob(nlist_data.jlist, {lmp_list.inum,max_num_neighbors}, int32_options); - at::Tensor nlist= firstneigh.to(torch::kInt64).to(device); - firstneigh_tensor = module.run_method("sort_neighbor_list",coord_wrapped_Tensor,nlist).toTensor(); - } - else - { - at::Tensor firstneigh = torch::from_blob(nlist_data.jlist, {1,lmp_list.inum,max_num_neighbors}, int32_options); - firstneigh_tensor = firstneigh.to(torch::kInt64).to(device); - } + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const InputNlist& lmp_list, + const int& ago) { + torch::Device device(torch::kCUDA, gpu_id); + std::vector coord_wrapped = coord; + int natoms = atype.size(); + auto options = torch::TensorOptions().dtype(torch::kFloat64); + auto int_options = torch::TensorOptions().dtype(torch::kInt64); + auto int32_options = torch::TensorOptions().dtype(torch::kInt32); + std::vector inputs; + at::Tensor coord_wrapped_Tensor = + torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options) + .to(device); + inputs.push_back(coord_wrapped_Tensor); + std::vector atype_64(atype.begin(), atype.end()); + at::Tensor atype_Tensor = + torch::from_blob(atype_64.data(), {1, natoms}, int_options).to(device); + inputs.push_back(atype_Tensor); + if (ago == 0) { + int64_t nnei = module.run_method("get_nnei").toInt(); + nlist_data.copy_from_nlist(lmp_list, max_num_neighbors, nnei); + if (max_num_neighbors > nnei) { + at::Tensor firstneigh = torch::from_blob( + nlist_data.jlist, {lmp_list.inum, max_num_neighbors}, int32_options); + at::Tensor nlist = firstneigh.to(torch::kInt64).to(device); + firstneigh_tensor = + module.run_method("sort_neighbor_list", coord_wrapped_Tensor, nlist) + .toTensor(); + } else { + at::Tensor firstneigh = torch::from_blob( + nlist_data.jlist, {1, lmp_list.inum, max_num_neighbors}, + int32_options); + firstneigh_tensor = firstneigh.to(torch::kInt64).to(device); } - inputs.push_back(firstneigh_tensor); - c10::Dict outputs = module.forward(inputs).toGenericDict(); - c10::IValue energy_ = outputs.at("energy"); - c10::IValue force_ = outputs.at("extended_force"); - c10::IValue virial_ = outputs.at("extended_virial"); - ener = energy_.toTensor().item(); - - torch::Tensor flat_force_ = force_.toTensor().view({-1}); - torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); - force.assign(cpu_force_.data_ptr(), cpu_force_.data_ptr() + cpu_force_.numel()); - - torch::Tensor flat_virial_ = virial_.toTensor().view({-1}); - torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); - virial.assign(cpu_virial_.data_ptr(), cpu_virial_.data_ptr() + cpu_virial_.numel()); - + } + inputs.push_back(firstneigh_tensor); + c10::Dict outputs = + module.forward(inputs).toGenericDict(); + c10::IValue energy_ = outputs.at("energy"); + c10::IValue force_ = outputs.at("extended_force"); + c10::IValue virial_ = outputs.at("extended_virial"); + ener = energy_.toTensor().item(); + + torch::Tensor flat_force_ = force_.toTensor().view({-1}); + torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); + force.assign(cpu_force_.data_ptr(), + cpu_force_.data_ptr() + cpu_force_.numel()); + + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); } -template void DeepPotPT::compute(double& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const InputNlist& lmp_list, - const int& ago); - +template void DeepPotPT::compute( + double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const InputNlist& lmp_list, + const int& ago); template void DeepPotPT::compute(ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box) -{ - auto device = torch::kCUDA; - module.to(device); - std::vector coord_wrapped = coord; - int natoms = atype.size(); - auto options = torch::TensorOptions().dtype(torch::kFloat64); - auto int_options = torch::TensorOptions().dtype(torch::kInt64); - std::vector inputs; - at::Tensor coord_wrapped_Tensor = torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options).to(device); - inputs.push_back(coord_wrapped_Tensor); - std::vector atype_64(atype.begin(), atype.end()); - at::Tensor atype_Tensor = torch::from_blob(atype_64.data(), {1, natoms}, int_options).to(device); - inputs.push_back(atype_Tensor); - at::Tensor box_Tensor = torch::from_blob(const_cast(box.data()), {1, 9}, options).to(device); - inputs.push_back(box_Tensor); - c10::Dict outputs = module.forward(inputs).toGenericDict(); - - - c10::IValue energy_ = outputs.at("energy"); - c10::IValue force_ = outputs.at("force"); - c10::IValue virial_ = outputs.at("virial"); - ener = energy_.toTensor().item(); - - torch::Tensor flat_force_ = force_.toTensor().view({-1}); - torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); - force.assign(cpu_force_.data_ptr(), cpu_force_.data_ptr() + cpu_force_.numel()); - - torch::Tensor flat_virial_ = virial_.toTensor().view({-1}); - torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); - virial.assign(cpu_virial_.data_ptr(), cpu_virial_.data_ptr() + cpu_virial_.numel()); - + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box) { + auto device = torch::kCUDA; + module.to(device); + std::vector coord_wrapped = coord; + int natoms = atype.size(); + auto options = torch::TensorOptions().dtype(torch::kFloat64); + auto int_options = torch::TensorOptions().dtype(torch::kInt64); + std::vector inputs; + at::Tensor coord_wrapped_Tensor = + torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options) + .to(device); + inputs.push_back(coord_wrapped_Tensor); + std::vector atype_64(atype.begin(), atype.end()); + at::Tensor atype_Tensor = + torch::from_blob(atype_64.data(), {1, natoms}, int_options).to(device); + inputs.push_back(atype_Tensor); + at::Tensor box_Tensor = + torch::from_blob(const_cast(box.data()), {1, 9}, options) + .to(device); + inputs.push_back(box_Tensor); + c10::Dict outputs = + module.forward(inputs).toGenericDict(); + + c10::IValue energy_ = outputs.at("energy"); + c10::IValue force_ = outputs.at("force"); + c10::IValue virial_ = outputs.at("virial"); + ener = energy_.toTensor().item(); + + torch::Tensor flat_force_ = force_.toTensor().view({-1}); + torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); + force.assign(cpu_force_.data_ptr(), + cpu_force_.data_ptr() + cpu_force_.numel()); + + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); } -template void DeepPotPT::compute(double& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box); +template void DeepPotPT::compute( + double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box); void DeepPotPT::get_type_map(std::string& type_map) { auto ret = module.run_method("get_type_map").toList(); for (const torch::IValue& element : ret) { - type_map += torch::str(element); // Convert each element to a string - type_map += " "; // Add a space between elements + type_map += torch::str(element); // Convert each element to a string + type_map += " "; // Add a space between elements } } @@ -169,7 +179,7 @@ void DeepPotPT::computew(std::vector& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam) { - //TODO: atomic compute unsupported + // TODO: atomic compute unsupported compute(ener, force, virial, coord, atype, box); } void DeepPotPT::computew(std::vector& ener, @@ -182,7 +192,7 @@ void DeepPotPT::computew(std::vector& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam) { - //TODO: atomic compute unsupported + // TODO: atomic compute unsupported compute(ener, force, virial, coord, atype, box); } void DeepPotPT::computew(std::vector& ener, @@ -198,8 +208,8 @@ void DeepPotPT::computew(std::vector& ener, const int& ago, const std::vector& fparam, const std::vector& aparam) { - //TODO: atomic compute unsupported - compute(ener, force, virial, coord, atype, box, inlist,ago); + // TODO: atomic compute unsupported + compute(ener, force, virial, coord, atype, box, inlist, ago); } void DeepPotPT::computew(std::vector& ener, std::vector& force, @@ -214,8 +224,8 @@ void DeepPotPT::computew(std::vector& ener, const int& ago, const std::vector& fparam, const std::vector& aparam) { - //TODO: atomic compute unsupported - compute(ener, force, virial, coord, atype, box, inlist,ago); + // TODO: atomic compute unsupported + compute(ener, force, virial, coord, atype, box, inlist, ago); } void DeepPotPT::computew_mixed_type(std::vector& ener, std::vector& force, @@ -228,8 +238,7 @@ void DeepPotPT::computew_mixed_type(std::vector& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam) { - throw; - + throw; } void DeepPotPT::computew_mixed_type(std::vector& ener, std::vector& force, diff --git a/source/api_cc/src/commonPT.cc b/source/api_cc/src/commonPT.cc index d4b3236064..aaa19b64f2 100644 --- a/source/api_cc/src/commonPT.cc +++ b/source/api_cc/src/commonPT.cc @@ -1,20 +1,24 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later #include "commonPT.h" -void NeighborListData::copy_from_nlist(const InputNlist& inlist, int& max_num_neighbors, int nnei) -{ +void NeighborListData::copy_from_nlist(const InputNlist& inlist, + int& max_num_neighbors, + int nnei) { int inum = inlist.inum; ilist.resize(inum); numneigh.resize(inum); memcpy(&ilist[0], inlist.ilist, inum * sizeof(int)); int* max_element = std::max_element(inlist.numneigh, inlist.numneigh + inum); max_num_neighbors = *max_element; - if (max_num_neighbors < nnei) + if (max_num_neighbors < nnei) { max_num_neighbors = nnei; + } jlist = (int*)malloc(inum * max_num_neighbors * sizeof(int)); - memset(jlist, -1 , inum * max_num_neighbors * sizeof(int)); + memset(jlist, -1, inum * max_num_neighbors * sizeof(int)); for (int ii = 0; ii < inum; ++ii) { int jnum = inlist.numneigh[ii]; numneigh[ii] = inlist.numneigh[ii]; - memcpy(&jlist[ii * max_num_neighbors], inlist.firstneigh[ii], jnum * sizeof(int)); + memcpy(&jlist[ii * max_num_neighbors], inlist.firstneigh[ii], + jnum * sizeof(int)); } -} \ No newline at end of file +}