Skip to content

Optimization: refactor some memory allocation#1356

Merged
thowell merged 4 commits into
google-deepmind:mainfrom
Kenny-Vilella:dev/kvilella/optimize_memset
Jun 3, 2026
Merged

Optimization: refactor some memory allocation#1356
thowell merged 4 commits into
google-deepmind:mainfrom
Kenny-Vilella:dev/kvilella/optimize_memset

Conversation

@Kenny-Vilella

@Kenny-Vilella Kenny-Vilella commented May 14, 2026

Copy link
Copy Markdown
Collaborator

Removed several zeroing in the sparse codepath.

Results on my RTX4080 running the benchmark 5 times and taking the median:

Optimization Main Median Branch Median Delta (%)
aloha_pot 1554434.93 1597390.69 +2.76
aloha_sdf 219188.60 223895.40 +2.15
aloha_cloth 81.48 83.82 +2.87
aloha_clutter 5863.09 6151.17 +4.91
cloth 771.69 765.44 -0.81
franka_emika_panda 15530391.57 15744625.83 +1.38
humanoid 3151787.27 3273520.82 +3.86
three_humanoids 509640.79 561189.55 +10.11
myoarm 792295.43 815093.72 +2.88
unitree_g1_flat 1346559.18 1403586.08 +4.24
unitree_g1_hfield 1000116.13 1058186.62 +5.81

Biggest win is in the sparse case but still some small win in the dense case.

@Kenny-Vilella

Kenny-Vilella commented May 14, 2026

Copy link
Copy Markdown
Collaborator Author

I still need to double check that we do not have any problems in all codepaths.
Will also run compute-sanitizer.

Kenny-Vilella
Kenny-Vilella commented May 14, 2026
View reviewed changes
Comment thread mujoco_warp/_src/solver.py Outdated
@Kenny-Vilella

Kenny-Vilella commented May 15, 2026

Copy link
Copy Markdown
Collaborator Author

I removed the dangerous trick and re-benchmark on L40 (more consistent results there).

Optimization Main Median Branch Median Delta (%)
aloha_pot 2083027.4577659 2094684.60997266 +0.56
aloha_sdf 353151.503929575 353761.276640722 +0.17
aloha_cloth 92.3157749372582 92.3501486093072 +0.04
aloha_clutter 9120.46618748854 9328.35970860687 +2.23
cloth 974.509169894073 975.156182326258 +0.07
franka_emika_panda 18547427.11419 18678812.8518133 +0.70
humanoid 4033081.4767262 4140494.60584687 +2.59
three_humanoids 681013.786672344 731548.288131534 +6.91
myoarm 1015078.63538665 1041717.72597511 +2.56
unitree_g1_flat 1678697.80799625 1745439.46772209 +3.82
unitree_g1_hfield 1245721.74167735 1280713.69450947 +2.73

@thowell

thowell commented May 15, 2026

Copy link
Copy Markdown
Collaborator

lgtm, thanks @Kenny-Vilella!

@Kenny-Vilella

Kenny-Vilella commented May 15, 2026

Copy link
Copy Markdown
Collaborator Author

Thanks @thowell !
I will do some compute sanitizer runs on Monday to be 100% sure then will switch to review.

@Kenny-Vilella Kenny-Vilella changed the title [WIP] Optimization: refactor some memory allocation Optimization: refactor some memory allocation May 18, 2026
@Kenny-Vilella Kenny-Vilella marked this pull request as ready for review May 18, 2026 05:13
@Kenny-Vilella

Kenny-Vilella commented May 18, 2026

Copy link
Copy Markdown
Collaborator Author

@thowell Did the compute sanitizer run on test and a selection of benchmark, everything came clean.
It should be ready to be merged.

@thowell thowell merged commit be83026 into google-deepmind:main Jun 3, 2026
14 checks passed
@adenzler-nvidia adenzler-nvidia mentioned this pull request Jun 3, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@thowell thowell thowell approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Kenny-Vilella @thowell