types.py

Go to the documentation of this file.

# ===-----------------------------------------------------------------------===#
#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
#  See https://llvm.org/LICENSE.txt for license information.
#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# ===-----------------------------------------------------------------------===#
#
# The structure of the Python classes and hierarchy roughly mirrors the C++
# side, but wraps the C++ objects. The wrapper classes sometimes add convenience
# functionality and serve to return wrapped versions of the returned objects.
#
# ===-----------------------------------------------------------------------===#
 
from __future__ import annotations
 
from . import esiCppAccel as cpp
 
from typing import TYPE_CHECKING
if TYPE_CHECKING:
  from .accelerator import HWModule
 
from concurrent.futures import Future
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 
def _get_esi_type(cpp_type: cpp.Type):
  """Get the wrapper class for a C++ type."""
  for cpp_type_cls, fn in __esi_mapping.items():
    if isinstance(cpp_type, cpp_type_cls):
      return fn(cpp_type)
  return ESIType(cpp_type)
 
 
# Mapping from C++ types to functions constructing the Python object
# corresponding to that type.
__esi_mapping: Dict[Type, Callable] = {
    cpp.ChannelType: lambda cpp_type: _get_esi_type(cpp_type.inner)
}
 
 
class ESIType:
 
  def __init__(self, cpp_type: cpp.Type):
    self.cpp_typecpp_type = cpp_type
 
  @property
  def supports_host(self) -> Tuple[bool, Optional[str]]:
    """Does this type support host communication via Python? Returns either
    '(True, None)' if it is, or '(False, reason)' if it is not."""
 
    if self.bit_widthbit_width % 8 != 0:
      return (False, "runtime only supports types with multiple of 8 bits")
    return (True, None)
 
  def is_valid(self, obj) -> Tuple[bool, Optional[str]]:
    """Is a Python object compatible with HW type?  Returns either '(True,
    None)' if it is, or '(False, reason)' if it is not."""
    assert False, "unimplemented"
 
  @property
  def bit_width(self) -> int:
    """Size of this type, in bits. Negative for unbounded types."""
    assert False, "unimplemented"
 
  @property
  def max_size(self) -> int:
    """Maximum size of a value of this type, in bytes."""
    bitwidth = int((self.bit_widthbit_width + 7) / 8)
    if bitwidth < 0:
      return bitwidth
    return bitwidth
 
  def serialize(self, obj) -> bytearray:
    """Convert a Python object to a bytearray."""
    assert False, "unimplemented"
 
  def deserialize(self, data: bytearray) -> Tuple[object, bytearray]:
    """Convert a bytearray to a Python object. Return the object and the
    leftover bytes."""
    assert False, "unimplemented"
 
  def __str__(self) -> str:
    return str(self.cpp_typecpp_type)
 
 
class VoidType(ESIType):
 
  def is_valid(self, obj) -> Tuple[bool, Optional[str]]:
    if obj is not None:
      return (False, f"void type cannot must represented by None, not {obj}")
    return (True, None)
 
  @property
  def bit_width(self) -> int:
    return 8
 
  def serialize(self, obj) -> bytearray:
    # By convention, void is represented by a single byte of value 0.
    return bytearray([0])
 
  def deserialize(self, data: bytearray) -> Tuple[object, bytearray]:
    if len(data) == 0:
      raise ValueError(f"void type cannot be represented by {data}")
    return (None, data[1:])
 
 
__esi_mapping[cpp.VoidType] = VoidType
 
 
class BitsType(ESIType):
 
  def __init__(self, cpp_type: cpp.BitsType):
    self.cpp_typecpp_type: cpp.BitsType = cpp_type
 
  def is_valid(self, obj) -> Tuple[bool, Optional[str]]:
    if not isinstance(obj, (bytearray, bytes, list)):
      return (False, f"invalid type: {type(obj)}")
    if isinstance(obj, list) and not all(
        [isinstance(b, int) and b.bit_length() <= 8 for b in obj]):
      return (False, f"list item too large: {obj}")
    if len(obj) != self.max_sizemax_size:
      return (False, f"wrong size: {len(obj)}")
    return (True, None)
 
  @property
  def bit_width(self) -> int:
    return self.cpp_typecpp_type.width
 
  def serialize(self, obj: Union[bytearray, bytes, List[int]]) -> bytearray:
    if isinstance(obj, bytearray):
      return obj
    if isinstance(obj, bytes) or isinstance(obj, list):
      return bytearray(obj)
    raise ValueError(f"cannot convert {obj} to bytearray")
 
  def deserialize(self, data: bytearray) -> Tuple[bytearray, bytearray]:
    return (data[0:self.max_sizemax_size], data[self.max_sizemax_size:])
 
 
__esi_mapping[cpp.BitsType] = BitsType
 
 
class IntType(ESIType):
 
  def __init__(self, cpp_type: cpp.IntegerType):
    self.cpp_typecpp_type: cpp.IntegerType = cpp_type
 
  @property
  def bit_width(self) -> int:
    return self.cpp_typecpp_type.width
 
 
