compiler

FinchJLKernel

Bases: AssemblyKernel

A callable Julia kernel.

Source code in src/finch/compile_jl/compiler.py

class FinchJLKernel(AssemblyKernel):
    """A callable Julia kernel."""

    def __init__(
        self,
        jl_code,
        type_,
        finch_program: ntn.Function,
        runtime: FinchJLRuntime,
        func_name: str,
        dynamic_args: tuple[int, ...] = (),
        extents: tuple[tuple[int, int], ...] = (),
    ):
        super().__init__(type_)
        # We store this code so that we can verify it in pytest
        self.jl_code = jl_code
        self.finch_program = finch_program
        self.runtime = runtime
        self.func_name = func_name
        self.dynamic_args = dynamic_args
        # The position and axis of the argument each trailing extent measures.
        self.extents = extents

    def __call__(self, *args):
        return self.runtime.kernel_call(self.func_name, self, args)

handle_fills(func)

Rewrite every Dynamic fill in func to a zero of its dtype, and report which argument positions carried a Dynamic fill. This is a necessary but potentially unsound rewrite which should be removed eventually.

Source code in src/finch/compile_jl/compiler.py

def handle_fills(func: ntn.Function) -> tuple[ntn.Function, tuple[int, ...]]:
    """Rewrite every Dynamic fill in `func` to a zero of its dtype, and report
    which argument positions carried a Dynamic fill. This is a necessary but
    potentially unsound rewrite which should be removed eventually.
    """
    dynamic_args = tuple(
        i
        for i, arg in enumerate(func.args)
        if is_dynamic(getattr(arg.type_, "fill_value", None))
    )

    def rule(node):
        match node:
            case ntn.Literal(DynamicFill() as fill):
                return ntn.Literal(ftype(fill.value)(0))
            case ntn.Literal(StaticFill() as fill):
                return ntn.Literal(fill.value)
        return None

    return Rewrite(PostWalk(rule))(func), dynamic_args