compile

AssemblyContext

Bases: Context

Compiles Finch Notation to Finch Assembly. Holds the state of the compilation process.

Source code in src/finch/compile/lower.py

class AssemblyContext(Context):
    """
    Compiles Finch Notation to Finch Assembly. Holds the state of the
    compilation process.
    """

    def __init__(
        self,
        namespace=None,
        preamble=None,
        epilogue=None,
        bindings=None,
        slots=None,
        access_modes=None,
        types=None,
        func_state: HaltState | None = None,
        mode: CompilerMode | None = None,
    ):
        super().__init__(
            namespace=namespace, preamble=preamble, epilogue=epilogue, mode=mode
        )
        if bindings is None:
            bindings = ScopedDict()
        if slots is None:
            slots = ScopedDict()
        if access_modes is None:
            access_modes = ScopedDict()
        if types is None:
            types = ScopedDict()
        self.bindings = bindings
        self.slots = slots
        self.access_modes = access_modes
        self.types = types
        self.func_state = func_state

    def _slot_expr(self, slot):
        match slot:
            case ntn.Fiber(lvl, _):
                return self(lvl.root)
            case _:
                raise TypeError(f"Unrecognized slot state: {slot}")

    @staticmethod
    def _slot_cursor(slot, type_):
        fields = dict(type_.struct_fields)
        if "pos" in fields:
            pos_t = fields["pos"]
        elif "val" in fields and hasattr(fields["val"], "length_type"):
            pos_t = fields["val"].length_type
        elif getattr(type_, "shape_type", None) == ():
            # 0-d tensors store their value directly rather than in a
            # positioned buffer; the cursor position is unused.
            pos_t = np.intp
        else:
            raise TypeError(f"Cannot create an assembly cursor for {type_}")
        return ntn.Fiber(ntn.Root(ntn.Value(slot, type_)), ntn.Literal(pos_t(0)))

    def block(self):
        """
        Create a new block. Preambles and epilogues will stay within this block.
        This is used to create a new context for compiling a block of code.
        """
        blk = super().block()
        blk.bindings = self.bindings
        blk.slots = self.slots
        blk.access_modes = self.access_modes
        blk.types = self.types
        blk.func_state = self.func_state
        return blk

    def scope(self):
        """
        Create a new scoped context that inherits from this one.
        """
        blk = self.block()
        blk.bindings = self.bindings.scope()
        blk.slots = self.slots.scope()
        blk.access_modes = self.access_modes.scope()
        blk.types = self.types.scope()
        return blk

    def should_halt(self):
        """
        Check if the current function should halt.
        This is used to determine if the function has returned.
        """
        assert self.func_state is not None
        return self.func_state.has_returned

    def emit(self):
        return self.preamble + self.epilogue

    def resolve(self, node):
        match node:
            case ntn.Slot(var_n, var_t):
                if var_n in self.slots:
                    view = self.slots[var_n]
                    if view.result_type != var_t:
                        raise TypeError(
                            f"Slot '{var_n}' is declared as type {view.result_type}, "
                            f"but used as type {var_t}."
                        )
                    if isinstance(view, ntn.Fiber):
                        return view
                    raise TypeError(f"Unrecognized slot state: {view}")
                raise KeyError(f"Slot {var_n} not found in context")
            case ntn.Fiber():
                return node
            case ntn.Value(_, _) | ntn.Full():
                return node
            case _:
                raise ValueError(
                    f"Expected Slot, Fiber, Value, or Full, got: {type(node)}"
                )

    def _freeze_tensor(self, tns_var: str, op: ntn.Literal | None) -> None:
        if op is None:
            assert tns_var not in self.access_modes
        else:
            assert self.access_modes[tns_var] == ntn.Update(op)
        self.access_modes[tns_var] = ntn.Read()

    def _thaw_tensor(self, tns_var: str, op: ntn.Literal) -> None:
        assert self.access_modes[tns_var] == ntn.Read()
        self.access_modes[tns_var] = ntn.Update(op)

    def _rm_tensor_from_accesses(self, tns_var: str) -> None:
        del self.access_modes[tns_var]

    # TODO: Move to .pyi file

    @overload
    def __call__(self, prgm: ntn.Module) -> asm.Module:
        pass

    @overload
    def __call__(self, prgm: ntn.Literal) -> asm.Literal:
        pass

    @overload
    def __call__(self, prgm: ntn.Variable) -> asm.Variable:
        pass

    @overload
    def __call__(self, prgm: ntn.NotationExpression) -> asm.AssemblyExpression:
        pass

    @overload
    def __call__(self, prgm: ntn.NotationStatement) -> None:
        pass

    def __call__(self, prgm):
        """
        Lower Finch Notation to Finch Assembly. First we check for early
        simplifications, then we call the normal lowering for the outermost
        node.
        """
        match prgm:
            case ntn.Literal(value):
                return asm.Literal(value)
            case ntn.Value(expr, _):
                return expr
            case ntn.Root(tns):
                return tns.result_type.get_child(self(tns), "lvl")
            case ntn.Child(parent, attr):
                return parent.result_type.level_get_child(self(parent), attr)
            case ntn.Call(f, args):
                f_e = self(f)
                args_e = tuple(self(arg) for arg in args)
                return asm.Call(f_e, args_e)
            case ntn.Assign(var, val):
                self.exec(asm.Assign(self(var), self(val)))
                return None
            case ntn.Variable(var_n, var_t):
                return asm.Variable(var_n, var_t)
            case ntn.Slot(var_n, var_t):
                if var_n in self.types:
                    def_t = self.types[var_n]
                    if def_t != var_t:
                        raise TypeError(
                            f"Slot '{var_n}' is declared as type {def_t}, "
                            f"but used as type {var_t}."
                        )
                if var_n in self.slots:
                    return self._slot_expr(self.slots[var_n])
                raise KeyError(f"Slot '{var_n}' is not defined in the current context.")
            case ntn.Fiber(lvl, _):
                return self(lvl.root)
            case ntn.Unpack(ntn.Slot(var_n, var_t), val):
                val_code = self(val)
                if val.result_type != var_t:
                    raise TypeError(f"Type mismatch: {val.result_type} != {var_t}")
                if var_n in self.slots:
                    raise KeyError(
                        f"Slot {var_n} already exists in context, cannot unpack"
                    )
                if var_n in self.types:
                    raise KeyError(
                        f"Variable '{var_n}' is already defined in the current"
                        f" context, cannot overwrite with slot."
                    )
                slot = asm.Slot(var_n, var_t)
                self.exec(asm.Unpack(slot, val_code))
                self.types[var_n] = var_t
                self._freeze_tensor(var_n, op=None)
                self.slots[var_n] = self._slot_cursor(slot, var_t)
                return None
            case ntn.Repack(ntn.Slot(var_n, var_t), _):
                if var_n not in self.slots or var_n not in self.types:
                    raise KeyError(f"Slot {var_n} not found in context, cannot repack")
                if var_t != self.types[var_n]:
                    raise TypeError(f"Type mismatch: {var_t} != {self.types[var_n]}")
                self._rm_tensor_from_accesses(var_n)
                self.exec(asm.Repack(asm.Slot(var_n, var_t)))
                return None
            case ntn.Unwrap(ntn.Full(val, ())):
                return self(val)
            case ntn.Unwrap(ntn.Access(ntn.Full(val, ()), ntn.Read(), ())):
                return self(val)
            case ntn.Dimension(ntn.Full(_, shape), ntn.Literal(r)):
                return self(shape[r])
            case ntn.Unwrap(ntn.Access(tns, mode, idxs)):
                assert isinstance(mode, ntn.Read)
                assert idxs == ()
                tns = self.resolve(tns)
                return tns.result_type.lower_unwrap(self, tns)
            case ntn.Increment(ntn.Access(tns, mode, idxs), val):
                assert isinstance(mode, ntn.Update)
                assert idxs == ()
                tns = self.resolve(tns)
                op = mode.op
                return tns.result_type.lower_increment(self, tns, op, val)
            case ntn.Block(bodies):
                for body in bodies:
                    self(body)
                return None
            case ntn.Loop(idx, ext, body):
                ext.result_type.lower_loop(
                    self, idx, SymbolicExtent.from_notation(ext), body
                )
                return None
            case ntn.Dimension(tns, ntn.Literal(r)):
                assert isinstance(r, int)
                tns = self.resolve(tns)
                return tns.result_type.lower_dim(self, tns, r)
            case ntn.FillValue(tns):
                tns = self.resolve(tns)
                return tns.result_type.lower_fill(self, tns)
            case ntn.Cached(arg, _):
                return self(arg)
            case ntn.Declare(tns, init, op, shape):
                self._thaw_tensor(tns.name, op)
                tns = self.resolve(tns)
                # A declaration may change a dynamic fill, but not a static one.
                fill = tns.result_type.fill_type
                if is_static(fill) and not (
                    isinstance(init, ntn.Literal)
                    and bool(ffuncs.same(init.val, fill.value))
                ):
                    raise ValueError(
                        f"Declaring {tns} with {init} would change its static "
                        f"fill {fill.value!r}"
                    )
                init_e = self(init)
                op_e = self(op)
                shape_e = [self(s) for s in shape]
                return tns.result_type.lower_declare(self, tns, init_e, op_e, shape_e)
            case ntn.Freeze(tns, op):
                self._freeze_tensor(tns.name, op)
                tns = self.resolve(tns)
                op_e = self(op)
                return tns.result_type.lower_freeze(self, tns, op_e)
            case ntn.Thaw(tns, op):
                self._thaw_tensor(tns.name, op)
                tns = self.resolve(tns)
                op_e = self(op)
                return tns.result_type.lower_thaw(self, tns, op_e)
            case ntn.If(cond, body):
                cond_e = self(cond)
                ctx_2 = self.scope()
                ctx_2(body)
                self.exec(asm.If(cond_e, asm.Block(ctx_2.emit())))
                return None
            case ntn.IfElse(cond, body, else_body):
                cond_e = self(cond)
                ctx_2 = self.scope()
                ctx_2(body)
                ctx_3 = self.scope()
                ctx_3(else_body)
                self.exec(
                    asm.IfElse(cond_e, asm.Block(ctx_2.emit()), asm.Block(ctx_3.emit()))
                )
                return None
            case ntn.Function(
                ntn.Variable(
                    func_n, asm.AssemblyKernelFType(result_type=ret_t) as func_type
                ),
                args,
                body,
            ):
                ctx = self.scope()
                ctx.func_state = HaltState(
                    return_var=asm.Variable(ctx.freshen(f"{func_n}_return"), ret_t)
                )
                blk = ctx.scope()
                blk(body)
                self.exec(
                    asm.Function(
                        asm.Variable(func_n, func_type),
                        tuple(ctx(arg) for arg in args),
                        asm.Block((*blk.emit(), asm.Return(ctx.func_state.return_var))),
                    )
                )
                return None
            case ntn.Return(value):
                if self.func_state is None:
                    raise ValueError("Return statement outside of function.")
                self.exec(asm.Assign(self.func_state.return_var, self(value)))
                return None
            case ntn.Module(funcs):
                ctx = self.scope()
                for func in funcs:
                    ctx(func)
                return asm.Module(ctx.emit())
            case other:
                raise Exception(f"{other} not recognized.")

