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, and every read of one, 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`, and every read of one, 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_type", None))
    )

    def rule(node):
        match node:
            case ntn.Literal(DynamicFill() | DynamicFillFType() as fill):
                return ntn.Literal(_zero_value(fill.element_type))
            case ntn.Literal(StaticFill() | StaticFillFType() as fill):
                return ntn.Literal(fill.value)
            case ntn.FillValue(tns):
                match tns.result_type.fill_type:
                    case DynamicFillFType() as fill:
                        return ntn.Literal(_zero_value(fill.element_type))
                    case StaticFillFType() as fill:
                        return ntn.Literal(fill.value)
        return None

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