codegen

CArgumentFType

Bases: ABC

Source code in src/finch/codegen/c_codegen/c.py

class CArgumentFType(ABC):
    @abstractmethod
    def c_type(self):
        """
        Return a ctypes type for this ftype.
        """
        ...

    @abstractmethod
    def serialize_to_c(self, obj):
        """
        Return a ctypes-compatible struct to be used in place of `obj`
        for the c backend.
        """
        ...

    @abstractmethod
    def deserialize_from_c(self, obj, res):
        """
        Update this `obj` based on how the c call modified `res`, the result
        of `serialize_to_c`.
        """
        ...

    @abstractmethod
    def construct_from_c(self, res):
        """
        Construct a new object based on the return value from c
        """

c_type() abstractmethod

Return a ctypes type for this ftype.

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def c_type(self):
    """
    Return a ctypes type for this ftype.
    """
    ...

construct_from_c(res) abstractmethod

Construct a new object based on the return value from c

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def construct_from_c(self, res):
    """
    Construct a new object based on the return value from c
    """

deserialize_from_c(obj, res) abstractmethod

Update this obj based on how the c call modified res, the result of serialize_to_c.

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def deserialize_from_c(self, obj, res):
    """
    Update this `obj` based on how the c call modified `res`, the result
    of `serialize_to_c`.
    """
    ...

serialize_to_c(obj) abstractmethod

Return a ctypes-compatible struct to be used in place of obj for the c backend.

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def serialize_to_c(self, obj):
    """
    Return a ctypes-compatible struct to be used in place of `obj`
    for the c backend.
    """
    ...

CBufferFType

Bases: BufferFType, CArgumentFType, ABC

Abstract base class for the ftype of buffers. The ftype defines how the data in a Buffer is organized and accessed.

Source code in src/finch/codegen/c_codegen/c.py

class CBufferFType(BufferFType, CArgumentFType, ABC):
    """
    Abstract base class for the ftype of buffers. The ftype defines how the
    data in a Buffer is organized and accessed.
    """

    @abstractmethod
    def c_length(self, ctx, buffer):
        """
        Return C code which loads a named buffer at the given index.
        """
        ...

    @abstractmethod
    def c_load(self, ctx, buffer, index):
        """
        Return C code which loads a named buffer at the given index.
        """
        ...

    @abstractmethod
    def c_store(self, ctx, buffer, index, value):
        """
        Return C code which stores a named buffer to the given index.
        """
        ...

    @abstractmethod
    def c_resize(self, ctx, buffer, new_length):
        """
        Return C code which resizes a named buffer to the given length.
        """
        ...

c_length(ctx, buffer) abstractmethod

Return C code which loads a named buffer at the given index.

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def c_length(self, ctx, buffer):
    """
    Return C code which loads a named buffer at the given index.
    """
    ...

c_load(ctx, buffer, index) abstractmethod

Return C code which loads a named buffer at the given index.

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def c_load(self, ctx, buffer, index):
    """
    Return C code which loads a named buffer at the given index.
    """
    ...

c_resize(ctx, buffer, new_length) abstractmethod

Return C code which resizes a named buffer to the given length.

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def c_resize(self, ctx, buffer, new_length):
    """
    Return C code which resizes a named buffer to the given length.
    """
    ...

c_store(ctx, buffer, index, value) abstractmethod

Return C code which stores a named buffer to the given index.

Source code in src/finch/codegen/c_codegen/c.py

@abstractmethod
def c_store(self, ctx, buffer, index, value):
    """
    Return C code which stores a named buffer to the given index.
    """
    ...

CCompiler

Bases: UnvalidatedForm, AssemblyLoader

A class to compile and run FinchAssembly.

Source code in src/finch/codegen/c_codegen/c.py

class CCompiler(UnvalidatedForm, asm.AssemblyLoader):
    """
    A class to compile and run FinchAssembly.
    """

    def __init__(
        self, ctx: CLowerer | None = None, cc=None, cflags=None, shared_cflags=None
    ):
        if cc is None:
            cc = config.get("cc")
        if cflags is None:
            cflags = config.get("cflags").split()
        if shared_cflags is None:
            shared_cflags = config.get("shared_cflags").split()
        self.cc = cc
        self.cflags = cflags
        self.shared_cflags = shared_cflags
        self.ctx: CLowerer = CGenerator() if ctx is None else ctx

    def lower(self, prgm: asm.Module, *, mode: CompilerMode | None = None) -> CLibrary:
        c_code = self.ctx(prgm, mode=mode).code
        logger.debug(f"Compiling C code:\n{c_code}")
        lib = load_shared_lib(
            c_code=c_code,
            cc=self.cc,
            cflags=(*self.cflags, *self.shared_cflags),
        )
        kernels = {}
        if prgm.head() != asm.Module:
            raise ValueError(
                "CCompiler expects a Module as the head of the program, "
                f"got {type(prgm.head())}"
            )
        for func in prgm.funcs:
            match func:
                case asm.Function(asm.Variable(func_name, func_type), _, _):
                    kern = CKernel(getattr(lib, func_name), func_type)
                    kernels[func_name] = kern
                case _:
                    raise NotImplementedError(
                        f"Unrecognized function type: {type(func)}"
                    )
        return CLibrary(lib, kernels)

CKernel

Bases: AssemblyKernel

A class to represent a C kernel.

Source code in src/finch/codegen/c_codegen/c.py

