[Feature Request] Code and build system for multiple backends

[njzjz]

Summary

This proposal describes how multiple backends are organized and built.

Detailed Description

Code structure

The code may organized as

(updated: Jan 27, 2024)

- deepmd
  - tf
  - pt
- source
  - api_cc (include tf and pt)
  - tests
     - common
     - tf
     - pt

deepmd is actually deepmd_tf, but we must keep compatibility.

Other codes are not affected.

CMake

Currently, CMake has TENSORFLOW_ROOT and USE_TF_PYTHON_LIBS. Now they are not required if one enables USE_PT_PYTHON_LIBS. However, an error should be raised if one does not enable any backend.
One or two backends can be built, as shown below:

Pip

By default, both backends are installed. This doesn't bring great changes because PT codes don't use C++.
Add a new environment variable option to disable TF.
Add optional dependencies for PyTorch, i.e.

pip install deepmd-kit[torch]

Pre-built packages with dependencies

As both TensorFlow and PyTorch are pretty large in size, we need to determine whether to bring them into the same pre-built package...

Further Information, Files, and Links

No response

[njzjz]

Pre-built packages with dependencies

As both TensorFlow and PyTorch are pretty large in size, we need to determine whether to bring them into the same pre-built package...

After investigation, the pre-built package seems impossible due to:

(1) PyTorch from the PyPI has incompatible CXX11_ABI with TF;
(2) PyTorch package built using CXX11_ABI=1 on their website has a higher GLIBC version, incompatible with the current docker image to make the pre-built package;
(3) PyTorch links the CUDA libraries instead of using dlopen(). This means we need to ship CUDA libraries when providing it... Otherwise, it cannot work even on CPUs.
(4) It's too big!

-rwxr-xr-x. 1 jz748 jz748  55M Dec 12 13:10 libnvrtc-b51b459d.so.12
-rwxr-xr-x. 1 jz748 jz748  68M Dec 12 13:10 libcudnn_ops_train.so.8
-rwxr-xr-x. 1 jz748 jz748  87M Dec 12 13:10 libcudnn_ops_infer.so.8
-rwxr-xr-x. 1 jz748 jz748 103M Dec 12 13:09 libcublas-37d11411.so.12
-rwxr-xr-x. 1 jz748 jz748 111M Dec 12 13:10 libcudnn_adv_train.so.8
-rwxr-xr-x. 1 jz748 jz748 120M Dec 12 13:10 libcudnn_adv_infer.so.8
-rwxr-xr-x. 1 jz748 jz748 127M Dec 12 13:10 libcudnn_cnn_train.so.8
-rwxr-xr-x. 1 jz748 jz748 129M Dec 12 13:10 libtorch_cuda_linalg.so
-rw-r--r--. 1 jz748 jz748 140M Dec 12 13:03 libdnnl.a
-rwxr-xr-x. 1 jz748 jz748 427M Dec 12 13:10 libtorch_cpu.so
-rwxr-xr-x. 1 jz748 jz748 492M Dec 12 13:08 libcublasLt-f97bfc2c.so.12
-rwxr-xr-x. 1 jz748 jz748 605M Dec 12 13:10 libcudnn_cnn_infer.so.8
-rwxr-xr-x. 1 jz748 jz748 1.3G Dec 12 13:10 libtorch_cuda.so