Large constants under compiled functions

[renecotyfanboy]

Hi, I have a question related to large arrays in jitted functions. I have issues with functions that interpolate within large static arrays, leading to extremely long compilation times because of constant folding or equivalent I guess. I know that a good practice to deal with this kind of issues would be to put these large table as parameters of the function I want, and it works like a charm when doing so. But since I am working with external frameworks, there is a point where I cannot pass these tables as input anymore, and since the framework is enforcing compilation, I have no choice but to find a way to prevent this constant folding.

large_array = np.random.randn(201, 30, 500_000)

def jax_func():
    return large_array

jax.jit(jax_func)() # Takes forever, how to prevent this ?

I think that solving this code would be sufficient for my problem. I could use pure_callbacks to return the tables but I am afraid to lose too much optimization doing so. I have also read about #21300, but I can't get the XLA flag to work, and I don't know how to manipulate jaxpr to do this. Any opinion ?

[renecotyfanboy]

At this time, I am using the solution proposed by @inversecrime in #21300 which consists in manipulating the jaxpr to set the constants as variables and prevent the folding. I am sharing below the script he sent which defines a custom @jit decorator doing this.

from functools import wraps
from typing import Any, Callable

import jax
import jax.core
import jax.numpy as jnp
import jax.interpreters.partial_eval
from jax import Array, ShapeDtypeStruct
from jax.core import ClosedJaxpr, Jaxpr


def abstract_like(x: Any) -> Any:
    return jax.tree.map(lambda leaf: ShapeDtypeStruct(shape=jnp.asarray(leaf).shape, dtype=jnp.asarray(leaf).dtype), x)


def clean_jaxpr(jaxpr: Jaxpr) -> Jaxpr:
    return Jaxpr(
        constvars=jaxpr.constvars,
        invars=jaxpr.invars,
        outvars=jaxpr.outvars,
        eqns=jaxpr.eqns
    )


def make_jaxpr[**P](func: Callable[P, Any]) -> Callable[P, ClosedJaxpr]:
    def wrapper(*args, **kwargs):
        closed_jaxpr = jax.make_jaxpr(func)(*args, **kwargs)
        closed_jaxpr = closed_jaxpr.replace(
            jaxpr=clean_jaxpr(closed_jaxpr.jaxpr)
        )
        return closed_jaxpr

    return wrapper  # type: ignore


def make_jaxpr_and_return_treedefs[**P](func: Callable[P, Array]) -> Callable[P, tuple[ClosedJaxpr, Any, Any]]:
    def wrapper(*args, **kwargs):
        (_, in_treedef) = jax.tree.flatten((args, kwargs))
        (closed_jaxpr, abstract_output) = jax.make_jaxpr(func, return_shape=True)(*args, **kwargs)
        closed_jaxpr = closed_jaxpr.replace(
            jaxpr=clean_jaxpr(closed_jaxpr.jaxpr)
        )
        (_, out_treedef) = jax.tree.flatten(abstract_output)
        return (closed_jaxpr, in_treedef, out_treedef)

    return wrapper  # type: ignore


def eval_jaxpr(jaxpr: Jaxpr, consts: list[Array], ins: list[Array]) -> list[Array]:
    return jax.core.eval_jaxpr(jaxpr, consts, *ins)


def move_invars_to_constvars(jaxpr: Jaxpr, *, n: int) -> Jaxpr:
    return jaxpr.replace(
        constvars=jaxpr.constvars + jaxpr.invars[:n],
        invars=jaxpr.invars[n:]
    )


def assume_known(closed_jaxpr: ClosedJaxpr, *, knowns: list[Array]) -> ClosedJaxpr:
    return closed_jaxpr.replace(
        jaxpr=move_invars_to_constvars(closed_jaxpr.jaxpr, n=len(knowns)),
        consts=closed_jaxpr.consts + knowns
    )