class CKernel(asm.AssemblyKernel):
    """
    A class to represent a C kernel.
    """

    def __init__(self, c_function, type_):
        super().__init__(type_)
        self.c_function = c_function
        ret_type, argtypes = type_.result_type, type_.arg_types
        self.ret_type = ret_type
        self.argtypes = argtypes
        self.c_args = tuple(
            (i, argtype)
            for i, argtype in enumerate(argtypes)
            if c_type(argtype) is not None
        )
        self.c_function.restype = c_type(ret_type)
        self.c_function.argtypes = tuple(c_type(t) for _, t in self.c_args)

    def __call__(self, *args):
        """
        Calls the C function with the given arguments.
        """
        if len(args) != len(self.argtypes):
            raise ValueError(
                f"Expected {len(self.argtypes)} arguments, got {len(args)}"
            )
        for argtype, arg in zip(self.argtypes, args, strict=False):
            if not fisinstance(arg, argtype):
                raise TypeError(f"Expected argument of type {argtype}, got {type(arg)}")
        serial_args = [serialize_to_c(t, args[i]) for i, t in self.c_args]
        res = self.c_function(*serial_args)
        for (i, type_), serial_arg in zip(self.c_args, serial_args, strict=True):
            deserialize_from_c(type_, args[i], serial_arg)
        if self.ret_type is algebra.none_:
            return None
        return construct_from_c(self.ret_type, res)

__call__(*args)

Calls the C function with the given arguments.

Source code in src/finch/codegen/c_codegen/c.py

def __call__(self, *args):
    """
    Calls the C function with the given arguments.
    """
    if len(args) != len(self.argtypes):
        raise ValueError(
            f"Expected {len(self.argtypes)} arguments, got {len(args)}"
        )
    for argtype, arg in zip(self.argtypes, args, strict=False):
        if not fisinstance(arg, argtype):
            raise TypeError(f"Expected argument of type {argtype}, got {type(arg)}")
    serial_args = [serialize_to_c(t, args[i]) for i, t in self.c_args]
    res = self.c_function(*serial_args)
    for (i, type_), serial_arg in zip(self.c_args, serial_args, strict=True):
        deserialize_from_c(type_, args[i], serial_arg)
    if self.ret_type is algebra.none_:
        return None
    return construct_from_c(self.ret_type, res)

CLibrary

Bases: AssemblyLibrary

A class to represent a C module.

Source code in src/finch/codegen/c_codegen/c.py

class CLibrary(asm.AssemblyLibrary):
    """
    A class to represent a C module.
    """

    def __init__(self, c_module, kernels):
        self.c_module = c_module
        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}"
        )

CLowerer

Bases: Stage

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

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

lower(prgm, *, mode=None) abstractmethod

Lower the given assembly program to C code.

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

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

MLIRBufferFType

Bases: BufferFType, MLIRArgumentFType, ABC

Abstract base class for the ftype of datastructures. The ftype defines how the data in an Buffer is organized and accessed.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRBufferFType(BufferFType, MLIRArgumentFType, ABC):
    """
    Abstract base class for the ftype of datastructures. The ftype defines how
    the data in an Buffer is organized and accessed.
    """

    @abstractmethod
    def mlir_length(self, ctx: MLIRContext, buffer):
        """
        Return MLIR code which loads a named buffer at the given index.
        """
        ...

    @abstractmethod
    def mlir_load(self, ctx: MLIRContext, buffer, index):
        """
        Return MLIR code which loads a named buffer at the given index.
        """
        ...

    @abstractmethod
    def mlir_store(self, ctx: MLIRContext, buffer, index, value=None):
        """
        Return MLIR code which stores a named buffer to the given index.
        """
        ...

    @abstractmethod
    def mlir_resize(self, ctx: MLIRContext, buffer, new_length):
        """Resize a buffer and store its updated descriptor in the slot box."""
        ...

mlir_length(ctx, buffer) abstractmethod

Return MLIR code which loads a named buffer at the given index.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_length(self, ctx: MLIRContext, buffer):
    """
    Return MLIR code which loads a named buffer at the given index.
    """
    ...

mlir_load(ctx, buffer, index) abstractmethod

Return MLIR code which loads a named buffer at the given index.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_load(self, ctx: MLIRContext, buffer, index):
    """
    Return MLIR code which loads a named buffer at the given index.
    """
    ...

mlir_resize(ctx, buffer, new_length) abstractmethod

Resize a buffer and store its updated descriptor in the slot box.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_resize(self, ctx: MLIRContext, buffer, new_length):
    """Resize a buffer and store its updated descriptor in the slot box."""
    ...

mlir_store(ctx, buffer, index, value=None) abstractmethod

Return MLIR code which stores a named buffer to the given index.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_store(self, ctx: MLIRContext, buffer, index, value=None):
    """
    Return MLIR code which stores a named buffer to the given index.
    """
    ...

MLIRCompiler

Bases: MLIRForm, AssemblyLoader

