Function codegen_kernel

warp/_src/codegen.py:5264–5360

(kernel, device, options)

def codegen_kernel(kernel, device, options):
    # Update the module's options with the ones defined on the kernel, if any.
    options = options | kernel.options

    adj = kernel.adj

    args_struct = ""
    if device == "cpu":
        args_struct = f"struct wp_args_{kernel.get_mangled_name()} {{\n"
        for i in adj.args:
            args_struct += f"    {i.ctype()} {i.label};\n"
        args_struct += "};\n"

    # Build line directive for function definition (subtract 1 to account for 1-indexing of AST line numbers)
    # This is used as a catch-all C-to-Python source line mapping for any code that does not have
    # a direct mapping to a Python source line.
    func_line_directive = ""
    if line_directive := adj.get_line_directive("", adj.fun_def_lineno - 1):
        func_line_directive = f"{line_directive}\n"

    if device == "cpu":
        template_forward = cpu_kernel_template_forward
        template_backward = cpu_kernel_template_backward
    elif device == "cuda":
        template_forward = cuda_kernel_template_forward
        template_backward = cuda_kernel_template_backward
    else:
        raise ValueError(f"Device {device} is not supported")

    template = ""
    template_fmt_args = {
        "name": kernel.get_mangled_name(),
        "launch_ndim": kernel.adj.kernel_dim,
    }

    # Generate launch_bounds string for CUDA kernels
    launch_bounds_str = ""
    if device == "cuda" and "launch_bounds" in options:
        launch_bounds = options["launch_bounds"]
        if isinstance(launch_bounds, int):
            launch_bounds_str = f"__launch_bounds__({launch_bounds}) "
        elif isinstance(launch_bounds, (tuple, list)):
            if len(launch_bounds) == 1:
                launch_bounds_str = f"__launch_bounds__({launch_bounds[0]}) "
            elif len(launch_bounds) == 2:
                launch_bounds_str = f"__launch_bounds__({launch_bounds[0]}, {launch_bounds[1]}) "
            else:
                raise ValueError(f"launch_bounds must be an int or a tuple/list of 1-2 ints, got {launch_bounds}")
        else:
            raise ValueError(f"launch_bounds must be an int or a tuple/list of 1-2 ints, got {type(launch_bounds)}")

    # build forward signature
    forward_args = [f"wp::launch_bounds_t<{adj.kernel_dim}> dim"]
    if device == "cpu":
        forward_args.append("size_t task_index")
    else:
        for arg in adj.args:
            forward_args.append(arg.ctype() + " var_" + arg.label)