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

[Osamaali313]

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.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Updates the Python stub generator to also scan CUDA source/header files so that pybind-bound functions defined or declared in .cu/.cuh don’t fall back to overly-generic stubs.

Changes:

• Expand the C/C++ file extension sets to include .cu and .cuh when indexing functions.
• Expand the scan for m.def(...) statements to include .cu and .cuh.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.