A class to compile and run FinchAssembly.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRCompiler(MLIRForm, asm.AssemblyLoader):
    """
    A class to compile and run FinchAssembly.

    """

    def __init__(self, ctx: MLIRLowerer | None = None):
        self.ctx: MLIRLowerer = MLIRGenerator() if ctx is None else ctx

    def lower(
        self, prgm: asm.Module, *, mode: CompilerMode | None = None
    ) -> MLIRLibrary:
        if prgm.head() != asm.Module:
            raise ValueError(
                "MLIRCompiler expects a Module as the head of the program, "
                f"got {type(prgm.head())}"
            )
        mlir_code = self.ctx(prgm, mode=mode).code
        logger.debug(f"Compiling MLIR code:\n{mlir_code}")
        context, module, engine = load_mlir_engine(mlir_code)
        kernels = {}
        for func in prgm.funcs:
            match func:
                case asm.Function(asm.Variable(func_name, func_type), _, _):
                    kernels[func_name] = MLIRKernel(
                        engine,
                        func_name,
                        func_type,
                    )
                case _:
                    raise NotImplementedError(
                        f"Unrecognized function type: {type(func)}"
                    )
        return MLIRLibrary((context, module, engine), kernels)

MLIRKernel

Bases: AssemblyKernel

A class to represent MLIR kernel. It validates and serializes arguments, invokes the compiled function, and reconstructs the result.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRKernel(asm.AssemblyKernel):
    """
    A class to represent MLIR kernel. It validates and serializes arguments,
    invokes the compiled function, and reconstructs the result.

    """

    def __init__(self, engine, func_name, type_):
        super().__init__(type_)
        self.engine = engine
        self.func_name = func_name
        self.ret_type = type_.result_type
        self.argtypes = type_.arg_types

    def __call__(self, *args):
        if len(args) != len(self.argtypes):
            raise ValueError(
                f"Expected {len(self.argtypes)} arguments, got {len(args)}"
            )
        for argtype, arg in zip(self.argtypes, args, strict=False):
            if not fisinstance(arg, argtype):
                raise TypeError(f"Expected argument of type {argtype}, got {type(arg)}")
        serial_args = list(map(serialize_to_mlir, self.argtypes, args))

        packed = []
        for t, sa in zip(self.argtypes, serial_args, strict=False):
            if isinstance(t, BufferFType) and mlir_type(t) != "!llvm.ptr":
                packed.append(ctypes.pointer(ctypes.pointer(sa)))
            else:
                packed.append(ctypes.pointer(sa))

        if self.ret_type == algebra.none_:
            res = None
            self.engine.invoke(self.func_name, *packed)
        elif isinstance(self.ret_type, StructFType):
            res = mlir_ctype(self.ret_type)()
            self.engine.invoke(
                self.func_name, *packed, ctypes.pointer(ctypes.pointer(res))
            )
        else:
            res = mlir_ctype(self.ret_type)()
            self.engine.invoke(self.func_name, *packed, ctypes.pointer(res))

        for arg_type, arg, serial_arg in zip(
            self.argtypes, args, serial_args, strict=False
        ):
            deserialize_from_mlir(arg_type, arg, serial_arg)

        if self.ret_type is type(None):
            return None
        return construct_from_mlir(self.ret_type, res)

MLIRLibrary

Bases: AssemblyLibrary

A class to represent MLIR module.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRLibrary(asm.AssemblyLibrary):
    """
    A class to represent MLIR module.

    """

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

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

MLIRLowerer

Bases: Stage

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

class MLIRLowerer(Stage):
    @abstractmethod
    def lower(self, prgm: asm.Module, *, mode: CompilerMode | None = None) -> MLIRCode:
        """
        Lower the given assembly program to textual MLIR.

        """

lower(prgm, *, mode=None) abstractmethod