def precompute_jaxpr(closed_jaxpr: ClosedJaxpr) -> ClosedJaxpr:
    (closed_known_jaxpr, closed_unknown_jaxpr, _, _) = jax.interpreters.partial_eval.partial_eval_jaxpr_nounits(closed_jaxpr, unknowns=[True] * len(closed_jaxpr.jaxpr.invars), instantiate=True)

    with jax.ensure_compile_time_eval():
        residuals = jax.jit(lambda consts, ins: eval_jaxpr(closed_known_jaxpr.jaxpr, consts, ins))(closed_known_jaxpr.consts, [])

    closed_unknown_jaxpr = assume_known(closed_unknown_jaxpr, knowns=residuals)

    return closed_unknown_jaxpr


def map_over_children(mapping: Callable[..., Any], tree: Any, *rest: Any) -> Any:
    return jax.tree.map(mapping, tree, *rest, is_leaf=lambda node: node is not tree)


def separate(tree: Any, *, separator: Any, replacement: Any | None = None) -> tuple[Any, Any]:
    (are_left, treedef) = jax.tree.flatten(separator)
    leaves = treedef.flatten_up_to(tree)
    leaves_left = [leaf if is_left else replacement for (leaf, is_left) in zip(leaves, are_left)]
    leaves_right = [replacement if is_left else leaf for (leaf, is_left) in zip(leaves, are_left)]
    tree_left = jax.tree.unflatten(treedef, leaves_left)
    tree_right = jax.tree.unflatten(treedef, leaves_right)
    return (tree_left, tree_right)


def combine(*trees: Any, replacement: Any | None = None) -> Any:
    if any(x is replacement for x in trees):
        (index, tree) = next(((index, tree) for (index, tree) in enumerate(trees) if tree is not replacement), (0, replacement))
        assert trees == (replacement,) * (index) + (tree,) + (replacement,) * (len(trees) - index - 1)
        return tree
    return map_over_children(lambda *children: combine(*children), *trees)


def jit[F:Callable](f: F, *, static_argnums: list[int] | None = None, static_argnames: list[str] | None = None, precompute: bool = True) -> F:
    if static_argnums is None:
        static_argnums = []
    if static_argnames is None:
        static_argnames = []

    cache_dict = {}

    @wraps(f)
    def wrapper(*args, **kwargs):
        assert set(static_argnums).issubset(range(len(args)))
        assert set(static_argnames).issubset(kwargs.keys())

        (static_args, dynamic_args) = separate(args, separator=tuple(argnum in static_argnums for argnum in range(len(args))))
        (static_kwargs, dynamic_kwargs) = separate(kwargs, separator={argname: argname in static_argnames for argname in kwargs.keys()})

        (dynamic_in_leaves, dynamic_in_treedef) = jax.tree.flatten((dynamic_args, dynamic_kwargs))
        abstract_dynamic_in_leaves = abstract_like(dynamic_in_leaves)

        cache_key = (static_args, frozenset(static_kwargs.items()), dynamic_in_treedef, tuple(abstract_dynamic_in_leaves))

        if cache_key in cache_dict.keys():
            (compiled_function, consts) = cache_dict[cache_key]
            result = compiled_function(consts, dynamic_in_leaves)
            jax.block_until_ready(result)
            return result

        def dynamic_function(dynamic_in_leaves):
            (dynamic_args, dynamic_kwargs) = jax.tree.unflatten(dynamic_in_treedef, dynamic_in_leaves)
            args = combine(static_args, dynamic_args)
            kwargs = combine(static_kwargs, dynamic_kwargs)
            return f(*args, **kwargs)

        (closed_jaxpr, _, out_treedef) = make_jaxpr_and_return_treedefs(dynamic_function)(dynamic_in_leaves)

        if precompute:
            closed_jaxpr = precompute_jaxpr(closed_jaxpr)

        compiled_function = jax.jit(lambda consts, dynamic_in_leaves: jax.tree.unflatten(out_treedef, eval_jaxpr(closed_jaxpr.jaxpr, consts, dynamic_in_leaves)))

        cache_dict[cache_key] = (compiled_function, closed_jaxpr.consts)

        return compiled_function(closed_jaxpr.consts, dynamic_in_leaves)

    return wrapper  # type: ignore