tile.py

from typing import List, Tuple
import dataclasses as dc
import numpy as np

import opcodes


@dc.dataclass
class Tile:
    """A portion of a tensor stored in SRAM."""

    shape: Tuple[int]
    dtype_in: np.dtype
    is_transposed: bool


@dc.dataclass
class TileOp:
    """Abstract base class for a tile operator"""

    def generate(self) -> List[opcodes.OpCode]:
        raise NotImplementedError


class Add(TileOp):
    tile_in_a: Tile
    tile_in_b: Tile
    tile_out: Tile


class MatMul(TileOp):
    tile_in_a: Tile
    tile_in_b: Tile
    tile_out: Tile

    def generate(self) -> List[opcodes.OpCode]:
        # Generate one dot product instruction for each of the 32x32 entries of
        # the result matrix.
        return [
            opcodes.DotProduct(
                opcodes.ScalarAdd(self.tile_in_a.address, i * 32),
                opcodes.ScalarAdd(self.tile_in_b.address, j * 32),
                opcodes.ScalarAdd(self.tile_out.address, i * 32 + j),
            )
            for i in range(32)
            for j in range(32)
        ]


class ExpSum(TileOp):
    tile_in: Tile
    tile_out_exps: Tile
    tile_out_accumulator: Tile


class FusedTileOperator(TileOp):
    tile_op1: TileOp
    tile_op2: TileOp

    def generate(self) -> List[opcodes.OpCode]:
        return self.tile_op1.generate() + self.tile_op2.generate()