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