src/ieee754/part/simd_scope.py

   1 # SPDX-License-Identifier: LGPL-3-or-later
   2 # See Notices.txt for copyright information
   3
   4
   5 from ieee754.part.util import (DEFAULT_FP_PART_COUNTS,
   6                                DEFAULT_INT_PART_COUNTS,
   7                                FpElWid, IntElWid, SimdMap)
   8 from nmigen.hdl.ast import Signal
   9
  10
  11 class SimdScope:
  12     """The global scope object for SimdSignal and friends
  13
  14     Members:
  15     * part_counts: SimdMap
  16         a map from `ElWid` values `k` to the number of parts in an element
  17         when `self.elwid == k`. Values should be minimized, since higher values
  18         often create bigger circuits.
  19
  20         Example:
  21         # here, an I8 element is 1 part wide
  22         part_counts = {ElWid.I8: 1, ElWid.I16: 2, ElWid.I32: 4, ElWid.I64: 8}
  23
  24         Another Example:
  25         # here, an F16 element is 1 part wide
  26         part_counts = {ElWid.F16: 1, ElWid.BF16: 1, ElWid.F32: 2, ElWid.F64: 4}
  27     * simd_full_width_hint: int
  28         the default value for SimdLayout's full_width argument, the full number
  29         of bits in a SIMD value.
  30     * elwid: ElWid or nmigen Value with a shape of some ElWid class
  31         the current elwid (simd element type)
  32     """
  33
  34     __SCOPE_STACK = []
  35
  36     @classmethod
  37     def get(cls):
  38         """get the current SimdScope.
  39
  40         Example:
  41         SimdScope.get(None) is None
  42         SimdScope.get() raises ValueError
  43         with SimdScope(...) as s:
  44             SimdScope.get() is s
  45         """
  46         if len(cls.__SCOPE_STACK) > 0:
  47             retval = cls.__SCOPE_STACK[-1]
  48             assert isinstance(retval, SimdScope), "inconsistent scope stack"
  49             return retval
  50         raise ValueError("not in a `with SimdScope()` statement")
  51
  52     def __enter__(self):
  53         self.__SCOPE_STACK.append(self)
  54         return self
  55
  56     def __exit__(self, exc_type, exc_value, traceback):
  57         assert self.__SCOPE_STACK.pop() is self, "inconsistent scope stack"
  58         return False
  59
  60     def __init__(self, *, simd_full_width_hint=64, elwid=None,
  61                  part_counts=None, elwid_type=IntElWid, scalar=False):
  62         # TODO: add more arguments/members and processing for integration with
  63         self.simd_full_width_hint = simd_full_width_hint
  64         if isinstance(elwid, (IntElWid, FpElWid)):
  65             elwid_type = type(elwid)
  66             if part_counts is None:
  67                 part_counts = SimdMap({elwid: 1})
  68         assert issubclass(elwid_type, (IntElWid, FpElWid))
  69         self.elwid_type = elwid_type
  70         scalar_elwid = elwid_type(0)
  71         if part_counts is None:
  72             if scalar:
  73                 part_counts = SimdMap({scalar_elwid: 1})
  74             elif issubclass(elwid_type, FpElWid):
  75                 part_counts = DEFAULT_FP_PART_COUNTS
  76             else:
  77                 part_counts = DEFAULT_INT_PART_COUNTS
  78
  79         def check(elwid, part_count):
  80             assert type(elwid) == elwid_type, "inconsistent ElWid types"
  81             part_count = int(part_count)
  82             assert part_count != 0 and (part_count & (part_count - 1)) == 0,\
  83                 "part_counts values must all be powers of two"
  84             return part_count
  85
  86         self.part_counts = SimdMap.map_with_elwid(check, part_counts)
  87         self.full_part_count = max(part_counts.values())
  88         assert self.simd_full_width_hint % self.full_part_count == 0,\
  89             "simd_full_width_hint must be a multiple of full_part_count"
  90         if elwid is not None:
  91             self.elwid = elwid
  92         elif scalar:
  93             self.elwid = scalar_elwid
  94         else:
  95             self.elwid = Signal(elwid_type)
  96
  97     def __repr__(self):
  98         return (f"SimdScope(\n"
  99                 f"        simd_full_width_hint={self.simd_full_width_hint},\n"
 100                 f"        elwid={self.elwid},\n"
 101                 f"        elwid_type={self.elwid_type},\n"
 102                 f"        part_counts={self.part_counts},\n"
 103                 f"        full_part_count={self.full_part_count})")