numba_codegen

NumbaArgumentFType

Bases: FType, ABC

Source code in src/finch/codegen/numba_codegen/numba.py

class NumbaArgumentFType(FType, ABC):
    @abstractmethod
    def numba_type(self) -> Any:
        """
        Return a Numba-compatible type for this ftype.
        """
        ...

    @abstractmethod
    def serialize_to_numba(self, obj):
        """
        Return a Numba-compatible object to be used in place of this argument
        for the Numba backend.
        """
        ...

    @abstractmethod
    def deserialize_from_numba(self, obj, numba_buffer):
        """
        Return an object from Numba returned value.
        """
        ...

    @abstractmethod
    def construct_from_numba(self, numba_buffer):
        """
        Construct and return an object from Numba returned value.
        """
        ...

construct_from_numba(numba_buffer) abstractmethod

Construct and return an object from Numba returned value.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def construct_from_numba(self, numba_buffer):
    """
    Construct and return an object from Numba returned value.
    """
    ...

deserialize_from_numba(obj, numba_buffer) abstractmethod

Return an object from Numba returned value.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def deserialize_from_numba(self, obj, numba_buffer):
    """
    Return an object from Numba returned value.
    """
    ...

numba_type() abstractmethod

Return a Numba-compatible type for this ftype.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def numba_type(self) -> Any:
    """
    Return a Numba-compatible type for this ftype.
    """
    ...

serialize_to_numba(obj) abstractmethod

Return a Numba-compatible object to be used in place of this argument for the Numba backend.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def serialize_to_numba(self, obj):
    """
    Return a Numba-compatible object to be used in place of this argument
    for the Numba backend.
    """
    ...

NumbaBufferFType

Bases: BufferFType, NumbaArgumentFType, ABC

Source code in src/finch/codegen/numba_codegen/numba.py

class NumbaBufferFType(BufferFType, NumbaArgumentFType, ABC):
    @abstractmethod
    def numba_length(self, ctx: "NumbaContext", buffer):
        """
        Return a Numba-compatible expression to get the length of the buffer.
        """
        ...

    @abstractmethod
    def numba_resize(self, ctx: "NumbaContext", buffer, size):
        """
        Return a Numba-compatible expression to resize the buffer to the given size.
        """
        ...

    @abstractmethod
    def numba_load(self, ctx: "NumbaContext", buffer, idx):
        """
        Return a Numba-compatible expression to load an element from the buffer
        at the given index.
        """
        ...

    @abstractmethod
    def numba_store(self, ctx: "NumbaContext", buffer, idx, value=None):
        """
        Return a Numba-compatible expression to store an element in the buffer
        at the given index. If value is None, it should store the length of the
        buffer.
        """
        ...

    @staticmethod
    def numba_name():
        return "list[numpy.ndarray]"

numba_length(ctx, buffer) abstractmethod

Return a Numba-compatible expression to get the length of the buffer.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def numba_length(self, ctx: "NumbaContext", buffer):
    """
    Return a Numba-compatible expression to get the length of the buffer.
    """
    ...

numba_load(ctx, buffer, idx) abstractmethod

Return a Numba-compatible expression to load an element from the buffer at the given index.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def numba_load(self, ctx: "NumbaContext", buffer, idx):
    """
    Return a Numba-compatible expression to load an element from the buffer
    at the given index.
    """
    ...

numba_resize(ctx, buffer, size) abstractmethod

Return a Numba-compatible expression to resize the buffer to the given size.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def numba_resize(self, ctx: "NumbaContext", buffer, size):
    """
    Return a Numba-compatible expression to resize the buffer to the given size.
    """
    ...

numba_store(ctx, buffer, idx, value=None) abstractmethod

Return a Numba-compatible expression to store an element in the buffer at the given index. If value is None, it should store the length of the buffer.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def numba_store(self, ctx: "NumbaContext", buffer, idx, value=None):
    """
    Return a Numba-compatible expression to store an element in the buffer
    at the given index. If value is None, it should store the length of the
    buffer.
    """
    ...

