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).  example: Signal(2)
  52         or Signal(IntElWid)
  53     """
  54
  55     __SCOPE_STACK = []
  56
  57     # XXX REMOVE THIS FUNCTION.  ITS USE IS DANGEROUS.
  58     @classmethod
  59     def get(cls):
  60         """get the current SimdScope. raises a ValueError outside of any
  61         SimdScope.
  62
  63         Example:
  64         with SimdScope(...) as s:
  65             assert SimdScope.get() is s
  66         """
  67         if len(cls.__SCOPE_STACK) > 0:
  68             retval = cls.__SCOPE_STACK[-1]
  69             assert isinstance(retval, SimdScope), "inconsistent scope stack"
  70             return retval
  71         raise ValueError("not in a `with SimdScope()` statement")
  72
  73     def __enter__(self):
  74         self.__SCOPE_STACK.append(self)
  75         return self
  76
  77     def __exit__(self, exc_type, exc_value, traceback):
  78         assert self.__SCOPE_STACK.pop() is self, "inconsistent scope stack"
  79         return False
  80
  81     def __init__(self, *, module, elwid=None,
  82                  vec_el_counts=None, elwid_type=IntElWid, scalar=False):
  83
  84         # in SIMD mode, must establish module as part of context and inform
  85         # the module to operate under "SIMD" Type 1 (AST) casting rules,
  86         # not the # default "Value.cast" rules.
  87         if not scalar:
  88             self.module = module
  89             from ieee754.part.partsig import SimdSignal
  90             module._setAstTypeCastFn(SimdSignal.cast)
  91
  92         # TODO, explain what this is about
  93         if isinstance(elwid, (IntElWid, FpElWid)):
  94             elwid_type = type(elwid)
  95             if vec_el_counts is None:
  96                 vec_el_counts = SimdMap({elwid: 1})
  97         assert issubclass(elwid_type, (IntElWid, FpElWid))
  98         self.elwid_type = elwid_type
  99         scalar_elwid = elwid_type(0)
 100
 101         # TODO, explain why this is needed.  Scalar should *NOT*
 102         # be doing anything other than *DIRECTLY* passing the
 103         # Signal() arguments *DIRECTLY* to nmigen.Signal.
 104         # UNDER NO CIRCUMSTANCES should ANY attempt be made to
 105         # treat SimdSignal as a "scalar Signal".  fuller explanation:
 106         # https://bugs.libre-soc.org/show_bug.cgi?id=734#c3
 107         if vec_el_counts is None:
 108             if scalar:
 109                 vec_el_counts = SimdMap({scalar_elwid: 1})
 110             elif issubclass(elwid_type, FpElWid):
 111                 vec_el_counts = DEFAULT_FP_VEC_EL_COUNTS
 112             else:
 113                 vec_el_counts = DEFAULT_INT_VEC_EL_COUNTS
 114
 115         # TODO, explain this function's purpose
 116         def check(elwid, vec_el_count):
 117             assert type(elwid) == elwid_type, "inconsistent ElWid types"
 118             vec_el_count = int(vec_el_count)
 119             assert vec_el_count != 0 \
 120                 and (vec_el_count & (vec_el_count - 1)) == 0,\
 121                 "vec_el_counts values must all be powers of two"
 122             return vec_el_count
 123
 124         # TODO, explain this
 125         self.vec_el_counts = SimdMap.map_with_elwid(check, vec_el_counts)
 126         self.full_el_count = max(self.vec_el_counts.values())
 127
 128         # TODO, explain this
 129         if elwid is not None:
 130             self.elwid = elwid
 131         elif scalar:
 132             self.elwid = scalar_elwid
 133         else:
 134             self.elwid = Signal(elwid_type)
 135
 136     def __repr__(self):
 137         return (f"SimdScope(\n"
 138                 f"        elwid={self.elwid},\n"
 139                 f"        elwid_type={self.elwid_type},\n"
 140                 f"        vec_el_counts={self.vec_el_counts},\n"
 141                 f"        full_el_count={self.full_el_count})")
 142
 143     ##################
 144     # from here, the functions are context-aware variants of standard
 145     # nmigen API (Signal, Signal.like, Shape) which are to be redirected
 146     # to either their standard scalar nmigen equivalents (verbatim)
 147     # or to the SimdSignal equivalents.  each one is to be documented
 148     # CAREFULLY and CLEARLY.
 149     ##################
 150
 151     def Signal(self, shape=None, *, name=None, reset=0, reset_less=False,
 152                  attrs=None, decoder=None, src_loc_at=0):
 153         if self.scalar:
 154             # scalar mode, just return a nmigen Signal.  THIS IS IMPORTANT.
 155             # when passing in SimdShape it should go "oh, this is
 156             # an isinstance Shape, i will just use its width and sign"
 157             # which is the entire reason why SimdShape had to derive
 158             # from Shape
 159             return Signal(shape=shape, name=name, reset=reset,
 160                           reset_less=reset_less, attrs=attrs,
 161                           decoder=decoder, src_loc_at=src_loc_at)
 162         else:
 163             # SIMD mode.  shape here can be either a SimdShape,
 164             # a Shape, or anything else that Signal can take (int or
 165             # a tuple (int,bool) for (width,sign)
 166             s = SimdSignal(mask=self, # should contain *all* context needed,
 167                                       # which goes all the way through to
 168                                       # the layout() function, passing
 169                                       # 1) elwid 2) vec_el_counts
 170                           shape=shape, # should contain the *secondary*
 171                                        # part of the context needed for
 172                                        # the layout() function:
 173                                        # 3) lane_shapes 4) fixed_width
 174                           name=name, reset=reset,
 175                           reset_less=reset_less, attrs=attrs,
 176                           decoder=decoder, src_loc_at=src_loc_at)
 177             # set the module context so that the SimdSignal can create
 178             # its own submodules during AST creation
 179             s.set_module(self.module)
 180
 181     # XXX TODO
 182     def Signal_like(self): pass
 183         #if self.scalar:
 184             # scalar mode, just return nmigen Signal.like.  THIS IS IMPORTANT.
 185         # else
 186             # simd mode.
 187
 188     # XXX TODO
 189     def Shape(self, width=1, signed=False):
 190         if self.scalar:
 191             # scalar mode, just return nmigen Shape.  THIS IS IMPORTANT.
 192             return Shape(width, signed)
 193         else:
 194             # SIMD mode. NOTE: for compatibility with Shape, the width
 195             # is assumed to be the widths_at_elwid parameter NOT the
 196             # fixed width.  this ensures that code that is converted
 197             # straight from scalar to SIMD will have the exact same
 198             # width at all elwidths, because layout() detects the integer
 199             # case and converts it, preserving the width at all elwidths
 200             return SimdShape(self, width=None, signed=signed,
 201                                    widths_at_elwid=width)