__call__(prgm)

__call__(prgm: ntn.Module) -> asm.Module

__call__(prgm: ntn.Literal) -> asm.Literal

__call__(prgm: ntn.Variable) -> asm.Variable

__call__(
    prgm: ntn.NotationExpression,
) -> asm.AssemblyExpression

__call__(prgm: ntn.NotationStatement) -> None

Lower Finch Notation to Finch Assembly. First we check for early simplifications, then we call the normal lowering for the outermost node.

Source code in src/finch/compile/lower.py

def __call__(self, prgm):
    """
    Lower Finch Notation to Finch Assembly. First we check for early
    simplifications, then we call the normal lowering for the outermost
    node.
    """
    match prgm:
        case ntn.Literal(value):
            return asm.Literal(value)
        case ntn.Value(expr, _):
            return expr
        case ntn.Root(tns):
            return tns.result_type.get_child(self(tns), "lvl")
        case ntn.Child(parent, attr):
            return parent.result_type.level_get_child(self(parent), attr)
        case ntn.Call(f, args):
            f_e = self(f)
            args_e = tuple(self(arg) for arg in args)
            return asm.Call(f_e, args_e)
        case ntn.Assign(var, val):
            self.exec(asm.Assign(self(var), self(val)))
            return None
        case ntn.Variable(var_n, var_t):
            return asm.Variable(var_n, var_t)
        case ntn.Slot(var_n, var_t):
            if var_n in self.types:
                def_t = self.types[var_n]
                if def_t != var_t:
                    raise TypeError(
                        f"Slot '{var_n}' is declared as type {def_t}, "
                        f"but used as type {var_t}."
                    )
            if var_n in self.slots:
                return self._slot_expr(self.slots[var_n])
            raise KeyError(f"Slot '{var_n}' is not defined in the current context.")
        case ntn.Fiber(lvl, _):
            return self(lvl.root)
        case ntn.Unpack(ntn.Slot(var_n, var_t), val):
            val_code = self(val)
            if val.result_type != var_t:
                raise TypeError(f"Type mismatch: {val.result_type} != {var_t}")
            if var_n in self.slots:
                raise KeyError(
                    f"Slot {var_n} already exists in context, cannot unpack"
                )
            if var_n in self.types:
                raise KeyError(
                    f"Variable '{var_n}' is already defined in the current"
                    f" context, cannot overwrite with slot."
                )
            slot = asm.Slot(var_n, var_t)
            self.exec(asm.Unpack(slot, val_code))
            self.types[var_n] = var_t
            self._freeze_tensor(var_n, op=None)
            self.slots[var_n] = self._slot_cursor(slot, var_t)
            return None
        case ntn.Repack(ntn.Slot(var_n, var_t), _):
            if var_n not in self.slots or var_n not in self.types:
                raise KeyError(f"Slot {var_n} not found in context, cannot repack")
            if var_t != self.types[var_n]:
                raise TypeError(f"Type mismatch: {var_t} != {self.types[var_n]}")
            self._rm_tensor_from_accesses(var_n)
            self.exec(asm.Repack(asm.Slot(var_n, var_t)))
            return None
        case ntn.Unwrap(ntn.Full(val, ())):
            return self(val)
        case ntn.Unwrap(ntn.Access(ntn.Full(val, ()), ntn.Read(), ())):
            return self(val)
        case ntn.Dimension(ntn.Full(_, shape), ntn.Literal(r)):
            return self(shape[r])
        case ntn.Unwrap(ntn.Access(tns, mode, idxs)):
            assert isinstance(mode, ntn.Read)
            assert idxs == ()
            tns = self.resolve(tns)
            return tns.result_type.lower_unwrap(self, tns)
        case ntn.Increment(ntn.Access(tns, mode, idxs), val):
            assert isinstance(mode, ntn.Update)
            assert idxs == ()
            tns = self.resolve(tns)
            op = mode.op
            return tns.result_type.lower_increment(self, tns, op, val)
        case ntn.Block(bodies):
            for body in bodies:
                self(body)
            return None
        case ntn.Loop(idx, ext, body):
            ext.result_type.lower_loop(
                self, idx, SymbolicExtent.from_notation(ext), body
            )
            return None
        case ntn.Dimension(tns, ntn.Literal(r)):
            assert isinstance(r, int)
            tns = self.resolve(tns)
            return tns.result_type.lower_dim(self, tns, r)
        case ntn.FillValue(tns):
            tns = self.resolve(tns)
            return tns.result_type.lower_fill(self, tns)
        case ntn.Cached(arg, _):
            return self(arg)
        case ntn.Declare(tns, init, op, shape):
            self._thaw_tensor(tns.name, op)
            tns = self.resolve(tns)
            # A declaration may change a dynamic fill, but not a static one.
            fill = tns.result_type.fill_type
            if is_static(fill) and not (
                isinstance(init, ntn.Literal)
                and bool(ffuncs.same(init.val, fill.value))
            ):
                raise ValueError(
                    f"Declaring {tns} with {init} would change its static "
                    f"fill {fill.value!r}"
                )
            init_e = self(init)
            op_e = self(op)
            shape_e = [self(s) for s in shape]
            return tns.result_type.lower_declare(self, tns, init_e, op_e, shape_e)
        case ntn.Freeze(tns, op):
            self._freeze_tensor(tns.name, op)
            tns = self.resolve(tns)
            op_e = self(op)
            return tns.result_type.lower_freeze(self, tns, op_e)
        case ntn.Thaw(tns, op):
            self._thaw_tensor(tns.name, op)
            tns = self.resolve(tns)
            op_e = self(op)
            return tns.result_type.lower_thaw(self, tns, op_e)
        case ntn.If(cond, body):
            cond_e = self(cond)
            ctx_2 = self.scope()
            ctx_2(body)
            self.exec(asm.If(cond_e, asm.Block(ctx_2.emit())))
            return None
        case ntn.IfElse(cond, body, else_body):
            cond_e = self(cond)
            ctx_2 = self.scope()
            ctx_2(body)
            ctx_3 = self.scope()
            ctx_3(else_body)
            self.exec(
                asm.IfElse(cond_e, asm.Block(ctx_2.emit()), asm.Block(ctx_3.emit()))
            )
            return None
        case ntn.Function(
            ntn.Variable(
                func_n, asm.AssemblyKernelFType(result_type=ret_t) as func_type
            ),
            args,
            body,
        ):
            ctx = self.scope()
            ctx.func_state = HaltState(
                return_var=asm.Variable(ctx.freshen(f"{func_n}_return"), ret_t)
            )
            blk = ctx.scope()
            blk(body)
            self.exec(
                asm.Function(
                    asm.Variable(func_n, func_type),
                    tuple(ctx(arg) for arg in args),
                    asm.Block((*blk.emit(), asm.Return(ctx.func_state.return_var))),
                )
            )
            return None
        case ntn.Return(value):
            if self.func_state is None:
                raise ValueError("Return statement outside of function.")
            self.exec(asm.Assign(self.func_state.return_var, self(value)))
            return None
        case ntn.Module(funcs):
            ctx = self.scope()
            for func in funcs:
                ctx(func)
            return asm.Module(ctx.emit())
        case other:
            raise Exception(f"{other} not recognized.")

