Skip to content

Support caching CUBINs generated with name_expressions - #9801

Merged
leofang merged 34 commits into
cupy:mainfrom
leofang:claude/investigate-caching-kernels-mangled-names
May 7, 2026
Merged

leofang merged 34 commits into
cupy:mainfrom
leofang:claude/investigate-caching-kernels-mangled-names

Conversation

@leofang

@leofang leofang commented Mar 13, 2026 •

Copy link
Copy Markdown
Member

Close #9667. Close #6159.

This helps reduce test time greatly. For example with test_mdspan.py:

main this PR
1st run (cold cache) 66s 66s
2nd run (warm cache) 36s 3s

Claude AI and others added 20 commits March 12, 2026 04:31
- Add cuModuleEnumerateFunctions and cuFuncGetName bindings (CUDA 11.6+)
- Implement C++ demangling using __cxa_demangle in function.pyx
- Add _enumerate_and_build_mapping to Module class
- Update compiler.py to use enumeration on cache hit for both CUDA and HIP
- HIP/ROCm: Raise NotImplementedError if enumeration not supported

Co-authored-by: leofang <[email protected]>
- Test cache hit with name_expressions
- Test C++ demangling and name matching with whitespace variations
- Verify mapping is built from cached CUBIN via enumeration APIs
- Skip on HIP/ROCm as it doesn't support function enumeration yet

Co-authored-by: leofang <[email protected]>
…evert HIP changes

- Replace __cxa_demangle with __cu_demangle from nv_decode.h and link libcufilt
- Change name_expressions parameter from list to tuple throughout
- Revert HIP enumeration changes, add TODO comment instead
- Functions are available in cuda.bindings (CUDA 11.6+ standard APIs)