class UIntType(IntType):
 
  def is_valid(self, obj) -> Tuple[bool, Optional[str]]:
    if not isinstance(obj, int):
      return (False, f"must be an int, not {type(obj)}")
    if obj < 0 or obj.bit_length() > self.bit_widthbit_widthbit_width:
      return (False, f"out of range: {obj}")
    return (True, None)
 
  def __str__(self) -> str:
    return f"uint{self.bit_width}"
 
  def serialize(self, obj: int) -> bytearray:
    return bytearray(int.to_bytes(obj, self.max_sizemax_size, "little"))
 
  def deserialize(self, data: bytearray) -> Tuple[int, bytearray]:
    return (int.from_bytes(data[0:self.max_sizemax_size],
                           "little"), data[self.max_sizemax_size:])
 
 
__esi_mapping[cpp.UIntType] = UIntType
 
 
class SIntType(IntType):
 
  def is_valid(self, obj) -> Tuple[bool, Optional[str]]:
    if not isinstance(obj, int):
      return (False, f"must be an int, not {type(obj)}")
    if obj < 0:
      if (-1 * obj) > 2**(self.bit_widthbit_widthbit_width - 1):
        return (False, f"out of range: {obj}")
    elif obj < 0:
      if obj >= 2**(self.bit_widthbit_widthbit_width - 1) - 1:
        return (False, f"out of range: {obj}")
    return (True, None)
 
  def __str__(self) -> str:
    return f"sint{self.bit_width}"
 
  def serialize(self, obj: int) -> bytearray:
    return bytearray(int.to_bytes(obj, self.max_sizemax_size, "little", signed=True))
 
  def deserialize(self, data: bytearray) -> Tuple[int, bytearray]:
    return (int.from_bytes(data[0:self.max_sizemax_size], "little",
                           signed=True), data[self.max_sizemax_size:])
 
 
__esi_mapping[cpp.SIntType] = SIntType
 
 
class StructType(ESIType):
 
  def __init__(self, cpp_type: cpp.StructType):
    self.cpp_typecpp_typecpp_type = cpp_type
    self.fields: List[Tuple[str, ESIType]] = [
        (name, _get_esi_type(ty)) for (name, ty) in cpp_type.fields
    ]
 
  @property
  def bit_width(self) -> int:
    widths = [ty.bit_width for (_, ty) in self.fields]
    if any([w < 0 for w in widths]):
      return -1
    return sum(widths)
 
  def is_valid(self, obj) -> Tuple[bool, Optional[str]]:
    fields_count = 0
    if not isinstance(obj, dict):
      obj = obj.__dict__
 
    for (fname, ftype) in self.fields:
      if fname not in obj:
        return (False, f"missing field '{fname}'")
      fvalid, reason = ftype.is_valid(obj[fname])
      if not fvalid:
        return (False, f"invalid field '{fname}': {reason}")
      fields_count += 1
    if fields_count != len(obj):
      return (False, "missing fields")
    return (True, None)
 
  def serialize(self, obj) -> bytearray:
    ret = bytearray()
    for (fname, ftype) in reversed(self.fields):
      fval = obj[fname]
      ret.extend(ftype.serialize(fval))
    return ret
 
  def deserialize(self, data: bytearray) -> Tuple[Dict[str, Any], bytearray]:
    ret = {}
    for (fname, ftype) in reversed(self.fields):
      (fval, data) = ftype.deserialize(data)
      ret[fname] = fval
    return (ret, data)
 
 
__esi_mapping[cpp.StructType] = StructType
 
 
class ArrayType(ESIType):
 
  def __init__(self, cpp_type: cpp.ArrayType):
    self.cpp_typecpp_typecpp_type = cpp_type
    self.element_typeelement_type = _get_esi_type(cpp_type.element)
    self.sizesize = cpp_type.size
 
  @property
  def bit_width(self) -> int:
    return self.element_typeelement_type.bit_width * self.sizesize
 
  def is_valid(self, obj) -> Tuple[bool, Optional[str]]:
    if not isinstance(obj, list):
      return (False, f"must be a list, not {type(obj)}")
    if len(obj) != self.sizesize:
      return (False, f"wrong size: expected {self.size} not {len(obj)}")
    for (idx, e) in enumerate(obj):
      evalid, reason = self.element_typeelement_type.is_valid(e)
      if not evalid:
        return (False, f"invalid element {idx}: {reason}")
    return (True, None)
 
  def serialize(self, lst: list) -> bytearray:
    ret = bytearray()
    for e in reversed(lst):
      ret.extend(self.element_typeelement_type.serialize(e))
    return ret
 
  def deserialize(self, data: bytearray) -> Tuple[List[Any], bytearray]:
    ret = []
    for _ in range(self.sizesize):
      (obj, data) = self.element_typeelement_type.deserialize(data)
      ret.append(obj)
    ret.reverse()
    return (ret, data)
 
 