NumbaContext

Bases: Context

Source code in src/finch/codegen/numba_codegen/numba.py

class NumbaContext(Context):
    def __init__(self, tab="    ", indent=0, types=None, slots=None, *, mode=None):
        if types is None:
            types = ScopedDict()
        if slots is None:
            slots = ScopedDict()

        super().__init__(mode=mode)

        self.tab = tab
        self.indent = indent
        self.types = types
        self.slots = slots

        self.imports = [
            "import _operator, builtins",
            "from numba import njit",
            "import numpy",
            "from numpy import int64, float64",
            "\n",
        ]

    @property
    def feed(self) -> str:
        return self.tab * self.indent

    def emit_global(self):
        """
        Emit the headers for the C code.
        """
        return "\n".join([*self.imports, self.emit()])

    def emit(self):
        return "\n".join([*self.preamble, *self.epilogue])

    def block(self) -> "NumbaContext":
        blk = super().block()
        blk.indent = self.indent
        blk.tab = self.tab
        blk.types = self.types
        blk.slots = self.slots
        return blk

    def subblock(self):
        blk = self.block()
        blk.indent = self.indent + 1
        blk.types = self.types.scope()
        blk.slots = self.slots.scope()
        return blk

    def cache(self, name, val):
        if isinstance(val, asm.Literal | asm.Variable):
            return val
        var_n = self.freshen(name)
        var_t = val.result_type
        self.exec(f"{self.feed}{var_n} = {self(val)}")
        return asm.Variable(var_n, var_t)

    def resolve(self, node):
        match node:
            case asm.Slot(var_n, _):
                if var_n in self.slots:
                    return self.slots[var_n]
                raise KeyError(f"Slot {var_n} not found in context")
            case _:
                raise ValueError(f"Expected Slot, got: {type(node)}")

    @staticmethod
    def full_name(val: Any) -> str:
        if isinstance(val, algebra.ftypes.FType):
            val = numba_type(val)
        if hasattr(val, "numba_name"):
            return val.numba_name()
        if hasattr(val, "name"):
            return val.name
        if not hasattr(val, "__module__") or not hasattr(val, "__name__"):
            return str(val)
        return f"{val.__module__}.{val.__name__}"

    def __call__(self, prgm: asm.AssemblyNode):
        feed = self.feed
        match prgm:
            case asm.Assert(exp):
                condition = self(exp)
                message = f"Finch assertion failed: {exp}"
                self.exec(
                    f"{feed}if not ({condition}):\n"
                    f"{feed}{self.tab}raise AssertionError({message!r})"
                )
                return None
            case asm.Literal(value):
                if is_dynamic(value):
                    # str() would silently emit broken source.
                    raise DynamicFillError(
                        "cannot emit a dynamic fill as a numba literal"
                    )
                return str(value)
            case asm.Variable(name, _) | asm.Slot(name, _):
                return name.replace("#", "_")
            case asm.Assign(asm.Variable(var_n, var_t) as var, val):
                val_code = self(val)
                var_code = self(var)
                if val.result_type != var_t:
                    raise TypeError(f"Type mismatch: {val.result_type} != {var_t}")
                if var_n in self.types:
                    assert var_t == self.types[var_n]
                    self.exec(f"{feed}{var_code} = {val_code}")
                else:
                    self.types[var_n] = var_t
                    self.exec(
                        f"{feed}{var_code}: "
                        f"{self.full_name(numba_type(var_t))} = {val_code}"
                    )
                return None
            case asm.GetAttr(obj, attr):
                obj_code = self(obj)
                obj_t = obj.result_type
                if not isinstance(obj_t, StructFType):
                    raise TypeError(f"Expected struct type, got: {obj_t}")
                if not obj_t.struct_hasattr(attr.val):
                    raise ValueError(f"trying to get missing attr: {attr}")
                return numba_getattr(obj_t, self, obj_code, attr.val)
            case asm.SetAttr(obj, attr, val):
                obj_code = self(obj)
                obj_t = obj.result_type
                if not isinstance(obj_t, StructFType):
                    raise TypeError(f"Expected struct type, got: {obj_t}")
                if not fisinstance(val, obj_t.struct_attrtype(attr.val)):
                    raise TypeError(
                        f"Type mismatch: {val.result_type} != "
                        f"{obj_t.struct_attrtype(attr.val)}"
                    )
                val_code = self(val)
                numba_setattr(obj_t, self, obj_code, attr.val, val_code)
                return None
            case asm.Call(op, args):
                return numba_function_call(op, self, *args)

            case asm.Unpack(asm.Slot(var_n, var_t) as slot, 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."
                    )
                var_code = self(slot)
                self.exec(f"{feed}{var_code} = {self(val)}")
                self.types[var_n] = var_t
                self.slots[var_n] = var_t.numba_unpack(
                    self, var_code, asm.Variable(var_n, var_t)
                )
                return None
            case asm.Repack(asm.Slot(var_n, var_t) as slot):
                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]}")
                obj = self.slots[var_n]
                var_code = self(slot)
                var_t.numba_repack(self, var_code, obj)
                return None
            case asm.Load(buf, idx):
                buf_t = buf.result_type
                if not isinstance(buf_t, NumbaBufferFType):
                    raise TypeError(f"Expected numba buffer type, got: {buf_t}")
                if self.mode.debug:
                    idx = self.cache("index", idx)
                    self(
                        asm.Assert(
                            asm.Call(
                                asm.Literal(ffuncs.and_),
                                (
                                    asm.Call(
                                        asm.Literal(ffuncs.ge),
                                        (idx, asm.Literal(idx.result_type(0))),
                                    ),
                                    asm.Call(
                                        asm.Literal(ffuncs.lt),
                                        (idx, asm.Length(buf)),
                                    ),
                                ),
                            )
                        )
                    )
                return buf_t.numba_load(self, self.resolve(buf), idx)
            case asm.Store(buf, idx, val):
                buf_t = buf.result_type
                if not isinstance(buf_t, NumbaBufferFType):
                    raise TypeError(f"Expected numba buffer type, got: {buf_t}")
                if self.mode.debug:
                    idx = self.cache("index", idx)
                    self(
                        asm.Assert(
                            asm.Call(
                                asm.Literal(ffuncs.and_),
                                (
                                    asm.Call(
                                        asm.Literal(ffuncs.ge),
                                        (idx, asm.Literal(idx.result_type(0))),
                                    ),
                                    asm.Call(
                                        asm.Literal(ffuncs.lt),
                                        (idx, asm.Length(buf)),
                                    ),
                                ),
                            )
                        )
                    )
                buf_t.numba_store(self, self.resolve(buf), idx, val)
                return None
            case asm.Resize(buf, size):
                buf_t = buf.result_type
                if not isinstance(buf_t, NumbaBufferFType):
                    raise TypeError(f"Expected numba buffer type, got: {buf_t}")
                buf_t.numba_resize(self, self.resolve(buf), size)
                return None
            case asm.Length(buf):
                buf_t = buf.result_type
                if not isinstance(buf_t, NumbaBufferFType):
                    raise TypeError(f"Expected numba buffer type, got: {buf_t}")
                return buf_t.numba_length(self, self.resolve(buf))
            case asm.Block(bodies):
                ctx_2 = self.block()
                if bodies == ():
                    self.exec(f"{feed}pass")
                else:
                    for body in bodies:
                        ctx_2(body)
                    self.exec(ctx_2.emit())
                return None
            case asm.ForLoop(asm.Variable(_, _) as var, start, end, body):
                var_2 = self(var)
                start = self(start)
                end = self(end)
                ctx_2 = self.subblock()
                ctx_2(body)
                ctx_2.types[var.name] = var.result_type
                body_code = ctx_2.emit()
                self.exec(f"{feed}for {var_2} in range({start}, {end}):\n{body_code}")
                return None
            case asm.BufferLoop(buf, var, body):
                raise NotImplementedError
            case asm.WhileLoop(cond, body):
                ctx_2 = self.subblock()
                cond_code = ctx_2(cond)
                if ctx_2.preamble or ctx_2.epilogue:
                    # Checks and cached indices in the condition must run on
                    # every iteration, including the final condition evaluation.
                    ctx_2.exec(f"{ctx_2.feed}if not ({cond_code}):")
                    ctx_2.exec(f"{ctx_2.feed}{self.tab}break")
                    cond_code = "True"
                ctx_2(body)
                body_code = ctx_2.emit()
                self.exec(f"{feed}while {cond_code}:\n{body_code}")
                return None
            case asm.If(cond, body):
                cond_code = self(cond)
                ctx_2 = self.subblock()
                ctx_2(body)
                body_code = ctx_2.emit()
                self.exec(f"{feed}if {cond_code}:\n{body_code}")
                return None
            case asm.IfElse(cond, body, else_body):
                cond_code = self(cond)
                ctx_2 = self.subblock()
                ctx_2(body)
                body_code = ctx_2.emit()
                ctx_3 = self.subblock()
                ctx_3(else_body)
                else_body_code = ctx_3.emit()
                self.exec(
                    f"{feed}if {cond_code}:\n{body_code}\n{feed}else:\n{else_body_code}"
                )
                return None
            case asm.Function(
                asm.Variable(func_name, asm.AssemblyKernelFType(result_type=return_t)),
                args,
                body,
            ):
                ctx_2 = self.subblock()
                arg_decls = []
                for arg in args:
                    match arg:
                        case asm.Variable(name, t) as var:
                            name_code = self(var)
                            arg_decls.append(
                                f"{name_code}: {self.full_name(numba_type(t))}"
                            )
                            ctx_2.types[name] = t
                        case _:
                            raise NotImplementedError(
                                f"Unrecognized argument type: {arg}"
                            )
                ctx_2(body)
                body_code = ctx_2.emit()
                feed = self.feed
                self.exec(
                    f"{feed}@njit\n"
                    f"{feed}def {func_name}({', '.join(arg_decls)}) -> "
                    f"{self.full_name(numba_type(return_t))}:\n"
                    f"{body_code}\n"
                )
                return None
            case asm.Return(value):
                value = self(value)
                self.exec(f"{feed}return {value}")
                return None
            case asm.Break():
                self.exec(f"{feed}break")
                return None
            case asm.Module(funcs):
                for func in funcs:
                    if not isinstance(func, asm.Function):
                        raise NotImplementedError(
                            f"Unrecognized function type: {type(func)}"
                        )
                    self(func)
                return None
            case node:
                raise NotImplementedError(f"Unrecognized node: {node}")