Lower the given assembly program to textual MLIR.

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

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

    """

MallocBuffer

Bases: Buffer

A buffer that uses buffers managed by malloc to store data.

Source code in src/finch/codegen/buffers/malloc_buffer.py

class MallocBuffer(Buffer):
    """
    A buffer that uses buffers managed by malloc to store data.
    """

    def __init__(self, length: int, dtype: FType, data=None):
        """
        Constructor for the MallocBuffer class.

        length (int): the length of the malloc array.
        dtype (FType): the type that the buffer will be based on.
        data (optional): a list of data to initialize the buffer with.
        """
        self._dtype = dtype
        self._c_dtype = c_type(self._dtype)
        self.buffer = ctypes.pointer(CMallocBufferStruct())

        MallocBufferBackend.library(self.ftype).init(
            self.buffer,
            serialize_to_c(self.length_type, ctypes.sizeof(self._c_dtype)),
            serialize_to_c(self.length_type, length),
        )
        if data is None:
            return
        if len(data) > length:
            raise IndexError

        for idx, elt in enumerate(data):
            self.castbuffer[idx] = serialize_to_c(self._dtype, elt)

    def __del__(self):
        """
        Frees the mallocbuffer stored inside.
        """
        if hasattr(self, "buffer"):
            MallocBufferBackend.library(self.ftype).free(self.buffer)

    @property
    def castbuffer(self):
        return ctypes.cast(self.buffer.contents.data, ctypes.POINTER(self._c_dtype))

    @property
    def ftype(self):
        """
        Returns the ftype of the buffer, which is a MallocBufferFType.
        """
        return MallocBufferFType(self._dtype)

    # TODO should be property
    def length(self):
        return np.intp(self.buffer.contents.length)

    def load(self, index):
        value = self.castbuffer[index]
        return construct_from_c(self.ftype.element_type, value)

    def store(self, index, value):
        value = serialize_to_c(self.ftype.element_type, value)
        self.castbuffer[index] = value

    def resize(self, new_length):
        # we effectively have to run an entire serialization cycle ourselves.
        old_length = serialize_to_c(self.ftype.length_type, self.length())
        new_length = serialize_to_c(self.ftype.length_type, new_length)
        newptr = MallocBufferBackend.library(self.ftype).resize(
            self.castbuffer, old_length, new_length
        )
        self.buffer.contents.data = ctypes.cast(newptr, ctypes.c_void_p)
        self.buffer.contents.length = new_length

    def __str__(self):
        array = self.castbuffer[: self.length()]
        return f"malloc_buf({array})"

ftype property

Returns the ftype of the buffer, which is a MallocBufferFType.

__del__()

Frees the mallocbuffer stored inside.

Source code in src/finch/codegen/buffers/malloc_buffer.py

def __del__(self):
    """
    Frees the mallocbuffer stored inside.
    """
    if hasattr(self, "buffer"):
        MallocBufferBackend.library(self.ftype).free(self.buffer)

__init__(length, dtype, data=None)

Constructor for the MallocBuffer class.

length (int): the length of the malloc array. dtype (FType): the type that the buffer will be based on. data (optional): a list of data to initialize the buffer with.

Source code in src/finch/codegen/buffers/malloc_buffer.py

def __init__(self, length: int, dtype: FType, data=None):
    """
    Constructor for the MallocBuffer class.

    length (int): the length of the malloc array.
    dtype (FType): the type that the buffer will be based on.
    data (optional): a list of data to initialize the buffer with.
    """
    self._dtype = dtype
    self._c_dtype = c_type(self._dtype)
    self.buffer = ctypes.pointer(CMallocBufferStruct())

    MallocBufferBackend.library(self.ftype).init(
        self.buffer,
        serialize_to_c(self.length_type, ctypes.sizeof(self._c_dtype)),
        serialize_to_c(self.length_type, length),
    )
    if data is None:
        return
    if len(data) > length:
        raise IndexError

    for idx, elt in enumerate(data):
        self.castbuffer[idx] = serialize_to_c(self._dtype, elt)

MallocBufferFType

Bases: CBufferFType, CUnpackableFType

A ftype for buffers that uses libc-provided malloc functions. This is a concrete implementation of the BufferFType class.

This does not support the numba backend.

Source code in src/finch/codegen/buffers/malloc_buffer.py

class MallocBufferFType(CBufferFType, CUnpackableFType):
    """
    A ftype for buffers that uses libc-provided malloc functions. This is a
    concrete implementation of the BufferFType class.

    This does not support the numba backend.
    """

    def __init__(self, element_type: FType):
        self._element_type = element_type

    def __eq__(self, other):
        if not isinstance(other, MallocBufferFType):
            return False
        return self._element_type == other._element_type

    def __str__(self):
        return f"malloc_buf_t({qual_str(self._element_type)})"

    def __repr__(self):
        return f"MallocBufferFType({qual_str(self._element_type)})"

    @property
    def length_type(self):
        """
        Returns the type used for the length of the buffer.
        """
        return ftypes.intp

    @property
    def element_type(self):
        """
        Returns the type of elements stored in the buffer. This will be a ctypes array.
        """
        return self._element_type

    def __hash__(self):
        return hash(("MallocBufferFType", self._element_type))

    def __call__(self, len: int = 0):
        return MallocBuffer(len, self._element_type)

    def c_type(self):
        return ctypes.POINTER(CMallocBufferStruct)

    def c_length(self, ctx: CContext, buf: CBufferFields):
        return buf.length

    def c_data(self, ctx: CContext, buf: CBufferFields):
        return buf.data

    def c_load(self, ctx: CContext, buf: CBufferFields, idx: AssemblyExpression):
        return f"({buf.data})[{ctx(idx)}]"

    def c_store(
        self,
        ctx: CContext,
        buf: CBufferFields,
        idx: AssemblyExpression,
        value: AssemblyExpression,
    ):
        ctx.exec(f"{ctx.feed}({buf.data})[{ctx(idx)}] = {ctx(value)};")

    def c_resize(self, ctx: CContext, buf: CBufferFields, new_len: AssemblyExpression):
        if self not in ctx.buffer_methods:
            raise Exception("A Mallocbuffer must be unpacked before being operated on!")

        methods: CMallocBufferMethods = ctx.buffer_methods[self]

        new_len = ctx(ctx.cache("len", new_len))
        data = buf.data
        length = buf.length

        ctx.exec(
            f"{ctx.feed}{data} = {methods.resize}({data}, {length}, {new_len});\n"
            f"{ctx.feed}{length} = {new_len};"
        )
        return

    def c_unpack(self, ctx: CContext, var_n, val):
        """
        Unpack the malloc buffer into C context.
        """
        data = ctx.freshen(var_n, "data")
        length = ctx.freshen(var_n, "length")
        t = ctx.ctype_name(c_type(self.element_type))
        ctx.add_header("#include <stddef.h>")

        if self not in ctx.buffer_methods:
            MallocBufferBackend.gen_code(ctx, self, inline=True)

        ctx.exec(
            f"{ctx.feed}{t}* {data} = ({t}*){ctx(val)}->data;\n"
            f"{ctx.feed}size_t {length} = {ctx(val)}->length;"
        )

        return CBufferFields(data, length, var_n)

    def c_repack(self, ctx, var_n: str, obj: CBufferFields):
        """
        Repack the buffer from C context.
        """
        ctx.exec(
            f"{ctx.feed}{var_n}->data = (void*){obj.data};\n"
            f"{ctx.feed}{var_n}->length = {obj.length};"
        )
        return

    def serialize_to_c(self, obj: MallocBuffer):
        """
        Serialize the Malloc buffer to a C-compatible structure.
        This is trivial.
        """
        return obj.buffer

    def deserialize_from_c(self, obj, c_buffer):
        pass
        # this is handled by the resize callback

    def construct_from_c(self, c_buffer):
        """
        Construct a MallocBuffer from a C-compatible structure.
        """
        return c_buffer.contents

element_type property

Returns the type of elements stored in the buffer. This will be a ctypes array.

length_type property

Returns the type used for the length of the buffer.

c_repack(ctx, var_n, obj)

Repack the buffer from C context.

Source code in src/finch/codegen/buffers/malloc_buffer.py

def c_repack(self, ctx, var_n: str, obj: CBufferFields):
    """
    Repack the buffer from C context.
    """
    ctx.exec(
        f"{ctx.feed}{var_n}->data = (void*){obj.data};\n"
        f"{ctx.feed}{var_n}->length = {obj.length};"
    )
    return

c_unpack(ctx, var_n, val)

Unpack the malloc buffer into C context.

Source code in src/finch/codegen/buffers/malloc_buffer.py

def c_unpack(self, ctx: CContext, var_n, val):
    """
    Unpack the malloc buffer into C context.
    """
    data = ctx.freshen(var_n, "data")
    length = ctx.freshen(var_n, "length")
    t = ctx.ctype_name(c_type(self.element_type))
    ctx.add_header("#include <stddef.h>")

    if self not in ctx.buffer_methods:
        MallocBufferBackend.gen_code(ctx, self, inline=True)

    ctx.exec(
        f"{ctx.feed}{t}* {data} = ({t}*){ctx(val)}->data;\n"
        f"{ctx.feed}size_t {length} = {ctx(val)}->length;"
    )

    return CBufferFields(data, length, var_n)

construct_from_c(c_buffer)

Construct a MallocBuffer from a C-compatible structure.

Source code in src/finch/codegen/buffers/malloc_buffer.py

def construct_from_c(self, c_buffer):
    """
    Construct a MallocBuffer from a C-compatible structure.
    """
    return c_buffer.contents

serialize_to_c(obj)

Serialize the Malloc buffer to a C-compatible structure. This is trivial.

Source code in src/finch/codegen/buffers/malloc_buffer.py

def serialize_to_c(self, obj: MallocBuffer):
    """
    Serialize the Malloc buffer to a C-compatible structure.
    This is trivial.
    """
    return obj.buffer

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.
    """

