Function codegen_func

warp/_src/codegen.py:5018–5167

(adj, c_func_name: str, device="cpu", options=None, forward_only=False, reverse_only=False)

def codegen_func(adj, c_func_name: str, device="cpu", options=None, forward_only=False, reverse_only=False):
    if options is None:
        options = {}

    # 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"

    # forward header
    if adj.return_var is not None and len(adj.return_var) == 1:
        return_type = adj.return_var[0].ctype()
    else:
        return_type = "void"

    has_multiple_outputs = adj.return_var is not None and len(adj.return_var) != 1

    forward_args = []
    reverse_args = []

    # Tile parameters are emitted as C++ template parameters so that the
    # same @wp.func can accept tiles with any storage type (register or
    # shared) without requiring separate overloads.  The Python-level tile
    # annotation (e.g. wp.tile[float, M, N]) defaults to
    # register storage, but at the call site the tile may actually live in
    # shared memory.  By generating ``template<typename tile_t>`` instead
    # of the concrete ``tile_register_t<...>`` type, C++ template argument
    # deduction resolves the correct storage type automatically.
    #
    # Tile parameters are passed by non-const reference (not by value)
    # for two reasons: (1) owning shared tiles (tile_shared_t with
    # Owner=true) cannot be copied (static_assert in copy constructor),
    # and (2) adjoint built-ins like adj_tile_sum() expect non-const
    # Tile& parameters.
    #
    # This is a semantic change for register tiles, which were previously
    # passed by value.  The difference is observable for in-place tile
    # operations (e.g., a += b where both are tiles), which mutate the
    # parameter directly.  Simple rebinding (a = expr) creates a new C++
    # variable: for register tiles this is a full value copy, for shared
    # tiles a non-owning handle to the same shared memory (element-level
    # writes through either variable affect the same data).
    # The pass-by-reference behavior for in-place ops is intentional:
    # it matches the Python semantics where augmented assignment on a
    # mutable object modifies it in place.
    template_params = []

    # forward args
    for i, arg in enumerate(adj.args):
        if is_tile(arg.type) or is_tile_stack(arg.type):
            tname = f"tile_{arg.label}"
            template_params.append(tname)
            s = f"{tname}& {arg.emit()}"
        else:
            s = f"{arg.ctype()} {arg.emit()}"
        forward_args.append(s)