emit_global()

Emit the headers for the C code.

Source code in src/finch/codegen/numba_codegen/numba.py

def emit_global(self):
    """
    Emit the headers for the C code.
    """
    return "\n".join([*self.imports, self.emit()])

NumbaLibrary

Bases: AssemblyLibrary

A class to represent a Numba module.

Source code in src/finch/codegen/numba_codegen/numba.py

class NumbaLibrary(asm.AssemblyLibrary):
    """
    A class to represent a Numba module.
    """

    def __init__(self, kernels):
        self.kernels = kernels

    def __getattr__(self, name):
        # Allow attribute access to kernels by name
        if name in self.kernels:
            return self.kernels[name]
        raise AttributeError(
            f"{type(self).__name__!r} object has no attribute {name!r}"
        )

NumbaLowerer

Bases: Stage

Source code in src/finch/codegen/numba_codegen/stages.py

class NumbaLowerer(Stage):
    @abstractmethod
    def lower(self, prgm: asm.Module, *, mode: CompilerMode | None = None) -> NumbaCode:
        """
        Lower the given assembly program to Numba code.
        """

lower(prgm, *, mode=None) abstractmethod

Lower the given assembly program to Numba code.

Source code in src/finch/codegen/numba_codegen/stages.py

@abstractmethod
def lower(self, prgm: asm.Module, *, mode: CompilerMode | None = None) -> NumbaCode:
    """
    Lower the given assembly program to Numba code.
    """

