feat: disable hashing for modules and kernels to allow loading from ptx and so

[YuyangLee]

Category

• New feature
• Bugfix
• Breaking change
• Refactoring
• Documentation
• Other (please explain)

Description

This PR tries to implement the feature of disabling hashing on .ptx and .so, so that compiled kernels in PTX and SO libs can be loaded without the source code of the kernel. The user sets the env variable WARP_DISABLE_HASHING_PREFIX to be the prefixes of modules for which the hashing is to be disabled (separated by a comma). For example: setting WARP_DISABLE_HASHING_PREFIX to "warp.tests.misc,warp.sim.dummy" will disable hashing for modules that starts with warp.sim.misc and warp.sim.dummy.

The main modifications are in warp/context.py, and an auxiliary unit test is in warp/tests/misc. The unit test can be run by the following script and is tested on my personal PC:

python -m warp.tests -s autodetect -p '*disable_hashing.py'

We have previously communicated with Mr. Miles about this feature and its downstream applications. We look forward to the Warp team's investigation.

Before your PR is "Ready for review"

• All commits are signed-off to indicate that your contribution adheres to the Developer Certificate of Origin requirements
• Necessary tests have been added
• Documentation is up-to-date
• Auto-generated files modified by compiling Warp and building the documentation have been updated (e.g. stubs.py, functions.rst)
• Code passes formatting and linting checks with pre-commit run -a

[mmacklin]

One problem I think we have to address is that modifying the kernel_cache_dir globally won't work well with other modules / libraries that are running in the same Python process.

I think what we really need here is a per-module cache path. e.g.:

wp.set_module_options({"kernel_cache_dir": "warp/tests/misc"})

We could combine this with the per-module disable_hashing / kernel_cache_force flag, so that users would need to do this at the top of their module:

wp.set_module_options({"kernel_cache_force": True, "kernel_cache_dir": "warp/tests/misc"})

[shi-eric]

Hi @YuyangLee, please sign off the commit as described by the GitHub action: https://github.com/NVIDIA/warp/pull/661/checks?check_run_id=40703124341

This is required before we can accept this.

[shi-eric]

@YuyangLee: I think if you can get this pull request to pass the DCO checker, we can handle making additional changes on top of your work to get it into Warp. We can have you review our proposed changes before they go in.

As for the changes we want to make, I think we will add some per-module options to specify the kernel output cache and whether hashing is disabled or not rather than relying on an undocumented environment variable.

I think it would also be helpful if you could remove the additions to warp/tests/misc/wp_warp.tests.misc.add_kernel/* and /.gitignoreso the only changes are forcontext.pyandtests/misc/*(which I recommend a more descriptive name, but we can take care of renames too). The reason is I think we can probably revise the test so that it can generate thewp_warp.tests.misc.add_kernel.sm86.ptx ` on-the-fly rather than having it in Git version control. But if you've already tried this and ran into difficulties, I'd like to hear about them.

Thanks!

[YuyangLee]

@shi-eric:

Hi Eric, sorry for the late reply! I have fixed the signoffs, and now the commits are signed by my GPG keys.

Also, I have removed the compiled caches and the .gitignore lines according to your suggestions. I'm not quite sure what the best way is to perform the on-the-fly compilation and test the loading-from-cache function, as my previous test was to first call the kernel to produce the caches, then remove the kernel function contents and add the environment variable to test the loading. I have left a TODO in warp/tests/misc/test_disable_hashing.py and maybe you can modify the test.

Thanks a lot!