Function generate_torch_binding_cpp

export_hf.py:367–398

Generate torch_binding.cpp with TORCH_LIBRARY_EXPAND registration. Uses the HuggingFace Kernels convention: torch/library.h + registration.h + REGISTER_EXTENSION macro for compatibility with the kernel-builder Nix build pipeline.

(
    name: str,
    functions: List[Dict[str, str]],
)

def generate_torch_binding_cpp(
    name: str,
    functions: List[Dict[str, str]],
) -> str:
    """
    Generate torch_binding.cpp with TORCH_LIBRARY_EXPAND registration.

    Uses the HuggingFace Kernels convention: torch/library.h + registration.h
    + REGISTER_EXTENSION macro for compatibility with the kernel-builder
    Nix build pipeline.
    """
    # Build ops.def() and ops.impl() lines
    ops_lines = []
    for func in functions:
        schema = _build_ops_schema(func)
        ops_lines.append(f'  ops.def("{schema}");')
        ops_lines.append(f'  ops.impl("{func["name"]}", torch::kCUDA, &{func["name"]});')

    ops_str = "\n".join(ops_lines)

    return (
        "#include <torch/library.h>\n"
        "\n"
        '#include "registration.h"\n'
        '#include "torch_binding.h"\n'
        "\n"
        "TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {\n"
        f"{ops_str}\n"
        "}\n"
        "\n"
        "REGISTER_EXTENSION(TORCH_EXTENSION_NAME)\n"
    )


def generate_torch_binding_h(