NumpyBuffer

Bases: Buffer

A buffer that uses NumPy arrays to store data. This is a concrete implementation of the Buffer class.

Source code in src/finch/codegen/buffers/numpy_buffer.py

class NumpyBuffer(Buffer):
    """
    A buffer that uses NumPy arrays to store data. This is a concrete implementation
    of the Buffer class.
    """

    def __init__(self, arr: np.ndarray):
        if not arr.flags["C_CONTIGUOUS"]:
            raise ValueError("NumPy array must be C-contiguous")
        self.arr = arr

    @property
    def ftype(self):
        """
        Returns the ftype of the buffer, which is a NumpyBufferFType.
        """
        return NumpyBufferFType(ftype(self.arr.dtype))

    # TODO should be property
    def length(self):
        return self.arr.size

    def load(self, index: int):
        value = self.arr[index]
        if isinstance(self.ftype.element_type, TupleFType):
            return tuple(
                value[name] for name in self.ftype.element_type.struct_fieldnames
            )
        return value

    def store(self, index: int, value):
        self.arr[index] = value

    def resize(self, new_length: int):
        self.arr = np.resize(self.arr, new_length)

    def __str__(self):
        arr_str = str(self.arr).replace("\n", "")
        return f"np_buf({arr_str})"

    def __repr__(self):
        arr_repr = repr(self.arr).replace("\n", "")
        return f"NumpyBuffer({arr_repr})"

ftype property

Returns the ftype of the buffer, which is a NumpyBufferFType.

NumpyBufferFType

Bases: CBufferFType, NumbaBufferFType, MLIRBufferFType, CUnpackableFType

A ftype for buffers that uses NumPy arrays. This is a concrete implementation of the BufferFType class.

Source code in src/finch/codegen/buffers/numpy_buffer.py

class NumpyBufferFType(
    CBufferFType, NumbaBufferFType, MLIRBufferFType, CUnpackableFType
):
    """
    A ftype for buffers that uses NumPy arrays. This is a concrete implementation
    of the BufferFType class.
    """

    def __init__(self, element_type: FType):
        self._element_type = ftype(np_dtype(ftype(element_type)))

    @property
    def _dtype(self):
        return np_dtype(self._element_type)

    def __eq__(self, other):
        if not isinstance(other, NumpyBufferFType):
            return False
        return self._element_type == other._element_type

    def __str__(self):
        return f"np_buf_t({qual_str(self._dtype.type)})"

    def __repr__(self):
        return f"NumpyBufferFType({repr(self._element_type)})"

    @property
    def length_type(self):
        """
        Returns the type used for the length of the buffer.
        """
        return ftypes.intp

    @property
    def element_type(self):
        """
        Returns the type of elements stored in the buffer.
        This is typically the same as the dtype used to create the buffer.
        """
        return self._element_type

    def __hash__(self):
        return hash(self._element_type)

    def __call__(self, len: int = 0, element_type: FType | None = None):
        return NumpyBuffer(np.zeros(len, dtype=self._dtype))

    def c_type(self):
        return ctypes.POINTER(CNumpyBuffer)

    def c_length(self, ctx: "CContext", buf: CBufferFields):
        return buf.length

    def c_data(self, ctx: "CContext", buf: CBufferFields):
        return buf.data

    def c_load(self, ctx: "CContext", buf: CBufferFields, idx: "AssemblyExpression"):
        return f"({buf.data})[{ctx(idx)}]"

    def c_store(
        self,
        ctx: "CContext",
        buf: CBufferFields,
        idx: "AssemblyExpression",
        value: "AssemblyExpression",
    ):
        ctx.exec(f"{ctx.feed}({buf.data})[{ctx(idx)}] = {ctx(value)};")

    def c_resize(self, ctx, buf: CBufferFields, new_len):
        new_len = ctx(ctx.cache("len", new_len))
        data = buf.data
        length = buf.length
        obj = buf.obj
        t = ctx.ctype_name(c_type(self.element_type))
        ctx.exec(
            f"{ctx.feed}{data} = ({t}*){obj}->resize(&{obj}->arr, {new_len});\n"
            f"{ctx.feed}{length} = {new_len};"
        )
        return

    def c_unpack(self, ctx, var_n, val):
        """
        Unpack the buffer into C context.
        """
        data = ctx.freshen(var_n, "data")
        length = ctx.freshen(var_n, "length")
        t = ctx.ctype_name(c_type(self.element_type))
        ctx.add_header("#include <stddef.h>")
        ctx.exec(
            f"{ctx.feed}{t}* {data} = ({t}*){ctx(val)}->data;\n"
            f"{ctx.feed}size_t {length} = {ctx(val)}->length;"
        )

        return CBufferFields(data, length, var_n)

    def c_repack(self, ctx, lhs, obj):
        """
        Repack the buffer from C context.
        """
        ctx.exec(
            f"{ctx.feed}{lhs}->data = (void*){obj.data};\n"
            f"{ctx.feed}{lhs}->length = {obj.length};"
        )
        return

    def serialize_to_c(self, obj):
        """
        Serialize the NumPy buffer to a C-compatible structure.
        """
        data = ctypes.c_void_p(obj.arr.ctypes.data)
        length = obj.arr.size
        obj._self_obj = ctypes.py_object(obj)
        obj._c_callback = numpy_buffer_resize_callback
        obj._c_buffer = CNumpyBuffer(obj._self_obj, data, length, obj._c_callback)
        return ctypes.pointer(obj._c_buffer)

    def deserialize_from_c(self, obj, c_buffer):
        """
        Update this buffer based on how the C call modified the CNumpyBuffer structure.
        """
        # this is handled by the resize callback

    def construct_from_c(self, c_buffer):
        """
        Construct a NumpyBuffer from a C-compatible structure.
        """
        return c_buffer.contents.arr

    def numba_type(self) -> type:
        return list[np.ndarray]

    def numba_jitclass_type(self) -> numba.types.Type:
        return numba.types.ListType(
            numba.types.Array(numba.from_dtype(self._dtype), 1, "C")
        )

    def numba_length(self, ctx, buf: NumbaBufferFields):
        arr = buf.arr
        return f"len({arr})"

    def numba_load(self, ctx, buf, idx):
        buf = cast(NumbaBufferFields, buf)
        arr = buf.arr
        if isinstance(self.element_type, TupleFType):
            idx = ctx(ctx.cache("idx", idx))
            fields = ", ".join(
                f"{arr}[{idx}]['{name}']"
                for name in self.element_type.struct_fieldnames
            )
            return f"({fields},)"
        return f"{arr}[{ctx(idx)}]"

    def numba_store(self, ctx, buf, idx, value=None):
        buf = cast(NumbaBufferFields, buf)
        arr = buf.arr
        if isinstance(self.element_type, TupleFType):
            idx = ctx(ctx.cache("idx", idx))
            value = ctx.cache("val", value)
            val_code = ctx(value)
            for i, name in enumerate(self.element_type.struct_fieldnames):
                ctx.exec(f"{ctx.feed}{arr}[{idx}]['{name}'] = {val_code}[{i}]")
            return
        ctx.exec(f"{ctx.feed}{arr}[{ctx(idx)}] = {ctx(value)}")

    def numba_resize(self, ctx, buf: NumbaBufferFields, new_len):
        arr = buf.arr
        ctx.exec(f"{ctx.feed}{arr} = numpy.resize({arr}, {ctx(new_len)})")

    def numba_unpack(self, ctx, var_n, val):
        """
        Unpack the buffer into Numba context.
        """
        arr = ctx.freshen(var_n, "arr")
        ctx.exec(f"{ctx.feed}{arr} = {ctx(val)}[0]")

        return NumbaBufferFields(arr, var_n)

    def numba_repack(self, ctx, lhs, obj):
        """
        Repack the buffer from Numba context.
        """
        ctx.exec(f"{ctx.feed}{lhs}[0] = {obj.arr}")
        return

    def serialize_to_numba(self, obj):
        """
        Serialize the NumPy buffer to a Numba-compatible object.
        """
        return numba.typed.List([obj.arr])

    def deserialize_from_numba(self, obj, numba_buffer):
        obj.arr = numba_buffer[0]
        return

    def construct_from_numba(self, numba_buffer):
        """
        Construct a NumpyBuffer from a Numba-compatible object.
        """
        return NumpyBuffer(numba_buffer[0])

    def mlir_type(self):
        return "!llvm.ptr"

    def mlir_buffer_type(self):
        return f"memref<?x{mlir_type(self.element_type)}>"

    # this is the rank-1 memref descriptor type
    def mlir_descriptor_type(self):
        return "!llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>"

    # this is the callback buffer type
    def llvm_callback_type(self):
        return "!llvm.struct<(ptr, ptr, i64, ptr)>"

    # Create our own rank-1 memeref descriptor from a data pointer and length.
    def mlir_memref_from_pointer(self, ctx, data, length):
        memref_t = self.mlir_buffer_type()
        desc_t = self.mlir_descriptor_type()

        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.mlir.undef : {desc_t}")

        for id, val in (
            ("0", data),
            ("1", data),
            ("2", ctx.constant(0, "i64")),
            ("3, 0", length),
            ("4, 0", ctx.constant(1, "i64")),
        ):
            fresh = ctx.new_ssa()
            ctx.exec(
                f"{ctx.feed}{fresh} = llvm.insertvalue {val}, {desc}[{id}] : {desc_t}"
            )
            desc = fresh

        buf = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buf} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {memref_t}"
        )
        return buf

    # load the length of the NumPy buffer
    def mlir_length(self, ctx: "MLIRContext", buf: MLIRBufferFields):
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {buf.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {self.mlir_buffer_type()}"
        )
        c0 = ctx.constant(0, "index")
        res = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{res} = memref.dim {buffer}, {c0} : {self.mlir_buffer_type()}"
        )
        return res

    # Convert index values to mlir index types
    def mlir_index(self, ctx: "MLIRContext", idx: "AssemblyExpression"):
        i = ctx(idx)
        t = mlir_type(idx.result_type)
        if t == "index":
            return i
        cast_ = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{cast_} = arith.index_cast {i} : {t} to index")
        return cast_

    # Load buffer values of NumPy buffer
    def mlir_load(
        self, ctx: "MLIRContext", buf: MLIRBufferFields, idx: "AssemblyExpression"
    ):
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {buf.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {self.mlir_buffer_type()}"
        )
        i = self.mlir_index(ctx, idx)
        c0 = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{c0} = memref.load {buffer}[{i}] : {self.mlir_buffer_type()}"
        )
        return c0

    # Store values in the buffer
    def mlir_store(
        self,
        ctx: "MLIRContext",
        buf: Any,
        idx: Any,
        value: Any = None,
    ):
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {buf.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {self.mlir_buffer_type()}"
        )
        new = ctx(value)
        i = self.mlir_index(ctx, idx)
        ctx.exec(
            f"{ctx.feed}memref.store {new}, {buffer}[{i}] : {self.mlir_buffer_type()}"
        )

    # Resize the NumPy buffer and update its memref descriptor.
    def mlir_resize(self, ctx, buf: MLIRBufferFields, new_len):
        len_idx = mlir_cast_value(ctx, ctx(new_len), new_len.result_type, ftypes.intp)
        len = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{len} = arith.index_cast {len_idx} : index to i64")
        arr_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{arr_ptr} = llvm.getelementptr {buf.obj}[0, 0] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )

        cb_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{cb_ptr} = llvm.getelementptr {buf.obj}[0, 3] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        cb = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{cb} = llvm.load {cb_ptr} : !llvm.ptr -> !llvm.ptr")

        data = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data} = llvm.call {cb}({arr_ptr}, {len}) "
            f": !llvm.ptr, (!llvm.ptr, i64) -> !llvm.ptr"
        )

        buffer = self.mlir_memref_from_pointer(ctx, data, len)
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{desc} = builtin.unrealized_conversion_cast "
            f"{buffer} : {self.mlir_buffer_type()} to {desc_t}"
        )
        ctx.exec(f"{ctx.feed}llvm.store {desc}, {buf.box} : {desc_t}, !llvm.ptr")

    # Unpack the NumPy buffer and create its memeref descriptor
    def mlir_unpack(self, ctx: MLIRContext, _var_n, val):
        obj = ctx(val)

        data_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data_ptr} = llvm.getelementptr {obj}[0, 1] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        data = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{data} = llvm.load {data_ptr} : !llvm.ptr -> !llvm.ptr")

        length_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{length_ptr} = llvm.getelementptr {obj}[0, 2] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        length = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{length} = llvm.load {length_ptr} : !llvm.ptr -> i64")

        buffer = self.mlir_memref_from_pointer(ctx, data, length)
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{desc} = builtin.unrealized_conversion_cast "
            f"{buffer} : {self.mlir_buffer_type()} to {desc_t}"
        )
        box = ctx.new_ssa()
        count = ctx.constant(1, "i64")
        ctx.exec(
            f"{ctx.feed}{box} = llvm.alloca {count} x {desc_t} : (i64) -> !llvm.ptr"
        )
        ctx.exec(f"{ctx.feed}llvm.store {desc}, {box} : {desc_t}, !llvm.ptr")

        return MLIRBufferFields(box, obj)

    # Update the buffer wrapper with the resized pointer and length.
    def mlir_repack(self, ctx: MLIRContext, _lhs, obj):
        memref_t = self.mlir_buffer_type()
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {obj.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {memref_t}"
        )

        len_ssa = self.mlir_length(ctx, obj)
        len_i64 = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{len_i64} = arith.index_cast {len_ssa} : index to i64")
        len_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{len_ptr} = llvm.getelementptr {obj.obj}[0, 2] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        ctx.exec(f"{ctx.feed}llvm.store {len_i64}, {len_ptr} : i64, !llvm.ptr")

        data_ssa = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data_ssa} = "
            f"memref.extract_aligned_pointer_as_index {buffer} "
            f": {memref_t} -> index"
        )
        data_i64 = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{data_i64} = arith.index_cast {data_ssa} : index to i64")
        data_llvm = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{data_llvm} = llvm.inttoptr {data_i64} : i64 to !llvm.ptr")
        data_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data_ptr} = llvm.getelementptr {obj.obj}[0, 1] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        ctx.exec(f"{ctx.feed}llvm.store {data_llvm}, {data_ptr} : !llvm.ptr, !llvm.ptr")

    # Package the NumPy array, data pointer, length, and resize callback for MLIR
    def serialize_to_mlir(self, obj):
        result = MLIRNumpyBuffer(
            ctypes.py_object(obj),
            ctypes.c_void_p(obj.arr.ctypes.data),
            obj.arr.size,
            numpy_buffer_resize_callback,
        )
        return ctypes.cast(ctypes.pointer(result), ctypes.c_void_p)

    def deserialize_from_mlir(self, obj, mlir_buffer):
        # The resize callback updates the owning NumPy buffer directly.
        pass

    def construct_from_mlir(self, mlir_buffer):
        # zero copy
        result = ctypes.cast(mlir_buffer, ctypes.POINTER(MLIRNumpyBuffer))
        return result.contents.arr

