Optimize GJK JIT time

[erikfrey]

• Could be some fixes on Warp side related to linker / .ptx file size
• Could be some work on our end to find a few spots in the kernel that we can break down

[c0d1f1ed]

Do you have some step-by-step instructions to run samples that see the slowest compile times? I've been working on some generic improvements that appear effective for other projects, that I'd like to test for MJWarp.

[thowell]

@c0d1f1ed

run testspeed

mjwarp-testspeed benchmark/aloha_pot/scene.xml --function=collision --clear_kernel_cache=True --nconmax=64 --njmax=512 --verbosity=1

output

...
Module ccd_kernel_builder__locals__ccd_kernel_e27ad55e e27ad55 load on device 'cuda:0' took 10091.27 ms  (compiled)
Module ccd_kernel_builder__locals__ccd_kernel_dad3c435 dad3c43 load on device 'cuda:0' took 11283.15 ms  (compiled)
Module ccd_kernel_builder__locals__ccd_kernel_d44e0c9d d44e0c9 load on device 'cuda:0' took 12242.70 ms  (compiled)
...

[c0d1f1ed]

Thanks @thowell! I'm running into an issue though:

...
       import mujoco
  File "C:\src\mujoco_warp\env\Lib\site-packages\mujoco\__init__.py", line 217, in <module>
    _load_all_bundled_plugins()
  File "C:\src\mujoco_warp\env\Lib\site-packages\mujoco\__init__.py", line 210, in _load_all_bundled_plugins
    PLUGIN_HANDLES.append(ctypes.CDLL(os.path.join(directory, filename)))
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Nicolas\miniconda3\Lib\ctypes\__init__.py", line 376, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: [WinError 1114] A dynamic link library (DLL) initialization routine failed

Looks like #2275.

[c0d1f1ed]

I got it up and running on Ubuntu 24.04. The first time I ran the benchmark I observed similar timings of ~10s per ccd_kernel. But every subsequent run they only take ~350ms. That's despite the --clear_kernel_cache=True. Manually clearing the Warp cache directory produced the same results.

I had to set CUDA_CACHE_DISABLE=1 or clear ~/.nv/ComputeCache, as described here, to get back to the worst-case compile times. I never observed the runtime's cache being effective...

Anyway, it's great to have a benchmark we can target for optimizing compile times!

[c0d1f1ed]

The use of Precompiled Headers (PCH) and avoiding aggressive NVRTC optimization with --Ofast-compile=min brings down the cold startup time from 32s to 22s.

I've seen cases where reducing the optimization level further (mid) provides bigger gains, but unfortunately it dramatically reduces kernel execution performance. --Ofast-compile=max often even causes compilation to fail. I'll see if we can make --Ofast-compile=min the default and if it should be under the application's control to make the tradeoffs...

What would help greatly is to compile multiple modules in parallel, instead of synchronously waiting for each module compilation during graph capture. While this may not be readily supported by CUDA graphs, we're considering adding functions to the Warp API to specify modules that provide kernels to be eagerly compiled in a multi-threaded way.