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_VEC_EL_COUNTS,
   6                                DEFAULT_INT_VEC_EL_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     * vec_el_counts: SimdMap
  16         a map from `ElWid` values `k` to the number of elements in a vector
  17         when `self.elwid == k`.
  18
  19         Example:
  20         vec_el_counts = SimdMap({
  21             IntElWid.I64: 1,
  22             IntElWid.I32: 2,
  23             IntElWid.I16: 4,
  24             IntElWid.I8: 8,
  25         })
  26
  27         Another Example:
  28         vec_el_counts = SimdMap({
  29             FpElWid.F64: 1,
  30             FpElWid.F32: 2,
  31             FpElWid.F16: 4,
  32             FpElWid.BF16: 4,
  33         })
  34     * simd_full_width_hint: int
  35         the default value for SimdLayout's full_width argument, the full number
  36         of bits in a SIMD value.
  37     * elwid: ElWid or nmigen Value with a shape of some ElWid class
  38         the current elwid (simd element type)
  39     """
  40
  41     __SCOPE_STACK = []
  42
  43     @classmethod
  44     def get(cls):
  45         """get the current SimdScope. raises a ValueError outside of any
  46         SimdScope.
  47
  48         Example:
  49         with SimdScope(...) as s:
  50             assert SimdScope.get() is s
  51         """
  52         if len(cls.__SCOPE_STACK) > 0:
  53             retval = cls.__SCOPE_STACK[-1]
  54             assert isinstance(retval, SimdScope), "inconsistent scope stack"
  55             return retval
  56         raise ValueError("not in a `with SimdScope()` statement")
  57
  58     def __enter__(self):
  59         self.__SCOPE_STACK.append(self)
  60         return self
  61
  62     def __exit__(self, exc_type, exc_value, traceback):
  63         assert self.__SCOPE_STACK.pop() is self, "inconsistent scope stack"
  64         return False
  65
  66     def __init__(self, *, simd_full_width_hint=64, elwid=None,
  67                  vec_el_counts=None, elwid_type=IntElWid, scalar=False):
  68         assert isinstance(simd_full_width_hint, int)
  69         self.simd_full_width_hint = simd_full_width_hint
  70         if isinstance(elwid, (IntElWid, FpElWid)):
  71             elwid_type = type(elwid)
  72             if vec_el_counts is None:
  73                 vec_el_counts = SimdMap({elwid: 1})
  74         assert issubclass(elwid_type, (IntElWid, FpElWid))
  75         self.elwid_type = elwid_type
  76         scalar_elwid = elwid_type(0)
  77         if vec_el_counts is None:
  78             if scalar:
  79                 vec_el_counts = SimdMap({scalar_elwid: 1})
  80             elif issubclass(elwid_type, FpElWid):
  81                 vec_el_counts = DEFAULT_FP_VEC_EL_COUNTS
  82             else:
  83                 vec_el_counts = DEFAULT_INT_VEC_EL_COUNTS
  84
  85         def check(elwid, vec_el_count):
  86             assert type(elwid) == elwid_type, "inconsistent ElWid types"
  87             vec_el_count = int(vec_el_count)
  88             assert vec_el_count != 0 \
  89                 and (vec_el_count & (vec_el_count - 1)) == 0,\
  90                 "vec_el_counts values must all be powers of two"
  91             return vec_el_count
  92
  93         self.vec_el_counts = SimdMap.map_with_elwid(check, vec_el_counts)
  94         self.full_el_count = max(self.vec_el_counts.values())
  95         assert self.simd_full_width_hint % self.full_el_count == 0,\
  96             "simd_full_width_hint must be a multiple of full_el_count"
  97         if elwid is not None:
  98             self.elwid = elwid
  99         elif scalar:
 100             self.elwid = scalar_elwid
 101         else:
 102             self.elwid = Signal(elwid_type)
 103
 104     def __repr__(self):
 105         return (f"SimdScope(\n"
 106                 f"        simd_full_width_hint={self.simd_full_width_hint},\n"
 107                 f"        elwid={self.elwid},\n"
 108                 f"        elwid_type={self.elwid_type},\n"
 109                 f"        vec_el_counts={self.vec_el_counts},\n"
 110                 f"        full_el_count={self.full_el_count})")