[PR #361] fix: index CUDA sources/headers when generating .pyi stubs #391

Open
opened 2026-10-26 21:17:10 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/361
Author: @Osamaali313
Created: 6/14/2026
Status: 🔄 Open

Base: main ← Head: fix/pyi-index-cuda-sources


📝 Commits (1)

  • 108820b fix: index CUDA sources/headers when generating .pyi stubs

📊 Changes

1 file changed (+6 additions, -2 deletions)

View changed files

📝 scripts/generate_pyi.py (+6 -2)

📄 Description

Problem

scripts/generate_pyi.py builds the type stubs (stubs/_C.pyi) consumed at build time by setup.py (generate_pyi_file(name='_C', root='./csrc', ...)). It locates each bound function's C++ signature by scanning source files, but the extension allowlists omit CUDA sources/headers:

# build_cpp_function_index
extensions = {'.cpp', '.cc', '.cxx', '.c', '.hpp', '.h'}      # no .cu / .cuh
# extract_m_def_statements
extensions = {'.hpp', '.cpp', '.h', '.cc'}                    # no .cu / .cuh

So any m.def-bound function whose C++ definition/declaration lives in a .cu or .cuh is not found, and its stub falls back to a generic def name(*args, **kwargs) -> Any: ... (the code even logs Warning: C++ function "<name>" not found in any .cpp file).

This is also internally inconsistent: build_cpp_function_index already classifies .cuh as a header in its is_header check —

is_header = file_path.suffix.lower() in {'.h', '.hpp', '.cuh'}

— but never reads .cuh files, so that branch is unreachable.

Fix

Add .cu and .cuh to both allowlists.

Validation

Ran the generator's parser over ./csrc before and after:

  • No regression on the current tree: all 38 resolvable bound signatures still resolve (the 12 lambdas are unaffected). The repo's only .cu file, csrc/indexing/main.cu, registers no bindings, so output is byte-identical today.
  • Defect reproduced and fixed with a minimal case — a binding registered in a .hpp whose C++ function is declared in a sibling .cuh:
    • before: my_kernel not indexed → Warning: ... not found in any .cpp file → generic *args stub
    • after: resolves to torch::Tensor my_kernel(torch::Tensor x) (typed stub)
    • same for a .cu-defined binding → int cu_fn(int a, int b)

The change is forward-looking: as more kernels/bindings move into .cu/.cuh, their public stubs stay typed instead of silently degrading to Any.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/361 **Author:** [@Osamaali313](https://github.com/Osamaali313) **Created:** 6/14/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/pyi-index-cuda-sources` --- ### 📝 Commits (1) - [`108820b`](https://github.com/deepseek-ai/DeepGEMM/commit/108820b795bd663a285401950a01a47c1dbcc7da) fix: index CUDA sources/headers when generating .pyi stubs ### 📊 Changes **1 file changed** (+6 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `scripts/generate_pyi.py` (+6 -2) </details> ### 📄 Description ## Problem `scripts/generate_pyi.py` builds the type stubs (`stubs/_C.pyi`) consumed at build time by `setup.py` (`generate_pyi_file(name='_C', root='./csrc', ...)`). It locates each bound function's C++ signature by scanning source files, but the extension allowlists omit CUDA sources/headers: ```python # build_cpp_function_index extensions = {'.cpp', '.cc', '.cxx', '.c', '.hpp', '.h'} # no .cu / .cuh # extract_m_def_statements extensions = {'.hpp', '.cpp', '.h', '.cc'} # no .cu / .cuh ``` So any `m.def`-bound function whose C++ definition/declaration lives in a `.cu` or `.cuh` is **not found**, and its stub falls back to a generic `def name(*args, **kwargs) -> Any: ...` (the code even logs `Warning: C++ function "<name>" not found in any .cpp file`). This is also internally inconsistent: `build_cpp_function_index` already classifies `.cuh` as a header in its `is_header` check — ```python is_header = file_path.suffix.lower() in {'.h', '.hpp', '.cuh'} ``` — but never *reads* `.cuh` files, so that branch is unreachable. ## Fix Add `.cu` and `.cuh` to both allowlists. ## Validation Ran the generator's parser over `./csrc` before and after: - **No regression on the current tree:** all 38 resolvable bound signatures still resolve (the 12 lambdas are unaffected). The repo's only `.cu` file, `csrc/indexing/main.cu`, registers no bindings, so output is byte-identical today. - **Defect reproduced and fixed** with a minimal case — a binding registered in a `.hpp` whose C++ function is declared in a sibling `.cuh`: - before: `my_kernel` not indexed → `Warning: ... not found in any .cpp file` → generic `*args` stub - after: resolves to `torch::Tensor my_kernel(torch::Tensor x)` (typed stub) - same for a `.cu`-*defined* binding → `int cu_fn(int a, int b)` The change is forward-looking: as more kernels/bindings move into `.cu`/`.cuh`, their public stubs stay typed instead of silently degrading to `Any`. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepGEMM#391
No description provided.