types

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

ftype_to_jl_constructor_str(ftype)

Julia source text constructing a minimal instance of ftype, for embedding directly in generated kernel source (see compiler.py).

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

def ftype_to_jl_constructor_str(ftype: FType) -> str:
    """Julia source text constructing a minimal instance of `ftype`, for
    embedding directly in generated kernel source (see compiler.py)."""
    if isinstance(ftype, ScalarFType):
        elem_t = _leaf_type_str(ftype.element_type)
        fill = _julia_literal(ftype.fill_value)
        return f"Finch.Tensor(Finch.ElementLevel({fill}, {elem_t}[]))"
    if isinstance(ftype, FiberTensorFType):
        return f"Finch.Tensor({_level_constructor_str(ftype.lvl_t)})"
    if isinstance(ftype, BufferizedNDArrayFType):
        # Matches interop.py's _ndarray_to_jl_tensor: a plain dense buffer
        # wrapped in `ndim` nested DenseLevels, each with plain Int64 shape
        # (regardless of the ftype's own dimension_type).
        elem_t = _leaf_type_str(ftype.element_type)
        fill = _julia_literal(ftype.fill_value)
        ctor = f"Finch.ElementLevel({fill}, {elem_t}[])"
        for _ in range(ftype.ndim):
            ctor = f"Finch.DenseLevel({ctor}, 1)"
        return f"Finch.Tensor({ctor})"
    if isinstance(ftype, FillTensorFType):
        # Matches interop.py, which passes a fill tensor as a Finch scalar.
        return f"{scalar_type_str(ftype.fill_value, ftype.element_type)}()"
    match ftype:
        case PatternTensorFType():
            shape = (1,) * ftype.ndim
            obj = ftype.construct(shape)
            # This matches the masks built by `_pattern_tensor_to_jl`.
            swizzle = False
            match obj:
                case EyeTensor():
                    ctor = "Finch.diagmask"
                    if obj._k:
                        ctor = f"Finch.offset({ctor}, 0, {int(obj._k)})"
                case UpperTriangleTensor():
                    ctor = f"Finch.offset(Finch.lotrimask, 0, {int(obj._k)})"
                case LowerTriangleTensor():
                    ctor = f"Finch.offset(Finch.uptrimask, 0, {int(obj._k)})"
                case PairSumTensor():
                    ctor = "Finch.repeatmask(2)"
                case PairCarryTensor():
                    ctor = "Finch.offset(Finch.pairsummask, 0, -1)"
                case ReverseTensor():
                    ctor = f"Finch.reversemask({int(obj.shape[1])})"
                case RollTensor():
                    ctor = f"Finch.rollmask({int(obj.shape[1])}, {-int(obj._k)})"
                case RepeatTensor():
                    ctor = f"Finch.repeatmask({int(obj._k)})"
                case ChunkMaskTensor():
                    ctor = f"Finch.chunkmask({int(obj.shape[1])}, {int(obj._b)})"
                case SplitMaskTensor():
                    ctor = f"Finch.splitmask({int(obj.shape[1])}, {int(obj.shape[0])})"
                case RandomMaskTensor():
                    ctor = (
                        f"Finch.randommask({shape}, {obj._p}; seed=UInt64({obj._seed}))"
                    )
                case OddEvenMergeSortPartnerMaskTensor():
                    ctor = (
                        "Finch.oddevenmergesortpartnermask("
                        f"{int(obj.shape[1])}, {int(obj._p)}, {int(obj._k)})"
                    )
                case OddEvenMergeSortLowerMaskTensor():
                    ctor = (
                        "Finch.oddevenmergesortlowermask("
                        f"{int(obj.shape[0])}, {int(obj._p)}, {int(obj._k)})"
                    )
                case OneHotMaskTensor():
                    ctor = f"Finch.onehotmask({int(obj._index) + 1})"
                case ParityMaskTensor():
                    ctor = f"Finch.paritymask({int(obj._parity)})"
                case ReshapeMaskTensor():
                    old_shape = tuple(int(dim) for dim in obj._old_shape)
                    new_shape = tuple(int(dim) for dim in obj._new_shape)
                    ctor = f"Finch.reshapemask({old_shape}, {new_shape})"
                    swizzle = True
                case _:
                    raise ValueError(
                        f"Unsupported Julia pattern tensor type: {type(obj)}"
                    )
            dims = ", ".join("Finch.Extent(1, 1)" for _ in shape)
            ctor = f"Finch.window({ctor}{', ' if dims else ''}{dims})"
            if swizzle and ftype.ndim > 1:
                axes = ", ".join(map(str, reversed(range(1, ftype.ndim + 1))))
                ctor = f"Finch.swizzle({ctor}, {axes})"
            return ctor
    raise NotImplementedError(
        f"ftype_to_jl_constructor_str: unsupported ftype kind {type(ftype).__name__}"
    )

ftype_to_jl_type_str(ftype)

Julia type as source text (e.g. for a cache key). Cached per ftype.

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

def ftype_to_jl_type_str(ftype: FType) -> str:
    """Julia type as source text (e.g. for a cache key). Cached per ftype."""
    cached = _TYPE_STR_CACHE.get(ftype)
    if cached is not None:
        return cached
    jl = get_jl()
    ctor = ftype_to_jl_constructor_str(ftype)
    type_str = str(jl.string(jl.typeof(jl.seval(ctor))))
    _TYPE_STR_CACHE[ftype] = type_str
    return type_str

scalar_type_str(fill, element_type)

The Julia type of a Finch scalar with the given fill and element type.

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

def scalar_type_str(fill: Any, element_type: Any) -> str:
    """The Julia type of a Finch scalar with the given fill and element type."""
    return f"Finch.Scalar{{{_julia_literal(fill)}, {_leaf_type_str(element_type)}}}"

to_fl_dtype(x)

Normalize a Julia DataType, numpy dtype/scalar type, Python builtin type, or finch FType into the corresponding finch FType.

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

def to_fl_dtype(x) -> FType:
    """Normalize a Julia DataType, numpy dtype/scalar type, Python builtin
    type, or finch FType into the corresponding finch FType."""
    if isinstance(x, FType):
        return x
    try:
        return ft.ftype(x)
    except NotImplementedError:
        pass
    try:
        return _jl_dtype_to_fl()[x]
    except (KeyError, TypeError):
        raise NotImplementedError(f"Cannot convert {x!r} to a Finch dtype") from None