fuse

This module provides functionality for array fusion and computation using lazy evaluation.

Overview:

Array fusion allows composing multiple array operations into a single kernel, enabling significant performance optimizations by letting the compiler optimize the entire operation at once.

Key Functions:

• defer: Marks an array as an input to a fused operation.
• compute: Executes the fused operation efficiently.
• fuse: Combines multiple operations into a single kernel.
• fused: A decorator for marking functions as fused.

Examples:

1. Basic Usage: >>> rng = np.random.default_rng(42) >>> A = rng.integers(0, 10, (10, 10)) >>> B = rng.integers(0, 10, (10, 10)) >>> C = defer(A) >>> D = defer(B) >>> E = (C + D) / 2 >>> compute(E) BufferizedNDArray(shape=(10, 10))

   In this example, E represents a fused operation that adds C and D together and divides the result by 2. The compute function optimizes and executes the operation efficiently.

2. Using fuse as a higher-order function: >>> result = fuse(lambda x, y: (x + y) / 2, A, B)

   Here, fuse combines the addition and division operations into a single fused kernel.

3. Using the fused decorator: >>> @fused ... def add_and_divide(x, y): ... return (x + y) / 2 >>> result = add_and_divide(A, B)

   The fused decorator enables automatic fusion of operations within the function.

Performance:

• Using defer and compute results in faster execution due to operation fusion.
• Different optimizers can be used with compute, such as the Galley optimizer, which adapts to the sparsity patterns of the inputs.
• The optimizer can be set using the ctx argument in compute, or via set_scheduler or with_scheduler.

compute(arg, ctx=None)

Executes a fused operation represented by LazyTensors. This function evaluates the entire operation in an optimized manner using the provided scheduler.

• arg: A lazy tensor or a tuple of lazy tensors representing the fused operation to be computed.
• ctx: The scheduler to use for computation. Defaults to the result of get_default_scheduler().

Returns: - A tensor or a list of tensors computed by the fused operation.

Source code in src/finch/interface/fuse.py

def compute(arg, ctx=None):
    """
    Executes a fused operation represented by LazyTensors. This function evaluates the
    entire operation in an optimized manner using the provided scheduler.

    Parameters:
    - arg: A lazy tensor or a tuple of lazy tensors representing the fused operation to
      be computed.
    - ctx: The scheduler to use for computation. Defaults to the result of
      `get_default_scheduler()`.

    Returns:
    - A tensor or a list of tensors computed by the fused operation.
    """
    if ctx is None:
        ctx = get_default_scheduler()

    args = arg if isinstance(arg, tuple) else (arg,)
    outputs = [None for _ in args]
    lazy_args = []
    lazy_arg_idxs = []
    for i, arg_i in enumerate(args):
        if not isinstance(arg_i, LazyTensor):
            outputs[i] = arg_i
        else:
            lazy_args.append(arg_i)
            lazy_arg_idxs.append(i)

    if lazy_args:
        vars = tuple(Alias(gensym("A")) for _ in lazy_args)
        ctx_2 = lazy_args[0].ctx.join(*[x.ctx for x in lazy_args[1:]])
        bodies = []
        for arg_i, var in zip(lazy_args, vars, strict=True):
            idxs = tuple(Field(gensym("i")) for _ in range(len(arg_i.shape)))
            bodies.append(Query(Table(var, idxs), Table(arg_i.data, idxs)))
        prgm = Plan((*ctx_2.trace(), *bodies, Produces(vars)))
        res = ctx(prgm)
        for lazy_idx, out_idx in enumerate(lazy_arg_idxs):
            device = lazy_args[lazy_idx].device
            if (
                len(res[lazy_idx].shape) == 0
            ):  # if the result is a scalar, extract the value and turn it into a
                # finch `Scalar`
                outputs[out_idx] = asarray(res[lazy_idx][()], device=device)
            else:
                outputs[out_idx] = asarray(res[lazy_idx], device=device)

    return tuple(outputs) if isinstance(arg, tuple) else outputs[0]

fuse(f, *args, ctx=None)

Fuses multiple array operations into a single kernel. This function allows for composing operations and executing them efficiently.

Parameters:

Name	Type	Description	Default
f	Callable	The function representing the operation to be fused, returning a tensor or tuple of tensor results.	required
*args	Any	The input arrays or LazyTensors to be fused.	()
ctx	Any	The scheduler to use for computation. Defaults to the result of get_default_scheduler().	None

Returns:

Type	Description
Any	The result of the fused operation, a tensor or tuple of tensors.

Source code in src/finch/interface/fuse.py

def fuse(f: Callable, *args: Any, ctx: Any = None) -> Any:
    """
    Fuses multiple array operations into a single kernel. This function allows for
    composing operations and executing them efficiently.

    Args:
        f: The function representing the operation to be fused, returning a tensor or
            tuple of tensor results.
        *args: The input arrays or LazyTensors to be fused.
        ctx: The scheduler to use for computation. Defaults to the result of
            `get_default_scheduler()`.

    Returns:
        The result of the fused operation, a tensor or tuple of tensors.
    """
    if ctx is None:
        ctx = get_default_scheduler()

    args = tuple(defer(arg) for arg in args)
    if len(args) == 1:
        return f(args[0])
    return compute(f(*args), ctx=ctx)

fused(f, /, ctx=None)

• fused(f): A decorator that marks a function as fused. This allows the function to be used with the fuse function for automatic fusion of operations.

Parameters: - f: The function to be marked as fused.

Returns: - A wrapper function that applies the fusion mechanism to the original function.

Source code in src/finch/interface/fuse.py

def fused(f, /, ctx=None):
    """
    - fused(f):
    A decorator that marks a function as fused. This allows the function to be used with
    the `fuse` function for automatic fusion of operations.

    Parameters:
    - f: The function to be marked as fused.

    Returns:
    - A wrapper function that applies the fusion mechanism to the original function.
    """
    if ctx is None:
        ctx = get_default_scheduler()

    def wrapper(*args):
        return fuse(f, *args, ctx=ctx)

    return wrapper