Function extract_function_signatures

export_hf.py:159–206

Find all torch::Tensor-returning function declarations in the CUDA source. Returns a list of dicts with keys: - 'return_type': e.g. 'torch::Tensor' - 'name': e.g. 'matmul_cuda' - 'params': e.g. 'torch::Tensor A, torch::Tensor B' - 'full_signature': the complete decl

(cuda_src: str)

# ---------------------------------------------------------------------------

def extract_function_signatures(cuda_src: str) -> List[Dict[str, str]]:
    """
    Find all torch::Tensor-returning function declarations in the CUDA source.

    Returns a list of dicts with keys:
      - 'return_type': e.g. 'torch::Tensor'
      - 'name': e.g. 'matmul_cuda'
      - 'params': e.g. 'torch::Tensor A, torch::Tensor B'
      - 'full_signature': the complete declaration

    Only extracts non-kernel functions (i.e., the C++ launcher functions that
    PyTorch binds to, not __global__ CUDA kernels).
    """
    # Match: torch::Tensor func_name(params) {
    # Also match at::Tensor, std::vector<torch::Tensor>, void
    pattern = (
        r"^((?:torch::Tensor|at::Tensor|std::vector<torch::Tensor>|void)\s+"
        r"(\w+)\s*\(([^)]*)\))\s*\{"
    )

    results = []
    for match in re.finditer(pattern, cuda_src, re.MULTILINE):
        full_sig = match.group(1).strip()
        func_name = match.group(2)
        params = match.group(3).strip()
        return_type = full_sig.split(func_name)[0].strip()

        # Skip __global__ kernels (they are called from launchers, not from Python)
        # Check if the line before has __global__
        start = match.start()
        preceding = cuda_src[max(0, start - 200) : start]
        if "__global__" in preceding.split("\n")[-1] if preceding else "":
            continue

        # Also skip if the function name suggests it's a device helper
        if func_name.startswith("__"):
            continue

        results.append(
            {
                "return_type": return_type,
                "name": func_name,
                "params": params,
                "full_signature": full_sig,
            }
        )

    return results


def _parse_param_list(params_str: str) -> List[Tuple[str, str]]: