compile_jl

DefaultFinchJLRuntime

Bases: FinchJLRuntime

Julia runtime with cached conversions and reusable result storage.

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

class DefaultFinchJLRuntime(FinchJLRuntime):
    """Julia runtime with cached conversions and reusable result storage."""

    def __init__(self) -> None:
        self._kernels: dict[Any, Any] = {}
        self._kernel_metadata: dict[int, _KernalMetadata] = {}
        self._translated_tensors: dict[_TranslationCacheKey, JuliaOwnedTensor] = {}
        self.free_pool = _StoragePool()

    def get_cached_kernel(self, key):
        return self._kernels.get(key)

    def cache_kernel(self, key, kernel):
        self._kernels[key] = kernel
        self._kernel_metadata[id(kernel)] = _KernalMetadata(
            reset_argument_positions(kernel.finch_program),
            returned_argument_positions(kernel.finch_program),
        )

    def kernel_call(self, func_name, kernel, args):
        metadata = self._kernel_metadata[id(kernel)]

        # Lease Julia storage only for resettable compiler-created outputs.
        julia_buf_args: list[JuliaOwnedTensor] = []
        for position, tensor in enumerate(args):
            pin_fill = position in getattr(kernel, "dynamic_args", ())
            if position in metadata.reset_positions and not isinstance(
                tensor, JuliaOwnedTensor
            ):
                lease = self.free_pool.acquire(
                    tensor.ftype,
                    tensor.shape,
                    pin_fill,
                )
                julia_buf = JuliaOwnedTensor(
                    tensor.ftype,
                    lease.key.shape,
                    self.release,
                    lease.raw,
                    lease.key.pin_fill,
                    lease,
                )
            else:
                julia_buf = self._to_julia_owned_tensor(tensor, pin_fill)
            julia_buf_args.append(julia_buf)

        raw_args = [arg.raw_julia_obj for arg in julia_buf_args]
        raw_args += [
            int(args[pos].shape[axis]) for pos, axis in getattr(kernel, "extents", ())
        ]
        getattr(jl, func_name)(*raw_args)

        # Associate returned tensors with their Python ownership handles.
        return tuple(
            julia_buf_args[position] for position in metadata.returned_positions
        )

    def _to_julia_owned_tensor(
        self, tensor: Tensor, pin_fill: bool = False
    ) -> JuliaOwnedTensor:
        if isinstance(tensor, JuliaOwnedTensor):
            if tensor._pin_fill == pin_fill:
                return tensor
            return self._to_julia_owned_tensor(tensor._as_tensor(), pin_fill)
        key = self._translation_cache_key(tensor, pin_fill)
        julia_tensor = self._translated_tensors.get(key)
        if julia_tensor is not None:
            return julia_tensor

        raw = tensor_to_jl(tensor, pin_fill=pin_fill)
        translation_finalizer = weakref.finalize(
            tensor, self._translated_tensors.pop, key, None
        )
        julia_tensor = JuliaOwnedTensor(
            tensor.ftype,
            tuple(int(dimension) for dimension in tensor.shape),
            self.release,
            raw,
            pin_fill,
            translation_finalizer=translation_finalizer,
        )
        self._translated_tensors[key] = julia_tensor
        return julia_tensor

    @staticmethod
    def _translation_cache_key(
        tensor: Tensor, pin_fill: bool = False
    ) -> _TranslationCacheKey:
        if isinstance(tensor, BufferizedNDArray):
            array = tensor.to_numpy()
            return _TranslationCacheKey(
                pin_fill,
                "numpy",
                # The first data-interface entry is the array's memory address.
                array.__array_interface__["data"][0],
                tuple(int(dimension) for dimension in array.shape),
                tuple(int(stride) for stride in array.strides),
                array.dtype.str,
            )
        if isinstance(tensor, NumPyWrapper):
            array = tensor._data
            return _TranslationCacheKey(
                pin_fill,
                "numpy",
                # The first data-interface entry is the array's memory address.
                array.__array_interface__["data"][0],
                tuple(int(dimension) for dimension in array.shape),
                tuple(int(stride) for stride in array.strides),
                array.dtype.str,
            )
        return _TranslationCacheKey(pin_fill, "object", id(tensor), None, None, None)

    def release(self, tensor: JuliaOwnedTensor) -> None:
        if tensor._lease is not None:
            self.free_pool.release_lease(tensor._lease)

    def close(self) -> None:
        self._kernels.clear()
        self._kernel_metadata.clear()
        self._translated_tensors.clear()

FinchJLRuntime

Bases: ABC

Runtime services used by the Julia compiler and its kernels.

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

class FinchJLRuntime(ABC):
    """Runtime services used by the Julia compiler and its kernels."""

    @abstractmethod
    def get_cached_kernel(self, key): ...

    @abstractmethod
    def cache_kernel(self, key, kernel): ...

    @abstractmethod
    def kernel_call(self, func_name, kernel, args): ...

JuliaElementFType

Bases: ABC

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

class JuliaElementFType(ABC):
    @abstractmethod
    def julia_type(self):
        """
        Return the Julia type used for elements with this ftype.
        """
        ...

    @abstractmethod
    def julia_value(self, value: Any, *, offset: int = 0):
        """
        Convert a Python value with this ftype into a Julia element value.
        """
        ...

    def julia_vector(self, values, *, offset: int = 0):
        jl = get_jl()
        return jc.convert(
            jl.Vector[self.julia_type()],
            [self.julia_value(value, offset=offset) for value in values],
        )

julia_type() abstractmethod

Return the Julia type used for elements with this ftype.

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

@abstractmethod
def julia_type(self):
    """
    Return the Julia type used for elements with this ftype.
    """
    ...

julia_value(value, *, offset=0) abstractmethod

Convert a Python value with this ftype into a Julia element value.

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

@abstractmethod
def julia_value(self, value: Any, *, offset: int = 0):
    """
    Convert a Python value with this ftype into a Julia element value.
    """
    ...