NumbaUnpackableFType

Bases: ABC

Abstract base class for unpackable formats in Numba. Unpackable formats must also support other functions with unpacked inputs in addition to variable ones.

Source code in src/finch/codegen/numba_codegen/numba.py

class NumbaUnpackableFType(ABC):
    """
    Abstract base class for unpackable formats in Numba. Unpackable formats must also
    support other functions with unpacked inputs in addition to variable ones.
    """

    @abstractmethod
    def numba_unpack(self, ctx, lhs, rhs):
        """
        Convert a value to an unpacked representation in Numba. Returns a NamedTuple
        of unpacked variable names, etc. The `lhs` is the variable namespace to
        assign to.
        """
        ...

    @abstractmethod
    def numba_repack(self, ctx, lhs, rhs):
        """
        Update an object based on an unpacked representation. The `rhs` is the
        unpacked representation to update from, and `lhs` is a variable name referring
        to the original object to update.
        """
        ...

numba_repack(ctx, lhs, rhs) abstractmethod

Update an object based on an unpacked representation. The rhs is the unpacked representation to update from, and lhs is a variable name referring to the original object to update.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def numba_repack(self, ctx, lhs, rhs):
    """
    Update an object based on an unpacked representation. The `rhs` is the
    unpacked representation to update from, and `lhs` is a variable name referring
    to the original object to update.
    """
    ...

