Function _generate_cpp_wrapper

kernels/cuda/_compile.py:103–146

Generate a pybind11-compatible C++ wrapper that forwards torch::Tensor arguments to the CUDA kernel launcher. arg_specs: list of (name, type_str) pairs. Supported types: - "tensor" -> torch::Tensor - "int" -> int64_t - "float" -> double - "bool" -> bool

(func_name: str, arg_specs: list)

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

def _generate_cpp_wrapper(func_name: str, arg_specs: list) -> str:
    """
    Generate a pybind11-compatible C++ wrapper that forwards torch::Tensor
    arguments to the CUDA kernel launcher.

    arg_specs: list of (name, type_str) pairs. Supported types:
        - "tensor" -> torch::Tensor
        - "int" -> int64_t
        - "float" -> double
        - "bool" -> bool
    """
    # Build function signature
    cpp_args = []
    for name, type_str in arg_specs:
        if type_str == "tensor":
            cpp_args.append(f"torch::Tensor {name}")
        elif type_str == "int":
            cpp_args.append(f"int64_t {name}")
        elif type_str == "float":
            cpp_args.append(f"double {name}")
        elif type_str == "bool":
            cpp_args.append(f"bool {name}")
        else:
            cpp_args.append(f"torch::Tensor {name}")

    args_str = ", ".join(cpp_args)
    forward_args = ", ".join(name for name, _ in arg_specs)

    wrapper = f"""
#include <torch/extension.h>

// Forward declaration of CUDA launcher (defined in .cu source)
torch::Tensor {func_name}_cuda({args_str});

// Python-facing wrapper
torch::Tensor {func_name}({args_str}) {{
    return {func_name}_cuda({forward_args});
}}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {{
    m.def("{func_name}", &{func_name}, "{func_name}");
}}
"""
    return wrapper


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