block()

Create a new block. Preambles and epilogues will stay within this block. This is used to create a new context for compiling a block of code.

Source code in src/finch/compile/lower.py

def block(self):
    """
    Create a new block. Preambles and epilogues will stay within this block.
    This is used to create a new context for compiling a block of code.
    """
    blk = super().block()
    blk.bindings = self.bindings
    blk.slots = self.slots
    blk.access_modes = self.access_modes
    blk.types = self.types
    blk.func_state = self.func_state
    return blk

scope()

Create a new scoped context that inherits from this one.

Source code in src/finch/compile/lower.py

def scope(self):
    """
    Create a new scoped context that inherits from this one.
    """
    blk = self.block()
    blk.bindings = self.bindings.scope()
    blk.slots = self.slots.scope()
    blk.access_modes = self.access_modes.scope()
    blk.types = self.types.scope()
    return blk

should_halt()

Check if the current function should halt. This is used to determine if the function has returned.

Source code in src/finch/compile/lower.py

def should_halt(self):
    """
    Check if the current function should halt.
    This is used to determine if the function has returned.
    """
    assert self.func_state is not None
    return self.func_state.has_returned

CompilerMode dataclass

Safe mode checks tensor extents; debug also checks buffer accesses.

Source code in src/finch/symbolic/stage.py

