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, vec_el_counts, scalar=False):
  82
  83         # in SIMD mode, must establish module as part of context and inform
  84         # the module to operate under "SIMD" Type 1 (AST) casting rules,
  85         # not the # default "Value.cast" rules.
  86         if not scalar:
  87             self.module = module
  88             from ieee754.part.partsig import SimdSignal
  89             module._setAstTypeCastFn(SimdSignal.cast)
  90
  91         self.elwid = elwid
  92         self.vec_el_counts = vec_el_counts
  93         self.scalar = scalar
  94
  95     def __repr__(self):
  96         return (f"SimdScope(\n"
  97                 f"        elwid={self.elwid},\n"
  98                 f"        vec_el_counts={self.vec_el_counts},\n")
  99
 100     ##################
 101     # from here, the functions are context-aware variants of standard
 102     # nmigen API (Signal, Signal.like, Shape) which are to be redirected
 103     # to either their standard scalar nmigen equivalents (verbatim)
 104     # or to the SimdSignal equivalents.  each one is to be documented
 105     # CAREFULLY and CLEARLY.
 106     ##################
 107
 108     def Signal(self, shape=None, *, name=None, reset=0, reset_less=False,
 109                  attrs=None, decoder=None, src_loc_at=0):
 110         if self.scalar:
 111             # scalar mode, just return a nmigen Signal.  THIS IS IMPORTANT.
 112             # when passing in SimdShape it should go "oh, this is
 113             # an isinstance Shape, i will just use its width and sign"
 114             # which is the entire reason why SimdShape had to derive
 115             # from Shape
 116             return Signal(shape=shape, name=name, reset=reset,
 117                           reset_less=reset_less, attrs=attrs,
 118                           decoder=decoder, src_loc_at=src_loc_at)
 119         else:
 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             # set the module context so that the SimdSignal can create
 135             # its own submodules during AST creation
 136             s.set_module(self.module)
 137
 138     # XXX TODO
 139     def Signal_like(self): pass
 140         #if self.scalar:
 141             # scalar mode, just return nmigen Signal.like.  THIS IS IMPORTANT.
 142         # else
 143             # simd mode.
 144
 145     # XXX TODO
 146     def Shape(self, width=1, signed=False):
 147         if self.scalar:
 148             # scalar mode, just return nmigen Shape.  THIS IS IMPORTANT.
 149             return Shape(width, signed)
 150         else:
 151             # SIMD mode. NOTE: for compatibility with Shape, the width
 152             # is assumed to be the widths_at_elwid parameter NOT the
 153             # fixed width.  this ensures that code that is converted
 154             # straight from scalar to SIMD will have the exact same
 155             # width at all elwidths, because layout() detects the integer
 156             # case and converts it, preserving the width at all elwidths
 157             # the names are preserved to ensure parameter-compatibility
 158             # with Shape()
 159             return SimdShape(self, width=width,   # actually widths_at_elwid
 160                                    signed=signed,
 161                                    fixed_width=None)