Co-authored-by: leofang <[email protected]>
- Include name_expressions in cache hash key (sorted, comma-separated)
- Explicitly unload module if enumeration fails to prevent CUBIN leak
- Move libcufilt from libraries to static_libraries (it's a static lib)
- Clarify comment: return mangled name on demangle failure for exact matching

Co-authored-by: leofang <[email protected]>
- Change moduleEnumerateFunctions to cdef returning std::vector<Function>
- Change funcGetName to cdef returning char*
- Remove match_name_expression, normalize_name, and demangle_cxx_name functions
- Simplify _enumerate_and_build_mapping to use identity mapping
- Remove unused imports (re, libc.stdlib.free)

Co-authored-by: leofang <[email protected]>
- Restore demangling code and name matching logic in function.pyx
- Swap HIP check before mapping check in _enumerate_and_build_mapping
- Change except * to except +* in moduleEnumerateFunctions
- Add back re import and libc.stdlib.free import

Co-authored-by: leofang <[email protected]>
- Add cdef declarations for moduleEnumerateFunctions and funcGetName in driver.pxd
- Simplify _enumerate_and_build_mapping to build demangled->mangled mapping directly
- Remove match_name_expression and normalize_name functions (no longer needed)
- Remove re import (no longer used)

Co-authored-by: leofang <[email protected]>
- Wrap nv_decode.h extern declaration with IF CUPY_HIP_VERSION == 0
- Guard demangle_cxx_name function body for CUDA vs HIP
- On HIP, return mangled name as-is since libcufilt is not available

Co-authored-by: leofang <[email protected]>
- Change demangle_cxx_name to take const char* instead of str
- Avoid encode→decode round-trip by passing char* directly
- Decode only once when building mapping
- Remove TestNameExpressionsCaching test class (tests unreliable due to multiple caching levels)

Co-authored-by: leofang <[email protected]>
- Change IF CUPY_HIP_VERSION == 0 to IF CUPY_CUDA_VERSION > 0
- Move cdef declarations outside try block to fix Cython compilation error
- "cdef statement not allowed here" error resolved

Co-authored-by: leofang <[email protected]>
- Change except +* to except + in moduleEnumerateFunctions (Cython syntax)
- Add _enumerate_and_build_mapping declaration to function.pxd
- Fixes compiler crash and missing method declaration errors

Co-authored-by: leofang <[email protected]>
@leofang

leofang commented Mar 13, 2026

Copy link
Copy Markdown
Member Author

/test mini

@leofang leofang added cat:enhancement Improvements to existing features to-be-backported Pull-requests to be backported to stable branch labels Mar 13, 2026
@leofang

leofang commented Mar 13, 2026

Copy link
Copy Markdown
Member Author

/test mini

@leofang

leofang commented Mar 13, 2026

Copy link
Copy Markdown
Member Author

/test mini,linux

leofang and others added 2 commits April 29, 2026 21:58
- Use r""" raw string for the cufilt C code block (reduces backslash escaping)
- Use Py_ssize_t for loop/count variables in _build_mangled_name_mapping
- Refactor moduleEnumerateFunctions to take vector as output parameter
  (avoids except+ on vector return; uses except?-1 per project convention)
- Drop except+ from moduleGetFunctionCount (default except* suffices)
- Update .pxd declarations to match

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
@leofang

leofang commented May 1, 2026

Copy link
Copy Markdown
Member Author

/test mini

@leofang

leofang commented May 1, 2026

Copy link
Copy Markdown
Member Author

This PR is ready for review. The Windows failure is already fixed in #9902.

@seberg seberg left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generally, looks good, I am still a bit wondering if integrating it into the file isn't better.
However, this doesn't block us from doing that in the future if we get into problems and it should be a pretty big improvement!

I'll be curious if we see speed improvements in CI here on re-runs now (that would be awesome and also a decent test that it works).

Claude found some nitpicks that I agree with, though (to me, I wonder if the try/except should be on the callers around self._enumerate_and_build_mapping() but either way seems fine).

Comment thread cupy/cuda/function.pyx Resolved
Comment thread cupy/cuda/function.pyx Resolved
Comment thread cupy/cuda/function.pyx Resolved
leofang and others added 2 commits May 6, 2026 21:54
- Thread safety: use std::call_once for Windows cufilt DLL init
  (eliminates TOCTOU race on _cupy_cufilt_loaded flag)
- Wide-char APIs: switch Windows path handling from ANSI to Wide
  (GetModuleFileNameW, LoadLibraryW, etc.) to handle paths beyond
  MAX_PATH; also check return value for truncation
- Incomplete mapping guard: only accept the demangled mapping when
  all name_expressions were matched (len check); if incomplete,
  fall through to recompilation which produces proper errors

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
@leofang leofang closed this May 7, 2026
@leofang leofang reopened this May 7, 2026
@seberg

seberg commented May 7, 2026

Copy link
Copy Markdown
Member

Close/re-open again, let's see if readthedocs recovers...

@seberg seberg closed this May 7, 2026
@seberg seberg reopened this May 7, 2026
@seberg

seberg commented May 7, 2026

Copy link
Copy Markdown
Member

/test mini

@seberg seberg left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, nice to use an output parameter for the function return there. Let's give this a shot (I am curious if it affects CI).

@leofang

leofang commented May 7, 2026

Copy link
Copy Markdown
Member Author

Test failure is irrelevant.

/test force-skip

@leofang
leofang enabled auto-merge May 7, 2026 12:26
@github-actions

github-actions Bot commented May 7, 2026

Copy link
Copy Markdown

The following tests were force-skipped:

@leofang
leofang merged commit 913572e into cupy:main May 7, 2026
76 of 77 checks passed
@leofang

leofang commented May 7, 2026

Copy link
Copy Markdown
Member Author

(I am curious if it affects CI).

Let's monitor the main branch. The per-PR CI runs don't seem to get faster much.

leofang added a commit that referenced this pull request May 8, 2026
…els-mangled-names

Support caching CUBINs generated with `name_expressions`
@kmaehashi kmaehashi added this to the v15.0.0a1 milestone May 21, 2026
@leofang
leofang deleted the claude/investigate-caching-kernels-mangled-names branch August 8, 2026 07:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cat:enhancement Improvements to existing features to-be-backported Pull-requests to be backported to stable branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Load kernels with mangled name from cache

5 participants