numba_unpack(ctx, lhs, rhs) abstractmethod

Convert a value to an unpacked representation in Numba. Returns a NamedTuple of unpacked variable names, etc. The lhs is the variable namespace to assign to.

Source code in src/finch/codegen/numba_codegen/numba.py

@abstractmethod
def numba_unpack(self, ctx, lhs, rhs):
    """
    Convert a value to an unpacked representation in Numba. Returns a NamedTuple
    of unpacked variable names, etc. The `lhs` is the variable namespace to
    assign to.
    """
    ...

construct_from_numba(fmt, numba_obj)

Construct an object from a Numba-compatible ftype.

Parameters:

Name	Type	Description	Default
fmt	FType	The ftype of the object.	required
numba_obj	Any	The Numba-compatible object to construct from.	required

Returns:

Type	Description
Any	An instance of the original object type.

Source code in src/finch/codegen/numba_codegen/numba.py

def construct_from_numba(fmt: FType, numba_obj: Any) -> Any:
    """
    Construct an object from a Numba-compatible ftype.

    Args:
        fmt: The ftype of the object.
        numba_obj: The Numba-compatible object to construct from.

    Returns:
        An instance of the original object type.
    """
    match fmt:
        case asm.AssemblyKernelFType():
            return NumbaKernel(numba_obj, fmt)
        case NumbaArgumentFType():
            return fmt.construct_from_numba(numba_obj)
        case algebra.none_:
            return None
        case algebra.ftypes.FDTypeNumpy():
            return fmt(numba_obj)
        case algebra.int_ | algebra.bool_ | algebra.float_:
            return numba_obj
        case TupleFType() | ImmutableStructFType():
            return immutable_construct_from_numba(fmt, numba_obj)
        case StructFType():
            return struct_construct_from_numba(fmt, numba_obj)
        case _:
            return fmt(numba_obj)

deserialize_from_numba(fmt, obj, numba_obj)

Deserialize a Numba-compatible object back to the original ftype.

Parameters:

Name	Type	Description	Default
fmt	FType	FType of obj	required
obj	Any	The original object to update.	required
numba_obj	Any	The Numba-compatible object to deserialize from.	required

Returns:

Type	Description
None	None

Source code in src/finch/codegen/numba_codegen/numba.py

