~10% speed up of the kernel at batch sizes 200-300
This is thanks to an improvement of the memory throughput as measured by Nsight-compute. Tried also doing warp swaps (each y block load WARP_SIZE atoms) in the y block, but that actually reduced the performance of the kernel to my surprise.
Benchmarks
There doesn't appear to be a significant improvement to the ns/day benchmarks, but do see a ~10% speed up in the water sampling example.
NS/Day Benchmarks
Only looking at the hif2a case, since it is the only benchmark with water sampling
2 / 10000 | density 991.6458074986836 | # of waters in spherical region 4 | md step: 0 | time per mc move: 4109.1798ns
4 / 20000 | density 990.2169202722931 | # of waters in spherical region 2 | md step: 10000 | time per mc move: 4040.6117ns
4 / 30000 | density 992.2573134334189 | # of waters in spherical region 2 | md step: 20000 | time per mc move: 3864.6266ns
4 / 40000 | density 987.5691282363993 | # of waters in spherical region 2 | md step: 30000 | time per mc move: 3865.3837ns
4 / 50000 | density 994.4733590379498 | # of waters in spherical region 2 | md step: 40000 | time per mc move: 3860.3977ns
PR
2 / 10000 | density 991.6458074986836 | # of waters in spherical region 4 | md step: 0 | time per mc move: 3741.9775ns
4 / 20000 | density 990.2169202722931 | # of waters in spherical region 2 | md step: 10000 | time per mc move: 3675.4506ns
4 / 30000 | density 992.2573134334189 | # of waters in spherical region 2 | md step: 20000 | time per mc move: 3516.9444ns
4 / 40000 | density 987.5691282363993 | # of waters in spherical region 2 | md step: 30000 | time per mc move: 3515.0994ns
4 / 50000 | density 994.4733590379498 | # of waters in spherical region 2 | md step: 40000 | time per mc move: 3511.0445ns
Kernel Timings
Looking at the output of the following command
nsys profile --stats=true -t nvtx,cuda python -u examples/water_sampling_mc.py --water_pdb timemachine/datasets/water_exchange/bb_6_waters.pdb --ligand_sdf timemachine/datasets/water_exchange/bb_centered_espaloma.sdf --out_cif traj_6_waters.cif --iterations 1 --equilibration_steps 0 --md_steps_per_batch 0 --mc_steps_per_batch 10000 --insertion_type targeted --use_hmr 1 --batch_size 250
k_atom_by_atom_energies went from ~250k ns to ~210k ns by improving the memory throughput.
Master
CUDA Kernel Statistics:
Time (%) Total Time (ns) Instances Avg (ns) Med (ns) Min (ns) Max (ns) StdDev (ns) Name
-------- --------------- --------- --------- --------- -------- -------- ----------- ----------------------------------------------------------------------------------------------------
53.4 25213920 98 257284.9 257227.5 254971 258428 470.6 void timemachine::k_atom_by_atom_energies<float>(int, int, const int *, const double *, const doubl…
8.0 3761820 49 76771.8 76735.0 76126 78143 373.1 void timemachine::k_adjust_energies<float, (bool)1>(int, int, int, int, const int *, const int *, c…
8.0 3760003 49 76734.8 76703.0 76223 77598 303.1 void timemachine::k_adjust_energies<float, (bool)0>(int, int, int, int, const int *, const int *, c…
7.0 3284267 49 67025.9 67007.0 66751 67327 154.6 void timemachine::k_set_sampled_energy_block<float, (int)512>(int, int, int, int, const int *, cons…
4.5 2127321 49 43414.7 41151.0 40927 54047 4835.1 timemachine::k_store_exchange_move(int, int, const int *, const int *, const int *, const int *, co…
4.1 1917985 1 1917985.0 1917985.0 1917985 1917985 0.0 void timemachine::k_compute_nonbonded_target_atom_energies<float, (int)64>(int, int, const int *, c…
2.4 1128244 99 11396.4 3744.0 3104 20352 8225.1 void timemachine::k_convert_energies_to_log_weights<float>(int, T1, const __int128 *, T1 *)
1.2 575702 4 143925.5 128669.5 128542 189821 30597.2 void generate_seed_pseudo<rng_config<curandStateXORWOW, (curandOrdering)101>>(unsigned long long, u…
1.2 553045 98 5643.3 5600.0 5535 6592 124.4 void timemachine::k_segmented_exp_sub_max<float>(int, const int *, const T1 *, const T1 *, T1 *)
0.9 414041 49 8449.8 8448.0 7584 8896 226.0 void timemachine::k_setup_destination_weights_for_targeted<float>(int, int, const int *, const int …
<truncated>
Benchmarks
There doesn't appear to be a significant improvement to the ns/day benchmarks, but do see a ~10% speed up in the water sampling example.
NS/Day Benchmarks
Only looking at the hif2a case, since it is the only benchmark with water sampling
Master
PR
Water Sampling Example
Result of running
python -u examples/water_sampling_mc.py --water_pdb timemachine/datasets/water_exchange/bb_6_waters.pdb --ligand_sdf timemachine/datasets/water_exchange/bb_centered_espaloma.sdf --out_cif traj_6_waters.cif --md_steps_per_batch 10000 --mc_steps_per_batch 10000 --insertion_type targeted --use_hmr 1 --batch_size 250Time per mc move is reduced by ~9%
Master
PR
Kernel Timings
Looking at the output of the following command
nsys profile --stats=true -t nvtx,cuda python -u examples/water_sampling_mc.py --water_pdb timemachine/datasets/water_exchange/bb_6_waters.pdb --ligand_sdf timemachine/datasets/water_exchange/bb_centered_espaloma.sdf --out_cif traj_6_waters.cif --iterations 1 --equilibration_steps 0 --md_steps_per_batch 0 --mc_steps_per_batch 10000 --insertion_type targeted --use_hmr 1 --batch_size 250k_atom_by_atom_energieswent from ~250k ns to ~210k ns by improving the memory throughput.Master
PR