add JIT option

[njzjz]

Here are my results (2080 Super, TF 2.9, CUDA 11.6):

(unit: ms/step)

	no JIT	JIT
water tanh fp32	8.45	7.78
water tanh fp64	33.37	34.09
water gelu fp32	19.69	14.23
water gelu fp64	48.39	49.19

Conclusion:

1. TF JIT only works with FP32 and slows down FP64 (or my card only has good FP32 performance?).
2. JIT has a great performance on TF's GeLU.
3. JIT only works for static shapes, so training with different natoms (or different atom types) needs a lot of compilation time (not sure if it's worth it).
4. JIT only works for a GPU.

Usage

export DP_JIT=1

(It sets TF_XLA_FLAGS=--tf_xla_auto_jit=2 but DP_JIT=1 is easier to remember)

[codecov-commenter]

Codecov Report

Base: 75.65% // Head: 75.66% // Increases project coverage by +0.00% 🎉

Coverage data is based on head (e6a4fb5) compared to base (3c32995).
Patch coverage: 50.00% of modified lines in pull request are covered.

Additional details and impacted files

@@           Coverage Diff           @@
##            devel    #1730   +/-   ##
=======================================
  Coverage   75.65%   75.66%           
=======================================
  Files         204      204           
  Lines       20258    20260    +2     
  Branches     1438     1438           
=======================================
+ Hits        15327    15329    +2     
- Misses       4025     4026    +1     
+ Partials      906      905    -1     

Impacted Files	Coverage Δ
deepmd/env.py	72.18% <50.00%> (-0.34%)	⬇️
source/lib/include/ComputeDescriptor.h	77.09% <0.00%> (+0.22%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[njzjz]

@wanghan-iapcm, any reviews for this PR?