element_type property

Returns the type of elements stored in the buffer. This is typically the same as the dtype used to create the buffer.

length_type property

Returns the type used for the length of the buffer.

c_repack(ctx, lhs, obj)

Repack the buffer from C context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def c_repack(self, ctx, lhs, obj):
    """
    Repack the buffer from C context.
    """
    ctx.exec(
        f"{ctx.feed}{lhs}->data = (void*){obj.data};\n"
        f"{ctx.feed}{lhs}->length = {obj.length};"
    )
    return

c_unpack(ctx, var_n, val)

Unpack the buffer into C context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def c_unpack(self, ctx, var_n, val):
    """
    Unpack the buffer into C context.
    """
    data = ctx.freshen(var_n, "data")
    length = ctx.freshen(var_n, "length")
    t = ctx.ctype_name(c_type(self.element_type))
    ctx.add_header("#include <stddef.h>")
    ctx.exec(
        f"{ctx.feed}{t}* {data} = ({t}*){ctx(val)}->data;\n"
        f"{ctx.feed}size_t {length} = {ctx(val)}->length;"
    )

    return CBufferFields(data, length, var_n)

construct_from_c(c_buffer)

Construct a NumpyBuffer from a C-compatible structure.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def construct_from_c(self, c_buffer):
    """
    Construct a NumpyBuffer from a C-compatible structure.
    """
    return c_buffer.contents.arr

