src/ieee754/part/simd_scope.py

   1 # SPDX-License-Identifier: LGPL-3-or-later
   2 # See Notices.txt for copyright information
   3 """
   4 SimdScope class - provides context for SIMD signals to make them useable
   5 under the exact same API as scalar nmigen Signals.
   6
   7 Copyright (C) 2021 Jacob Lifshay
   8 Copyright (C) 2021 Luke Kenneth Casson Leighton
   9
  10 use as:
  11
  12     m = Module()
  13     with SimdScope(m, elwid) as s:
  14         a = s.Signal(width=64, ....)
  15
  16     m.d.comb += a.eq(...)
  17
  18 """
  19
  20
  21 from ieee754.part.util import (DEFAULT_FP_VEC_EL_COUNTS,
  22                                DEFAULT_INT_VEC_EL_COUNTS,
  23                                FpElWid, IntElWid, SimdMap)
  24 from nmigen.hdl.ast import Signal
  25
  26
  27 class SimdScope:
  28     """The global scope object for SimdSignal and friends
  29
  30     Members:
  31     * vec_el_counts: SimdMap
  32         a map from `ElWid` values `k` to the number of elements in a vector
  33         when `self.elwid == k`.
  34
  35         Example:
  36         vec_el_counts = SimdMap({
  37             IntElWid.I64: 1,
  38             IntElWid.I32: 2,
  39             IntElWid.I16: 4,
  40             IntElWid.I8: 8,
  41         })
  42
  43         Another Example:
  44         vec_el_counts = SimdMap({
  45             FpElWid.F64: 1,
  46             FpElWid.F32: 2,
  47             FpElWid.F16: 4,
  48             FpElWid.BF16: 4,
  49         })
  50     * elwid: ElWid or nmigen Value with a shape of some ElWid class
  51         the current elwid (simd element type)
  52     """
  53
  54     __SCOPE_STACK = []
  55
  56     @classmethod
  57     def get(cls):
  58         """get the current SimdScope. raises a ValueError outside of any
  59         SimdScope.
  60
  61         Example:
  62         with SimdScope(...) as s:
  63             assert SimdScope.get() is s
  64         """
  65         if len(cls.__SCOPE_STACK) > 0:
  66             retval = cls.__SCOPE_STACK[-1]
  67             assert isinstance(retval, SimdScope), "inconsistent scope stack"
  68             return retval
  69         raise ValueError("not in a `with SimdScope()` statement")
  70
  71     def __enter__(self):
  72         self.__SCOPE_STACK.append(self)
  73         return self
  74
  75     def __exit__(self, exc_type, exc_value, traceback):
  76         assert self.__SCOPE_STACK.pop() is self, "inconsistent scope stack"
  77         return False
  78
  79     def __init__(self, *, module, elwid=None,
  80                  vec_el_counts=None, elwid_type=IntElWid, scalar=False):
  81
  82         # in SIMD mode, must establish module as part of context and inform
  83         # the module to operate under "SIMD" Type 1 (AST) casting rules,
  84         # not the # default "Value.cast" rules.
  85         if not scalar:
  86             self.module = module
  87             from ieee754.part.partsig import SimdSignal
  88             module._setAstTypeCastFn(SimdSignal.cast)
  89
  90         # TODO, explain what this is about
  91         if isinstance(elwid, (IntElWid, FpElWid)):
  92             elwid_type = type(elwid)
  93             if vec_el_counts is None:
  94                 vec_el_counts = SimdMap({elwid: 1})
  95         assert issubclass(elwid_type, (IntElWid, FpElWid))
  96         self.elwid_type = elwid_type
  97         scalar_elwid = elwid_type(0)
  98
  99         # TODO, explain why this is needed.  Scalar should *NOT*
 100         # be doing anything other than *DIRECTLY* passing the
 101         # Signal() arguments *DIRECTLY* to nmigen.Signal.
 102         # UNDER NO CIRCUMSTANCES should ANY attempt be made to
 103         # treat SimdSignal as a "scalar Signal".  fuller explanation:
 104         # https://bugs.libre-soc.org/show_bug.cgi?id=734#c3
 105         if vec_el_counts is None:
 106             if scalar:
 107                 vec_el_counts = SimdMap({scalar_elwid: 1})
 108             elif issubclass(elwid_type, FpElWid):
 109                 vec_el_counts = DEFAULT_FP_VEC_EL_COUNTS
 110             else:
 111                 vec_el_counts = DEFAULT_INT_VEC_EL_COUNTS
 112
 113         # TODO, explain this function's purpose
 114         def check(elwid, vec_el_count):
 115             assert type(elwid) == elwid_type, "inconsistent ElWid types"
 116             vec_el_count = int(vec_el_count)
 117             assert vec_el_count != 0 \
 118                 and (vec_el_count & (vec_el_count - 1)) == 0,\
 119                 "vec_el_counts values must all be powers of two"
 120             return vec_el_count
 121
 122         # TODO, explain this
 123         self.vec_el_counts = SimdMap.map_with_elwid(check, vec_el_counts)
 124         self.full_el_count = max(self.vec_el_counts.values())
 125
 126         # TODO, explain this
 127         if elwid is not None:
 128             self.elwid = elwid
 129         elif scalar:
 130             self.elwid = scalar_elwid
 131         else:
 132             self.elwid = Signal(elwid_type)
 133
 134     def __repr__(self):
 135         return (f"SimdScope(\n"
 136                 f"        elwid={self.elwid},\n"
 137                 f"        elwid_type={self.elwid_type},\n"
 138                 f"        vec_el_counts={self.vec_el_counts},\n"
 139                 f"        full_el_count={self.full_el_count})")
 140
 141     ##################
 142     # from here, the functions are context-aware variants of standard
 143     # nmigen API (Signal, Signal.like, Shape) which are to be redirected
 144     # to either their standard scalar nmigen equivalents (verbatim)
 145     # or to the SimdSignal equivalents.  each one is to be documented
 146     # CAREFULLY and CLEARLY.
 147     ##################
 148
 149     def Signal(self, shape=None, *, name=None, reset=0, reset_less=False,
 150                  attrs=None, decoder=None, src_loc_at=0):
 151         if self.scalar:
 152             # scalar mode, just return a nmigen Signal.  THIS IS IMPORTANT.
 153             # when passing in SimdShape it should go "oh, this is
 154             # an isinstance Shape, i will just use its width and sign"
 155             # which is the entire reason why SimdShape had to derive
 156             # from Shape
 157             return Signal(shape=shape, name=name, reset=reset,
 158                           reset_less=reset_less, attrs=attrs,
 159                           decoder=decoder, src_loc_at=src_loc_at)
 160         else:
 161             # SIMD mode.  shape here can be either a SimdShape,
 162             # a Shape, or anything else that Signal can take (int or
 163             # a tuple (int,bool) for (width,sign)
 164             s = SimdSignal(mask=self, # should contain *all* context needed,
 165                                       # which goes all the way through to
 166                                       # the layout() function, passing
 167                                       # 1) elwid 2) vec_el_counts
 168                           shape=shape, # should contain the *secondary*
 169                                        # part of the context needed for
 170                                        # the layout() function:
 171                                        # 3) lane_shapes 4) fixed_width
 172                           name=name, reset=reset,
 173                           reset_less=reset_less, attrs=attrs,
 174                           decoder=decoder, src_loc_at=src_loc_at)
 175             # set the module context so that the SimdSignal can create
 176             # its own submodules during AST creation
 177             s.set_module(self.module)
 178
 179     # XXX TODO
 180     def Signal_like(self): pass
 181         #if self.scalar:
 182             # scalar mode, just return nmigen Signal.like.  THIS IS IMPORTANT.
 183         # else
 184             # simd mode.
 185
 186     # XXX TODO
 187     def Shape(self): pass
 188         #if self.scalar:
 189             # scalar mode, just return nmigen Shape.  THIS IS IMPORTANT.
 190         # else
 191             # simd mode.