vector.py

from typing import List, Tuple
from numpy.typing import ArrayLike, DTypeLike

import dataclasses as dc
import numpy as np

import compute_graph as cg
import scalar_lang as sl
import opcodes


@dc.dataclass
class Vector(cg.Data):
    "A one-dimension array whose size and type is known at compile time."
    length: int
    dtype: np.dtype
    address: sl.AllocateMemory

    def set_name(self, num_allocations: int) -> None:
        self.name = f"v{num_allocations}"

    def referenced_ops(self) -> Tuple[sl.ScalarOp]:
        return (self.address,)

    def free(self, memory: sl.ScalarOp) -> sl.FreeMemory:
        return sl.FreeMemory(memory, self.address)


@dc.dataclass
class ConstVector(Vector):
    value: np.ndarray
    address: sl.AllocateMemory

    def set_name(self, num_allocations: int) -> None:
        self.name = f"cv{num_allocations}"


class VectorOp(sl.ScalarOp):
    def __hash__(self) -> int:
        return id(self)

    def AddressOf(self) -> sl.AllocateMemory:
        return self.out.address

    def BoundsCheck(self, offset: sl.ScalarOp) -> sl.ScalarOp:
        bounds_checkings = False
        if bounds_checkings:
            offset = sl.immediate_optionally(offset)
            return sl.If(offset >= self.out.length, sl.Halt(-1), self)
        else:
            return self

    def __getitem__(self, i):
        if isinstance(i, slice):
            if (i.step != 1) and (i.step is not None):
                raise NotImplementedError("steps other than 1 are not yet supported")
            return Slice(self, i.start, i.stop)
        else:
            return LoadElement(self, i)

    def new_value(self, parent: "VectorOp") -> "NewEmptyVector":
        return NewEmptyVector(parent)

    def copy_from(self, src: "VectorOp") -> "CopyByReference":
        return CopyByReference(self, src)


class NewEmptyVector(VectorOp):
    "Make a new vector with the same length, dtype, and address as `parent`."

    def __init__(self, parent: VectorOp):
        super().__init__(
            out=Vector(
                # Data args
                name="<unnamed vector for phi>",
                # Vector args
                length=parent.out.length,
                dtype=parent.out.dtype,
                address=sl.NewValue(parent.out.address),
            ),
            input_ops=(parent,),
        )


class CopyByReference(VectorOp):
    "Copy the address of one vector to the address of another."

    def __init__(self, dst: VectorOp, src: VectorOp):
        if dst.out.dtype != src.out.dtype:
            raise TypeError("dst and src must have the same dtype")
        if dst.out.length != src.out.length:
            raise ValueError("dst and src must have the same length")
        if isinstance(dst, Slice) or isinstance(src, Slice):
            raise NotImplementedError("cannot copy slices by reference yet")

        super().__init__(
            out=dst.out,
            input_ops=(dst, src, sl.CopyValue(dst.out.address, src.out.address)),
        )


class Allocate(VectorOp):
    "Create a vector with the given length and dtype and allocate storage for its data."

    def __init__(self, memory: sl.ScalarOp, length: int, dtype: DTypeLike):
        allocate_memory = sl.AllocateMemory(memory, length, dtype)
        vec = Vector("<unnamed>", length, dtype, allocate_memory)
        super().__init__(out=vec, input_ops=(allocate_memory,))

    def extra_args(self) -> str:
        return f"length={self.out.length}, dtype={self.out.dtype}"


class Constant(VectorOp):
    def __init__(self, memory: sl.ScalarOp, value: ArrayLike):
        value = np.asarray(value)
        if value.ndim != 1:
            raise ValueError("value must be a 1-dimensional array")

        address = sl.AllocateMemory(memory, len(value), value.dtype)
        copy = sl.CopyLiteralArray(address, value)
        super().__init__(
            out=ConstVector(
                # Data args
                name="<unnamed>",
                # Vector args
                length=len(value),
                dtype=value.dtype,
                # ConstVector args
                value=value,
                address=address,
            ),
            input_ops=(copy, memory),
        )

    def extra_args(self) -> str:
        return f"value={self.out.value}"


class Slice(VectorOp):
    def __init__(self, vec: VectorOp, start: sl.ScalarOp, end: sl.ScalarOp):
        start = sl.immediate_optionally(start)
        end = sl.immediate_optionally(end)
        super().__init__(out=vec.out, input_ops=(vec, start, end))

    def BoundsCheck(self, offset: sl.ScalarOp) -> sl.ScalarOp:
        _, start, _ = self.input_ops
        return super().BoundsCheck(start + offset)

    def AddressOf(self) -> sl.ScalarOp:
        _, start, _ = self.input_ops
        return super().AddressOf() + start


def Copy(
    dst: VectorOp,
    src: VectorOp,
    size: sl.ScalarOp,
    dst_offset: sl.ScalarOp = 0,
    src_offset: sl.ScalarOp = 0,
) -> VectorOp:
    if isinstance(dst.out, ConstVector):
        raise TypeError("dst must be mutable")
    if isinstance(dst, Slice) or isinstance(src, Slice):
        raise NotImplementedError("cannot copy slices yet")

    copy = sl.CopyMemory(
        dst.BoundsCheck(dst_offset + size),
        dst.AddressOf() + dst_offset,
        src.AddressOf() + src_offset,
        size,
    )
    # TODO: just return the copy here instead of creating a new VectorOp.
    return VectorOp(out=dst.out, input_ops=(copy, src))


def LoadElement(src: VectorOp, offset: sl.ScalarOp) -> sl.ScalarOp:
    return sl.Load(src.BoundsCheck(offset), src.AddressOf() + offset)


# TODO: This should be a VectorOp class. Otherwise it creates an unhelpful VectorOp.
def StoreElement(dst: VectorOp, offset: sl.ScalarOp, value: sl.ScalarOp) -> VectorOp:
    if isinstance(dst.out, ConstVector):
        raise TypeError("dst must be mutable")
    store = sl.Store(dst.BoundsCheck(offset), dst.AddressOf() + offset, value)
    return VectorOp(out=dst.out, input_ops=(store,))


def Clone(v: VectorOp) -> VectorOp:
    return Copy(dst=Allocate(v, v.out.length, v.out.dtype), src=v, size=v.out.length)


def lower(program: cg.Operator) -> List[opcodes.OpCode]:
    return sl.lower(program)