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 from nmigen.hdl.ast import Signal
  21
  22 class SimdScope:
  23     """The global scope object for SimdSignal and friends
  24
  25     Members:
  26     * vec_el_counts: SimdMap
  27         a map from `ElWid` values `k` to the number of elements in a vector
  28         when `self.elwid == k`.
  29
  30         Example:
  31         vec_el_counts = SimdMap({
  32             IntElWid.I64: 1,
  33             IntElWid.I32: 2,
  34             IntElWid.I16: 4,
  35             IntElWid.I8: 8,
  36         })
  37
  38         Another Example:
  39         vec_el_counts = SimdMap({
  40             FpElWid.F64: 1,
  41             FpElWid.F32: 2,
  42             FpElWid.F16: 4,
  43             FpElWid.BF16: 4,
  44         })
  45     * elwid: ElWid or nmigen Value with a shape of some ElWid class
  46         the current elwid (simd element type).  example: Signal(2)
  47         or Signal(IntElWid)
  48     """
  49
  50     __SCOPE_STACK = []
  51
  52     # XXX REMOVE THIS FUNCTION.  ITS USE IS DANGEROUS.
  53     @classmethod
  54     def get(cls):
  55         """get the current SimdScope. raises a ValueError outside of any
  56         SimdScope.
  57
  58         Example:
  59         with SimdScope(...) as s:
  60             assert SimdScope.get() is s
  61         """
  62         if len(cls.__SCOPE_STACK) > 0:
  63             retval = cls.__SCOPE_STACK[-1]
  64             assert isinstance(retval, SimdScope), "inconsistent scope stack"
  65             return retval
  66         raise ValueError("not in a `with SimdScope()` statement")
  67
  68     def __enter__(self):
  69         self.__SCOPE_STACK.append(self)
  70         return self
  71
  72     def __exit__(self, exc_type, exc_value, traceback):
  73         assert self.__SCOPE_STACK.pop() is self, "inconsistent scope stack"
  74         return False
  75
  76     def __init__(self, module, elwid, vec_el_counts, scalar=False):
  77
  78         # in SIMD mode, must establish module as part of context and inform
  79         # the module to operate under "SIMD" Type 1 (AST) casting rules,
  80         # not the # default "Value.cast" rules.
  81         if not scalar:
  82             self.module = module
  83             from ieee754.part.partsig import SimdSignal
  84             module._setAstTypeCastFn(SimdSignal.cast)
  85
  86         self.elwid = elwid
  87         self.vec_el_counts = vec_el_counts
  88         self.scalar = scalar
  89
  90     def __repr__(self):
  91         return (f"SimdScope(\n"
  92                 f"        elwid={self.elwid},\n"
  93                 f"        vec_el_counts={self.vec_el_counts},\n")
  94
  95     ##################
  96     # from here, the functions are context-aware variants of standard
  97     # nmigen API (Signal, Signal.like, Shape) which are to be redirected
  98     # to either their standard scalar nmigen equivalents (verbatim)
  99     # or to the SimdSignal equivalents.  each one is to be documented
 100     # CAREFULLY and CLEARLY.
 101     ##################
 102
 103     def Signal(self, shape=None, *, name=None, reset=0, reset_less=False,
 104                  attrs=None, decoder=None, src_loc_at=0,
 105                  fixed_width=None): # TODO: *REMOVE* THIS. work out how.
 106                                     # BE CAREFUL when using this param
 107                                     # it is NOT available in scalar mode
 108         if self.scalar:
 109             # scalar mode, just return a nmigen Signal.  THIS IS IMPORTANT.
 110             # when passing in SimdShape it should go "oh, this is
 111             # an isinstance Shape, i will just use its width and sign"
 112             # which is the entire reason why SimdShape had to derive
 113             # from Shape
 114             return Signal(shape=shape, name=name, reset=reset,
 115                           reset_less=reset_less, attrs=attrs,
 116                           decoder=decoder, src_loc_at=src_loc_at)
 117         else:
 118             # recursive module import resolution
 119             from ieee754.part.partsig import SimdSignal
 120             # SIMD mode.  shape here can be either a SimdShape,
 121             # a Shape, or anything else that Signal can take (int or
 122             # a tuple (int,bool) for (width,sign)
 123             s = SimdSignal(mask=self, # should contain *all* context needed,
 124                                       # which goes all the way through to
 125                                       # the layout() function, passing
 126                                       # 1) elwid 2) vec_el_counts
 127                           shape=shape, # should contain the *secondary*
 128                                        # part of the context needed for
 129                                        # the layout() function:
 130                                        # 3) lane_shapes 4) fixed_width
 131                           name=name, reset=reset,
 132                           reset_less=reset_less, attrs=attrs,
 133                           decoder=decoder, src_loc_at=src_loc_at,
 134                           fixed_width=fixed_width)
 135             # set the module context so that the SimdSignal can create
 136             # its own submodules during AST creation
 137             s.set_module(self.module)
 138             return s
 139
 140     # XXX TODO
 141     def Signal_like(self): pass
 142         #if self.scalar:
 143             # scalar mode, just return nmigen Signal.like.  THIS IS IMPORTANT.
 144         # else
 145             # simd mode.
 146
 147     # XXX TODO
 148     def Shape(self, width=1, signed=False):
 149         if self.scalar:
 150             # scalar mode, just return nmigen Shape.  THIS IS IMPORTANT.
 151             return Shape(width, signed)
 152         else:
 153             # SIMD mode. NOTE: for compatibility with Shape, the width
 154             # is assumed to be the widths_at_elwid parameter NOT the
 155             # fixed width.  this ensures that code that is converted
 156             # straight from scalar to SIMD will have the exact same
 157             # width at all elwidths, because layout() detects the integer
 158             # case and converts it, preserving the width at all elwidths
 159             # the names are preserved to ensure parameter-compatibility
 160             # with Shape()
 161             return SimdShape(self, width=width,   # actually widths_at_elwid
 162                                    signed=signed,
 163                                    fixed_width=None)