interop

JuliaBufferContext

Keep Julia-owned tensor buffers alive across kernel invocations.

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

class JuliaBufferContext:
    """Keep Julia-owned tensor buffers alive across kernel invocations."""

    def __init__(self):
        self._tensors: dict[tuple[Any, ...], tuple[Any, Any]] = {}

    @staticmethod
    def _cache_key(obj):
        # defer() can create a fresh wrapper around the same NumPy
        # allocation, so wrapper identity alone would miss reuse.
        if isinstance(obj, BufferizedNDArray):
            arr = obj.to_numpy()
            pointer = arr.__array_interface__["data"][0]
            return ("numpy", pointer, arr.shape, arr.strides, arr.dtype.str)
        if isinstance(obj, NumPyWrapper):
            arr = obj._data
            pointer = arr.__array_interface__["data"][0]
            return ("numpy", pointer, arr.shape, arr.strides, arr.dtype.str)
        # FiberTensors reuse their ids so we restrict cache keys to id.
        return ("object", id(obj))

    def tensor_to_jl(self, obj, *, pin_fill: bool = False):
        key = (self._cache_key(obj), pin_fill)
        cached = self._tensors.get(key)
        if cached is not None:
            return cached[1]

        jl_obj = tensor_to_jl(obj, pin_fill=pin_fill)
        self._tensors[key] = (obj, jl_obj)
        return jl_obj

    def tensor_to_python(self, obj):
        result = jl_tensor_to_python(obj)
        if isinstance(result, FiberTensor):
            self._tensors[self._cache_key(result)] = (result, obj)
        return result

    def close(self):
        self._tensors.clear()

level_to_jl(level, pin_fill=False)

Convert a level to its Julia counterpart, optionally pinning its fill.

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

def level_to_jl(level: Level, pin_fill: bool = False):
    """Convert a level to its Julia counterpart, optionally pinning its fill."""
    match level:
        case ElementLevel():
            fill = level.fill_value
            if pin_fill:
                fill = jl_dtypes._zero_value(ftype(fill))
            return jl.ElementLevel(
                _as_julia_scalar(fill),
                _buffer_to_jl(level.val),
            )
        case DenseLevel(lvl=lvl, dimension=dimension):
            return jl.DenseLevel(level_to_jl(lvl, pin_fill), int(dimension))
        case SparseListLevel(lvl=lvl, dimension=dimension, ptr=ptr, idx=idx):
            if ptr is None or idx is None:
                raise ValueError("SparseListLevel must have ptr and idx buffers")
            return jl.SparseListLevel(
                level_to_jl(lvl, pin_fill),
                int(dimension),
                _plus_one_buffer_to_jl(cast(Buffer, ptr)),
                _plus_one_buffer_to_jl(cast(Buffer, idx)),
            )
        case SparseByteMapLevel(
            lvl=lvl, dimension=dimension, ptr=ptr, tbl=tbl, srt=srt
        ):
            if ptr is None or tbl is None or srt is None:
                raise ValueError(
                    "SparseByteMapLevel must have ptr, tbl, and srt buffers"
                )
            return jl.SparseByteMapLevel(
                level_to_jl(lvl, pin_fill),
                int(dimension),
                _plus_one_buffer_to_jl(cast(Buffer, ptr)),
                _buffer_to_jl(cast(Buffer, tbl)),
                _plus_one_buffer_to_jl(cast(Buffer, srt)),
            )
        case SparseCOOLevel(lvl=lvl, coo_shape=coo_shape, ptr=ptr, tbl=tbl):
            assert isinstance(tbl, tuple)
            # Julia's COO lists its dims fastest first, Python's slowest first
            return jl.SparseCOOLevel(
                level_to_jl(lvl, pin_fill),
                tuple(int(dim) for dim in reversed(coo_shape)),
                _plus_one_buffer_to_jl(ptr),
                tuple(_plus_one_buffer_to_jl(idx) for idx in reversed(tbl)),
            )
        case SparseHashLevel(
            lvl=lvl,
            dimension=dimension,
            ptr=ptr,
            tbl_ctrl=tbl_ctrl,
            tbl=tbl,
            pool=pool,
            perm=perm,
            subtables=subtables,
            single_writer=single_writer,
        ):
            if (
                ptr is None
                or tbl_ctrl is None
                or tbl is None
                or pool is None
                or perm is None
            ):
                raise ValueError(
                    "SparseHashLevel must have ptr, tbl_ctrl, tbl, pool, and perm "
                    "buffers"
                )
            dimension = _as_julia_scalar(np.asarray(dimension).item())
            constructor = jl.SparseHashLevel[(jl.typeof(dimension), single_writer)]
            return constructor(
                level_to_jl(lvl, pin_fill),
                dimension,
                int(subtables),
                _plus_one_buffer_to_jl(cast(Buffer, ptr)),
                _buffer_to_jl(cast(Buffer, tbl_ctrl)),
                _buffer_to_jl(cast(Buffer, tbl), offset=1),
                _buffer_to_jl(cast(Buffer, pool)),
                _plus_one_buffer_to_jl(cast(Buffer, perm)),
            )
        case _:
            raise ValueError(f"Unsupported Finch level type: {type(level)}")

tensor_to_jl(obj, pin_fill=False)

Convert a tensor to its Julia counterpart, optionally pinning its fill.

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

def tensor_to_jl(obj, pin_fill: bool = False):
    """Convert a tensor to its Julia counterpart, optionally pinning its fill."""
    if hasattr(obj, "raw_julia_obj"):
        return obj.raw_julia_obj
    if is_julia_obj(obj) and jl.isa(obj, jl.Finch.Tensor):
        return obj
    if isinstance(obj, FiberTensor):
        if obj.pos != 0:
            raise ValueError("Only root-position FiberTensor objects can use Julia")
        return jl.Tensor(level_to_jl(obj.lvl, pin_fill))
    if isinstance(obj, BufferizedNDArray):
        fill = (
            jl_dtypes._zero_value(ftype(obj.fill_value)) if pin_fill else obj.fill_value
        )
        return _ndarray_to_jl_tensor(obj.to_numpy(), fill, copy=False)
    if isinstance(obj, NumPyWrapper):
        fill = (
            jl_dtypes._zero_value(ftype(obj.fill_value)) if pin_fill else obj.fill_value
        )
        return _ndarray_to_jl_tensor(obj._data, fill, copy=False)
    if isinstance(obj, Scalar):
        return scalar_to_jl(obj.val, pin_fill=pin_fill)
    if isinstance(obj, PatternTensor):
        return _pattern_tensor_to_jl(obj)
    if isinstance(obj, FillTensor):
        # A fill tensor is its fill everywhere, and its extents are passed to
        # kernels separately, so it is passed as a Finch scalar.
        fill = jl_dtypes._zero_value(obj.element_type) if pin_fill else obj.fill_value
        scalar_t = jl.seval(jl_dtypes.scalar_type_str(fill, obj.element_type))
        return scalar_t(_as_julia_scalar(obj.fill_value))
    if isinstance(obj, np.ndarray):
        fill = np.asarray(0, dtype=obj.dtype)[()]
        return _ndarray_to_jl_tensor(obj, fill, copy=False)
    if np.isscalar(obj):
        return scalar_to_jl(obj, pin_fill=pin_fill)
    raise ValueError(f"Unsupported Julia backend argument type: {type(obj)}")