pt: support dpa2 model parallel inference

[CaRoLZhangxy]

No description provided.

[CaRoLZhangxy]

TODO: add ut

[iProzd]

Docstr of comm_dict missed in all the python classes.

no need

[chazeon]

I am unable to compile comm.cc after this pull request. The error I get is

-> % cmake --build . --parallel 8
[  3%] Built target deepmd_ipi
[  5%] Built target gtest
[ 19%] Built target deepmd
[ 20%] Built target gtest_main
[ 21%] Built target gmock
[ 22%] Built target gmock_main
[ 28%] Built target deepmd_cc_test_no_backend
[ 46%] Built target runUnitTests_lib
[ 52%] Built target deepmd_cc
[ 73%] Built target deepmd_op
[ 74%] Linking CXX shared module libdeepmd_op_pt.so
[ 76%] Built target deepmd_c
[ 88%] Built target runUnitTests_cc
[ 89%] Built target deepmd_gromacs
[ 90%] Built target dp_ipi
Undefined symbols for architecture arm64:
  "long* at::TensorBase::data_ptr<long>() const", referenced from:
      Border::unpack_communicator(at::Tensor const&, int&) in comm.cc.o
ld: symbol(s) not found for architecture arm64
clang-16: error: linker command failed with exit code 1 (use -v to see invocation)
[ 99%] Built target runUnitTests_c
make[2]: *** [op/pt/CMakeFiles/deepmd_op_pt.dir/build.make:187: op/pt/libdeepmd_op_pt.so] Error 1
make[1]: *** [CMakeFiles/Makefile2:549: op/pt/CMakeFiles/deepmd_op_pt.dir/all] Error 2
make: *** [Makefile:146: all] Error 2

The error is related to these two lines:

https://github.com/CaRoLZhangxy/deepmd-kit/blob/beba142f64a51a9d6e65d6a621d30519679aa43d/source/op/pt/comm.cc#L323-L326

It seems at::TensorBase::data_ptr<long>() is not defined with long, but for int, long long etc in my libtorch.
I tried replacing long int with either long long or int. Both can be compiled.
Is there a specific reason why long int is used here?

-> % nm -gU libtorch_cpu.dylib | grep TensorBase | llvm-cxxfilt | grep ::data_ptr
0000000002153398 T c10::Float8_e5m2* at::TensorBase::data_ptr<c10::Float8_e5m2>() const
0000000002153524 T c10::Float8_e4m3fn* at::TensorBase::data_ptr<c10::Float8_e4m3fn>() const
0000000002152738 T c10::Half* at::TensorBase::data_ptr<c10::Half>() const
00000000021536b0 T c10::qint8* at::TensorBase::data_ptr<c10::qint8>() const
00000000021539c8 T c10::qint32* at::TensorBase::data_ptr<c10::qint32>() const
000000000215383c T c10::quint8* at::TensorBase::data_ptr<c10::quint8>() const
0000000002152bdc T c10::complex<c10::Half>* at::TensorBase::data_ptr<c10::complex<c10::Half>>() const
0000000002152ef4 T c10::complex<double>* at::TensorBase::data_ptr<c10::complex<double>>() const
0000000002152d68 T c10::complex<float>* at::TensorBase::data_ptr<c10::complex<float>>() const
000000000215320c T c10::BFloat16* at::TensorBase::data_ptr<c10::BFloat16>() const
0000000002153ce0 T c10::quint2x4* at::TensorBase::data_ptr<c10::quint2x4>() const
0000000002153b54 T c10::quint4x2* at::TensorBase::data_ptr<c10::quint4x2>() const
0000000002152108 T signed char* at::TensorBase::data_ptr<signed char>() const
0000000002153080 T bool* at::TensorBase::data_ptr<bool>() const
0000000002152a50 T double* at::TensorBase::data_ptr<double>() const
00000000021528c4 T float* at::TensorBase::data_ptr<float>() const
0000000002151f7c T unsigned char* at::TensorBase::data_ptr<unsigned char>() const
0000000002152420 T int* at::TensorBase::data_ptr<int>() const
0000000002152294 T short* at::TensorBase::data_ptr<short>() const
00000000021525ac T long long* at::TensorBase::data_ptr<long long>() const

My environment is:

     active environment : deepmd-dev
    active env location : /Users/chazeon/miniforge3/envs/deepmd-dev
            shell level : 1
       user config file : /Users/chazeon/.condarc
 populated config files : /Users/chazeon/miniforge3/.condarc
                          /Users/chazeon/.condarc
          conda version : 23.3.1
    conda-build version : not installed
         python version : 3.10.12.final.0
       virtual packages : __archspec=1=arm64
                          __osx=14.4.1=0
                          __unix=0=0
       base environment : /Users/chazeon/miniforge3  (writable)
      conda av data dir : /Users/chazeon/miniforge3/etc/conda
  conda av metadata url : None
           channel URLs : https://conda.anaconda.org/conda-forge/osx-arm64
                          https://conda.anaconda.org/conda-forge/noarch
                          https://repo.anaconda.com/pkgs/main/osx-arm64
                          https://repo.anaconda.com/pkgs/main/noarch
                          https://repo.anaconda.com/pkgs/r/osx-arm64
                          https://repo.anaconda.com/pkgs/r/noarch
          package cache : /Users/chazeon/miniforge3/pkgs
                          /Users/chazeon/.conda/pkgs
       envs directories : /Users/chazeon/miniforge3/envs
                          /Users/chazeon/.conda/envs
               platform : osx-arm64
             user-agent : conda/23.3.1 requests/2.31.0 CPython/3.10.12 Darwin/23.4.0 OSX/14.4.1
                UID:GID : 501:20
             netrc file : None
           offline mode : False

This might related to: https://stackoverflow.com/questions/67584843/pytorch-tensordata-ptrlong-long-not-working-on-linux

[njzjz]

@chazeon could you send a PR?

In a 64-bit system, long and long long have the same size, which might be why long can be compiled on Linux. But it looks dangerous.

[njzjz]

long used here also requires to be reviewed:

deepmd-kit/source/api_cc/src/DeepPotPT.cc

Lines 199 to 213 in ebd809b

	if (!fparam.empty()) {
	fparam_tensor =
	torch::from_blob(const_cast<VALUETYPE*>(fparam.data()),
	{1, static_cast<long int>(fparam.size())}, options)
	.to(device);
	}
	c10::optional<torch::Tensor> aparam_tensor;
	if (!aparam_.empty()) {
	aparam_tensor = torch::from_blob(
	const_cast<VALUETYPE*>(aparam_.data()),
	{1, lmp_list.inum,
	static_cast<long int>(aparam_.size()) / lmp_list.inum},
	options)
	.to(device);
	}