@dataclass(frozen=True)
class CompilerMode:
    """Safe mode checks tensor extents; debug also checks buffer accesses."""

    safe: bool = False
    debug: bool = False

Extent dataclass

Bases: FTyped[ExtentFType]

A class to represent the extent of a loop variable. This is used to define the start and end values of a loop.

The extent is start-inclusive and end-exclusive, meaning that the loop variable will take on values from start to end-1.

Source code in src/finch/compile/lower.py

@dataclass(frozen=True)
class Extent(FTyped[ExtentFType]):
    """
    A class to represent the extent of a loop variable. This is used to define
    the start and end values of a loop.

    The extent is start-inclusive and end-exclusive, meaning that the loop variable
    will take on values from start to end-1.
    """

    start: Any
    end: Any

    def loop_interpreter(self, ctx, idx, body):
        for idx_e in range(self.start, self.end):
            # Create a new scope for each iteration
            ctx_2 = ctx.scope(loop_state=HaltState())
            # Assign the loop variable
            ctx_2.bindings[idx.name] = idx.type_(idx_e)
            # Execute the body of the loop
            ctx_2(body)

    @property
    def ftype(self):
        return ExtentFType(ftype(self.start), ftype(self.end))

ExtentFType dataclass

Bases: ImmutableStructFType

