Function _export_cuda_kernel

export_hf.py:492–580

Export a CUDA C++ kernel to HF Kernels format.

(
    source: str,
    name: str,
    output_dir: str,
    repo_id: str,
)

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

def _export_cuda_kernel(
    source: str,
    name: str,
    output_dir: str,
    repo_id: str,
) -> None:
    """Export a CUDA C++ kernel to HF Kernels format."""

    # Extract the CUDA source string
    cuda_src = extract_cuda_source(source)
    if cuda_src is None:
        print("ERROR: Could not extract CUDA_SRC from kernel file.")
        print("       Expected a CUDA_SRC = r\"\"\"...\"\"\" string assignment.")
        sys.exit(1)

    # Parse function signatures from the CUDA source
    functions = extract_function_signatures(cuda_src)
    if not functions:
        # Try to detect from compile_cuda call
        func_name = extract_function_name_from_compile(source)
        if func_name:
            print(
                f"WARNING: Could not parse function signatures from CUDA source. "
                f"Using function name from compile_cuda call: {func_name}"
            )
            # Create a placeholder signature -- the user may need to adjust
            functions = [
                {
                    "return_type": "torch::Tensor",
                    "name": func_name,
                    "params": "torch::Tensor input",
                    "full_signature": f"torch::Tensor {func_name}(torch::Tensor input)",
                }
            ]
        else:
            print("ERROR: Could not find any torch::Tensor-returning functions in CUDA source.")
            print("       The CUDA source should contain launcher functions like:")
            print("         torch::Tensor my_kernel_cuda(torch::Tensor A, torch::Tensor B) { ... }")
            sys.exit(1)

    # Create directory structure (matches kernels-community convention)
    project_dir = os.path.join(output_dir, name)
    kernel_cuda_dir = os.path.join(project_dir, f"{name}_cuda")
    torch_ext_dir = os.path.join(project_dir, "torch-ext")

    os.makedirs(kernel_cuda_dir, exist_ok=True)
    os.makedirs(torch_ext_dir, exist_ok=True)

    # 1. Write kernel.cu
    kernel_cu_path = os.path.join(kernel_cuda_dir, "kernel.cu")
    cuda_src_clean = cuda_src.strip()
    with open(kernel_cu_path, "w", encoding="utf-8") as f:
        f.write(cuda_src_clean)
        f.write("\n")
    print(f"  Created {os.path.relpath(kernel_cu_path, output_dir)}")

    # 2. Write build.toml
    build_toml_path = os.path.join(project_dir, "build.toml")