Class Field

tinygrad/runtime/support/c.py:56–81

  return cls

class Field:
  def __init__(self, typ, off, bit_width=None, bit_off=0, *, name=None, idx=0):
    self.typ, self.off, self.bit_width, self.bit_off, self.name, self.idx = typ, off, bit_width, bit_off, name, idx

  def __set_name__(self, owner, name):
    entry = (name, self.typ, self.off) + ((self.bit_width, self.bit_off) if self.bit_width else ())
    if hasattr(owner, "_real_fields_"): owner._real_fields_.append(entry)
    else: setattr(owner, "_real_fields_", [entry])
    self.name, self.idx = name, len(owner._real_fields_) - 1

  # lazily resolve field descriptors
  def _resolve(self, cls):
    if self.bit_width: # handle bitfields ourselves
      sl, set_mask = slice(self.off, self.off+(sz:=ceildiv(self.bit_width+self.bit_off, 8))), ~((mask:=(1 << self.bit_width) - 1) << self.bit_off)
      def b2i(obj): return int.from_bytes(memoryview(obj).cast("B")[sl], sys.byteorder)
      def bset(obj, v): memoryview(obj).cast("B")[sl] = ((b2i(obj) & set_mask) | v << self.bit_off).to_bytes(sz, sys.byteorder)
      # FIXME: signedness
      cf = property(lambda obj: b2i(obj) >> self.bit_off & mask, bset)
    # pull the CField descriptor from a dummy class, zero length arrays are so ctypes manages references to child objects for us
    else: cf = type(self.name, (ctypes.Structure,), {"_layout_": "ms", "_pack_": 1, "_fields_": [(str(i), ctypes.c_byte*0) for i in range(self.idx)] +
                                                                                                [("_", ctypes.c_byte * self.off), ("v", self.typ)]}).v # type: ignore
    setattr(cls, self.name, cf)
    return cf

  def __get__(self, obj, objtype=None): return self._resolve(objtype).__get__(obj, objtype) if objtype else self
  def __set__(self, obj, value): self._resolve(obj.__class__).__set__(obj, value)

@functools.cache
def init_c_struct_t(sz:int, fields: tuple[tuple, ...]):