def deserialize_from_numba(fmt: FType, obj: Any, numba_obj: Any) -> None:
    """
    Deserialize a Numba-compatible object back to the original ftype.

    Args:
        fmt: FType of obj
        obj: The original object to update.
        numba_obj: The Numba-compatible object to deserialize from.

    Returns:
        None
    """
    match fmt:
        case NumbaArgumentFType():
            fmt.deserialize_from_numba(obj, numba_obj)
        case algebra.none_ | algebra.ftypes.FDTypeNumpy():
            return
        case StructFType():
            _deserialize_asm_struct_from_numba(fmt, obj, numba_obj)
        case _:
            return

numba_jitclass_type(t)

Returns the Numba jitclass spec type/ftype after serialization.

Parameters:

Name	Type	Description	Default
t	FType	The Python type/ftype.	required

Returns:

Type	Description
Any	The corresponding Numba jitclass spec type.

Source code in src/finch/codegen/numba_codegen/numba.py

def numba_jitclass_type(t: FType) -> Any:
    """
    Returns the Numba jitclass spec type/ftype after serialization.

    Args:
        t: The Python type/ftype.

    Returns:
        The corresponding Numba jitclass spec type.
    """
    match t:
        case asm.AssemblyKernelFType():
            ret = (
                numba.void
                if t.result_type == algebra.none_
                else numba_jitclass_type(t.result_type)
            )
            return numba.types.FunctionType(
                ret(*(numba_jitclass_type(arg) for arg in t.arg_types))
            )
        case _ if hasattr(t, "numba_jitclass_type"):
            return t.numba_jitclass_type()  # ty: ignore[call-non-callable]
        case algebra.ftypes.FDTypeNumpy():
            return numba.from_dtype(t.dtype)
        case algebra.int_:
            return numba.int32
        case algebra.bool_:
            return numba.boolean
        case algebra.float_:
            return numba.float64
        case TupleFType() | ImmutableStructFType():
            return immutable_struct_jitclass_type(t)
        case StructFType():
            return assembly_struct_numba_jitclass_type(t)
        case _:
            raise NotImplementedError(f"No Numba jitclass type mapping for {t}")

numba_type(t)

Returns the Numba type/ftype after serialization.

Parameters:

Name	Type	Description	Default
t	FType	The Python type/ftype.	required

Returns:

Type	Description
Any	The corresponding Numba type.

Source code in src/finch/codegen/numba_codegen/numba.py

def numba_type(t: FType) -> Any:
    """
    Returns the Numba type/ftype after serialization.

    Args:
        t: The Python type/ftype.

    Returns:
        The corresponding Numba type.
    """
    match t:
        case asm.AssemblyKernelFType():
            return object
        case NumbaArgumentFType():
            return t.numba_type()
        case algebra.ftypes.FDTypeNumpy():
            return numba.from_dtype(t.dtype)
        case algebra.ftypes.FDTypeBuiltin():
            return t.type
        case TupleFType() | ImmutableStructFType():
            return tuple
        case StructFType():
            return assembly_struct_numba_type(t)
        case _:
            return t

serialize_to_numba(fmt, obj)

Serialize an object to a Numba-compatible ftype.

Parameters:

Name	Type	Description	Default
fmt	FType	FType of obj	required
obj	Any	The object to serialize.	required

Returns:

Type	Description
Any	A Numba-compatible object.

Source code in src/finch/codegen/numba_codegen/numba.py

def serialize_to_numba(fmt: FType, obj: Any) -> Any:
    """
    Serialize an object to a Numba-compatible ftype.

    Args:
        fmt: FType of obj
        obj: The object to serialize.

    Returns:
        A Numba-compatible object.
    """
    match fmt:
        case asm.AssemblyKernelFType():
            return obj.numba_func
        case NumbaArgumentFType():
            return fmt.serialize_to_numba(obj)
        case algebra.none_:
            return None
        case (
            algebra.ftypes.FDTypeNumpy() | algebra.int_ | algebra.bool_ | algebra.float_
        ):
            return obj
        case TupleFType() | ImmutableStructFType():
            return serialize_immutable_to_numba(fmt, obj)
        case StructFType():
            return _serialize_asm_struct_to_numba(fmt, obj)
        case _:
            raise NotImplementedError(f"No Numba serialization mapping for {fmt}")