__esi_mapping[cpp.ArrayType] = ArrayType
 
 
class Port:
  """A unidirectional communication channel. This is the basic communication
  method with an accelerator."""
 
  def __init__(self, owner: BundlePort, cpp_port: cpp.ChannelPort):
    self.ownerowner = owner
    self.cpp_portcpp_port = cpp_port
    self.typetype = _get_esi_type(cpp_port.type)
    (supports_host, reason) = self.typetype.supports_host
    if not supports_host:
      raise TypeError(f"unsupported type: {reason}")
 
  def connect(self):
    self.cpp_portcpp_port.connect()
    return self
 
 
class WritePort(Port):
  """A unidirectional communication channel from the host to the accelerator."""
 
  def __init__(self, owner: BundlePort, cpp_port: cpp.WriteChannelPort):
    super().__init__(owner, cpp_port)
    self.cpp_portcpp_port: cpp.WriteChannelPort = cpp_port
 
  def write(self, msg=None) -> bool:
    """Write a typed message to the channel. Attempts to serialize 'msg' to what
    the accelerator expects, but will fail if the object is not convertible to
    the port type."""
 
    valid, reason = self.typetype.is_valid(msg)
    if not valid:
      raise ValueError(
          f"'{msg}' cannot be converted to '{self.type}': {reason}")
    msg_bytes: bytearray = self.typetype.serialize(msg)
    self.cpp_portcpp_port.write(msg_bytes)
    return True
 
 
class ReadPort(Port):
  """A unidirectional communication channel from the accelerator to the host."""
 
  def __init__(self, owner: BundlePort, cpp_port: cpp.ReadChannelPort):
    super().__init__(owner, cpp_port)
    self.cpp_portcpp_port: cpp.ReadChannelPort = cpp_port
 
  def read(self) -> Tuple[bool, Optional[object]]:
    """Read a typed message from the channel. Returns a deserialized object of a
    type defined by the port type."""
 
    buffer = self.cpp_portcpp_port.read()
    if buffer is None:
      return (False, None)
    (msg, leftover) = self.typetype.deserialize(buffer)
    if len(leftover) != 0:
      raise ValueError(f"leftover bytes: {leftover}")
    return (True, msg)
 
 
class BundlePort:
  """A collections of named, unidirectional communication channels."""
 
  # When creating a new port, we need to determine if it is a service port and
  # instantiate it correctly.
  def __new__(cls, owner: HWModule, cpp_port: cpp.BundlePort):
    # TODO: add a proper registration mechanism for service ports.
    if isinstance(cpp_port, cpp.Function):
      return super().__new__(FunctionPort)
    return super().__new__(cls)
 
  def __init__(self, owner: HWModule, cpp_port: cpp.BundlePort):
    self.ownerowner = owner
    self.cpp_portcpp_port = cpp_port
 
  def write_port(self, channel_name: str) -> WritePort:
    return WritePort(self, self.cpp_portcpp_port.getWrite(channel_name))
 
  def read_port(self, channel_name: str) -> ReadPort:
    return ReadPort(self, self.cpp_portcpp_port.getRead(channel_name))
 
 
class MessageFuture(Future):
  """A specialization of `Future` for ESI messages. Wraps the cpp object and
  deserializes the result.  Hopefully overrides all the methods necessary for
  proper operation, which is assumed to be not all of them."""
 
  def __init__(self, result_type: Type, cpp_future: cpp.MessageDataFuture):
    self.result_typeresult_type = result_type
    self.cpp_futurecpp_future = cpp_future
 
  def running(self) -> bool:
    return True
 
  def done(self) -> bool:
    return self.cpp_futurecpp_future.valid()
 
  def result(self, timeout: Optional[Union[int, float]] = None) -> Any:
    # TODO: respect timeout
    self.cpp_futurecpp_future.wait()
    result_bytes = self.cpp_futurecpp_future.get()
    (msg, leftover) = self.result_typeresult_type.deserialize(result_bytes)
    if len(leftover) != 0:
      raise ValueError(f"leftover bytes: {leftover}")
    return msg
 
  def add_done_callback(self, fn: Callable[[Future], object]) -> None:
    raise NotImplementedError("add_done_callback is not implemented")
 
 
class FunctionPort(BundlePort):
  """A pair of channels which carry the input and output of a function."""
 
  def __init__(self, owner: HWModule, cpp_port: cpp.BundlePort):
    super().__init__(owner, cpp_port)
    self.arg_typearg_type = self.write_portwrite_port("arg").type
    self.result_typeresult_type = self.read_portread_port("result").type
    self.connectedconnected = False
 
  def connect(self):
    self.cpp_portcpp_port.connect()
    self.connectedconnected = True
 
  def call(self, **kwargs: Any) -> Future:
    """Call the function with the given argument and returns a future of the
    result."""
    valid, reason = self.arg_typearg_type.is_valid(kwargs)
    if not valid:
      raise ValueError(
          f"'{kwargs}' cannot be converted to '{self.arg_type}': {reason}")
    arg_bytes: bytearray = self.arg_typearg_type.serialize(kwargs)
    cpp_future = self.cpp_portcpp_port.call(arg_bytes)
    return MessageFuture(self.result_typeresult_type, cpp_future)
 
  def __call__(self, *args: Any, **kwds: Any) -> Future:
    return self.callcall(*args, **kwds)