Source code in src/finch/compile/lower.py

@dataclass(eq=True, frozen=True)
class ExtentFType(ImmutableStructFType):
    start_t: FType
    end_t: FType

    def __repr__(self):
        return f"ExtentFType(start={self.start_t}, end={self.end_t})"

    @property
    def struct_name(self):
        return "Extent"

    @property
    def struct_fields(self):
        return [("start", self.start_t), ("end", self.end_t)]

    def from_fields(self, start, end) -> Extent:
        return Extent(start, end)

    def __call__(self, *args):
        raise TypeError(f"{self.struct_name} is not callable")

    def lower_loop(
        self,
        ctx: AssemblyContext,
        idx: ntn.Variable,
        ext: SymbolicExtent,
        body: ntn.NotationExpression,
    ):
        """
        Lower a loop with the given index and body.
        This is used to compile the loop into assembly.
        """
        lower_looplets(ctx, idx, ext, body)
        return

    def default_loop(self, ctx, idx, ext: SymbolicExtent, body):
        def assert_lowered(node):
            match node:
                case ntn.Access(_, _, (j, *_)):
                    if j == idx:
                        raise FinchCompileError(
                            node, f"Access with {j} should have been lowered already"
                        )
            return

        for node in PostOrderDFS(body):
            assert_lowered(node)

        ctx_2 = ctx.scope()
        ctx_2(body)
        body_3 = asm.Block(ctx_2.emit())
        ctx.exec(
            asm.ForLoop(
                ctx(idx),
                ctx(ext.get_start()),
                ctx(ext.get_end()),
                body_3,
            )
        )
        return

lower_loop(ctx, idx, ext, body)

Lower a loop with the given index and body. This is used to compile the loop into assembly.

Source code in src/finch/compile/lower.py

def lower_loop(
    self,
    ctx: AssemblyContext,
    idx: ntn.Variable,
    ext: SymbolicExtent,
    body: ntn.NotationExpression,
):
    """
    Lower a loop with the given index and body.
    This is used to compile the loop into assembly.
    """
    lower_looplets(ctx, idx, ext, body)
    return