construct_from_numba(numba_buffer)

Construct a NumpyBuffer from a Numba-compatible object.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def construct_from_numba(self, numba_buffer):
    """
    Construct a NumpyBuffer from a Numba-compatible object.
    """
    return NumpyBuffer(numba_buffer[0])

deserialize_from_c(obj, c_buffer)

Update this buffer based on how the C call modified the CNumpyBuffer structure.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def deserialize_from_c(self, obj, c_buffer):
    """
    Update this buffer based on how the C call modified the CNumpyBuffer structure.
    """

numba_repack(ctx, lhs, obj)

Repack the buffer from Numba context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def numba_repack(self, ctx, lhs, obj):
    """
    Repack the buffer from Numba context.
    """
    ctx.exec(f"{ctx.feed}{lhs}[0] = {obj.arr}")
    return

numba_unpack(ctx, var_n, val)

Unpack the buffer into Numba context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def numba_unpack(self, ctx, var_n, val):
    """
    Unpack the buffer into Numba context.
    """
    arr = ctx.freshen(var_n, "arr")
    ctx.exec(f"{ctx.feed}{arr} = {ctx(val)}[0]")

    return NumbaBufferFields(arr, var_n)

serialize_to_c(obj)

Serialize the NumPy buffer to a C-compatible structure.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def serialize_to_c(self, obj):
    """
    Serialize the NumPy buffer to a C-compatible structure.
    """
    data = ctypes.c_void_p(obj.arr.ctypes.data)
    length = obj.arr.size
    obj._self_obj = ctypes.py_object(obj)
    obj._c_callback = numpy_buffer_resize_callback
    obj._c_buffer = CNumpyBuffer(obj._self_obj, data, length, obj._c_callback)
    return ctypes.pointer(obj._c_buffer)

serialize_to_numba(obj)

Serialize the NumPy buffer to a Numba-compatible object.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def serialize_to_numba(self, obj):
    """
    Serialize the NumPy buffer to a Numba-compatible object.
    """
    return numba.typed.List([obj.arr])