move part_mask
[ieee754fpu.git] / src / ieee754 / part_mul_add / multiply.py
index 2cf37e84f1f3d6f322d4af48ede1e65e38f2d977..a67f4b4ccf7b8aed40f204a5d64159d17888e7e9 100644 (file)
@@ -309,13 +309,34 @@ class AddReduceData:
             for i in range(n_inputs)]
         self.reg_partition_points = ppoints.like()
 
-    def eq(self, rhs):
-        return [self.reg_partition_points.eq(rhs.reg_partition_points)] + \
-               [self.inputs[i].eq(rhs.inputs[i])
+    def eq_from(self, reg_partition_points, inputs, part_ops):
+        return [self.reg_partition_points.eq(reg_partition_points)] + \
+               [self.inputs[i].eq(inputs[i])
                                      for i in range(len(self.inputs))] + \
-               [self.part_ops[i].eq(rhs.part_ops[i])
+               [self.part_ops[i].eq(part_ops[i])
                                      for i in range(len(self.part_ops))]
 
+    def eq(self, rhs):
+        return self.eq_from(rhs.reg_partition_points, rhs.inputs, rhs.part_ops)
+
+
+class FinalReduceData:
+
+    def __init__(self, ppoints, output_width, n_parts):
+        self.part_ops = [Signal(2, name=f"part_ops_{i}")
+                          for i in range(n_parts)]
+        self.output = Signal(output_width)
+        self.reg_partition_points = ppoints.like()
+
+    def eq_from(self, reg_partition_points, output, part_ops):
+        return [self.reg_partition_points.eq(reg_partition_points)] + \
+               [self.output.eq(output)] + \
+               [self.part_ops[i].eq(part_ops[i])
+                                     for i in range(len(self.part_ops))]
+
+    def eq(self, rhs):
+        return self.eq_from(rhs.reg_partition_points, rhs.output, rhs.part_ops)
+
 
 class FinalAdd(Elaboratable):
     """ Final stage of add reduce
@@ -325,37 +346,40 @@ class FinalAdd(Elaboratable):
                        partition_points):
         self.i = AddReduceData(partition_points, n_inputs,
                                output_width, n_parts)
+        self.o = FinalReduceData(partition_points, output_width, n_parts)
+        self.output_width = output_width
         self.n_inputs = n_inputs
         self.n_parts = n_parts
-        self.out_part_ops = self.i.part_ops
-        self._resized_inputs = self.i.inputs
         self.register_levels = list(register_levels)
-        self.output = Signal(output_width)
         self.partition_points = PartitionPoints(partition_points)
         if not self.partition_points.fits_in_width(output_width):
             raise ValueError("partition_points doesn't fit in output_width")
-        self._reg_partition_points = self.i.reg_partition_points
-        self.intermediate_terms = []
 
     def elaborate(self, platform):
         """Elaborate this module."""
         m = Module()
 
+        output_width = self.output_width
+        output = Signal(output_width)
         if self.n_inputs == 0:
             # use 0 as the default output value
-            m.d.comb += self.output.eq(0)
+            m.d.comb += output.eq(0)
         elif self.n_inputs == 1:
             # handle single input
-            m.d.comb += self.output.eq(self._resized_inputs[0])
+            m.d.comb += output.eq(self.i.inputs[0])
         else:
             # base case for adding 2 inputs
             assert self.n_inputs == 2
-            adder = PartitionedAdder(len(self.output),
-                                     self._reg_partition_points)
+            adder = PartitionedAdder(output_width, self.i.reg_partition_points)
             m.submodules.final_adder = adder
-            m.d.comb += adder.a.eq(self._resized_inputs[0])
-            m.d.comb += adder.b.eq(self._resized_inputs[1])
-            m.d.comb += self.output.eq(adder.output)
+            m.d.comb += adder.a.eq(self.i.inputs[0])
+            m.d.comb += adder.b.eq(self.i.inputs[1])
+            m.d.comb += output.eq(adder.output)
+
+        # create output
+        m.d.comb += self.o.eq_from(self.i.reg_partition_points, output,
+                                   self.i.part_ops)
+
         return m
 
 
@@ -384,16 +408,12 @@ class AddReduceSingle(Elaboratable):
         self.n_inputs = n_inputs
         self.n_parts = n_parts
         self.output_width = output_width
-        self.out_part_ops = [Signal(2, name=f"out_part_ops_{i}")
-                          for i in range(n_parts)]
-        self._resized_inputs = [
-            Signal(output_width, name=f"resized_inputs[{i}]")
-            for i in range(n_inputs)]
+        self.i = AddReduceData(partition_points, n_inputs,
+                               output_width, n_parts)
         self.register_levels = list(register_levels)
         self.partition_points = PartitionPoints(partition_points)
         if not self.partition_points.fits_in_width(output_width):
             raise ValueError("partition_points doesn't fit in output_width")
-        self._reg_partition_points = self.partition_points.like()
 
         max_level = AddReduceSingle.get_max_level(n_inputs)
         for level in self.register_levels:
@@ -407,9 +427,13 @@ class AddReduceSingle(Elaboratable):
         # etc because this is not in elaboratable.
         self.groups = AddReduceSingle.full_adder_groups(n_inputs)
         self._intermediate_terms = []
+        self.adders = []
         if len(self.groups) != 0:
             self.create_next_terms()
 
+        self.o = AddReduceData(partition_points, len(self._intermediate_terms),
+                               output_width, n_parts)
+
     @staticmethod
     def get_max_level(input_count):
         """Get the maximum level.
@@ -437,42 +461,41 @@ class AddReduceSingle(Elaboratable):
         """Elaborate this module."""
         m = Module()
 
-        for (value, term) in self._intermediate_terms:
-            m.d.comb += term.eq(value)
+        # copy the intermediate terms to the output
+        for i, value in enumerate(self._intermediate_terms):
+            m.d.comb += self.o.inputs[i].eq(value)
+
+        # copy reg part points and part ops to output
+        m.d.comb += self.o.reg_partition_points.eq(self.i.reg_partition_points)
+        m.d.comb += [self.o.part_ops[i].eq(self.i.part_ops[i])
+                                     for i in range(len(self.i.part_ops))]
 
-        mask = self._reg_partition_points.as_mask(self.output_width)
-        m.d.comb += self.part_mask.eq(mask)
+        # set up the partition mask (for the adders)
+        part_mask = Signal(self.output_width, reset_less=True)
+
+        mask = self.i.reg_partition_points.as_mask(self.output_width)
+        m.d.comb += part_mask.eq(mask)
 
         # add and link the intermediate term modules
         for i, (iidx, adder_i) in enumerate(self.adders):
             setattr(m.submodules, f"adder_{i}", adder_i)
 
-            m.d.comb += adder_i.in0.eq(self._resized_inputs[iidx])
-            m.d.comb += adder_i.in1.eq(self._resized_inputs[iidx + 1])
-            m.d.comb += adder_i.in2.eq(self._resized_inputs[iidx + 2])
-            m.d.comb += adder_i.mask.eq(self.part_mask)
+            m.d.comb += adder_i.in0.eq(self.i.inputs[iidx])
+            m.d.comb += adder_i.in1.eq(self.i.inputs[iidx + 1])
+            m.d.comb += adder_i.in2.eq(self.i.inputs[iidx + 2])
+            m.d.comb += adder_i.mask.eq(part_mask)
 
         return m
 
     def create_next_terms(self):
 
-        # go on to prepare recursive case
-        intermediate_terms = []
         _intermediate_terms = []
 
         def add_intermediate_term(value):
-            intermediate_term = Signal(
-                self.output_width,
-                name=f"intermediate_terms[{len(intermediate_terms)}]")
-            _intermediate_terms.append((value, intermediate_term))
-            intermediate_terms.append(intermediate_term)
-
-        # store mask in intermediary (simplifies graph)
-        self.part_mask = Signal(self.output_width, reset_less=True)
+            _intermediate_terms.append(value)
 
         # create full adders for this recursive level.
         # this shrinks N terms to 2 * (N // 3) plus the remainder
-        self.adders = []
         for i in self.groups:
             adder_i = MaskedFullAdder(self.output_width)
             self.adders.append((i, adder_i))
@@ -482,17 +505,16 @@ class AddReduceSingle(Elaboratable):
             add_intermediate_term(adder_i.mcarry)
         # handle the remaining inputs.
         if self.n_inputs % FULL_ADDER_INPUT_COUNT == 1:
-            add_intermediate_term(self._resized_inputs[-1])
+            add_intermediate_term(self.i.inputs[-1])
         elif self.n_inputs % FULL_ADDER_INPUT_COUNT == 2:
             # Just pass the terms to the next layer, since we wouldn't gain
             # anything by using a half adder since there would still be 2 terms
             # and just passing the terms to the next layer saves gates.
-            add_intermediate_term(self._resized_inputs[-2])
-            add_intermediate_term(self._resized_inputs[-1])
+            add_intermediate_term(self.i.inputs[-2])
+            add_intermediate_term(self.i.inputs[-1])
         else:
             assert self.n_inputs % FULL_ADDER_INPUT_COUNT == 0
 
-        self.intermediate_terms = intermediate_terms
         self._intermediate_terms = _intermediate_terms
 
 
@@ -520,9 +542,8 @@ class AddReduce(Elaboratable):
         """
         self.inputs = inputs
         self.part_ops = part_ops
-        self.out_part_ops = [Signal(2, name=f"out_part_ops_{i}")
-                          for i in range(len(part_ops))]
-        self.output = Signal(output_width)
+        n_parts = len(part_ops)
+        self.o = FinalReduceData(partition_points, output_width, n_parts)
         self.output_width = output_width
         self.register_levels = register_levels
         self.partition_points = partition_points
@@ -546,27 +567,26 @@ class AddReduce(Elaboratable):
         mods = []
         next_levels = self.register_levels
         partition_points = self.partition_points
-        inputs = self.inputs
         part_ops = self.part_ops
         n_parts = len(part_ops)
+        inputs = self.inputs
+        ilen = len(inputs)
         while True:
-            ilen = len(inputs)
             next_level = AddReduceSingle(ilen, self.output_width, n_parts,
                                          next_levels, partition_points)
             mods.append(next_level)
             next_levels = list(AddReduce.next_register_levels(next_levels))
-            partition_points = next_level._reg_partition_points
-            inputs = next_level.intermediate_terms
+            partition_points = next_level.i.reg_partition_points
+            inputs = next_level.o.inputs
             ilen = len(inputs)
-            part_ops = next_level.out_part_ops
+            part_ops = next_level.i.part_ops
             groups = AddReduceSingle.full_adder_groups(len(inputs))
             if len(groups) == 0:
                 break
 
-        if ilen != 0:
-            next_level = FinalAdd(ilen, self.output_width, n_parts,
-                                  next_levels, partition_points)
-            mods.append(next_level)
+        next_level = FinalAdd(ilen, self.output_width, n_parts,
+                              next_levels, partition_points)
+        mods.append(next_level)
 
         self.levels = mods
 
@@ -580,29 +600,21 @@ class AddReduce(Elaboratable):
         partition_points = self.partition_points
         inputs = self.inputs
         part_ops = self.part_ops
-        for i in range(len(self.levels)):
-            mcur = self.levels[i]
-            inassign = [mcur._resized_inputs[i].eq(inputs[i])
-                                         for i in range(len(inputs))]
-            copy_part_ops = [mcur.out_part_ops[i].eq(part_ops[i])
-                                         for i in range(len(part_ops))]
+        n_parts = len(part_ops)
+        n_inputs = len(inputs)
+        output_width = self.output_width
+        i = AddReduceData(partition_points, n_inputs, output_width, n_parts)
+        m.d.comb += i.eq_from(partition_points, inputs, part_ops)
+        for idx in range(len(self.levels)):
+            mcur = self.levels[idx]
             if 0 in mcur.register_levels:
-                m.d.sync += copy_part_ops
-                m.d.sync += inassign
-                m.d.sync += mcur._reg_partition_points.eq(partition_points)
+                m.d.sync += mcur.i.eq(i)
             else:
-                m.d.comb += copy_part_ops
-                m.d.comb += inassign
-                m.d.comb += mcur._reg_partition_points.eq(partition_points)
-            partition_points = mcur._reg_partition_points
-            inputs = mcur.intermediate_terms
-            part_ops = mcur.out_part_ops
+                m.d.comb += mcur.i.eq(i)
+            i = mcur.o # for next loop
 
         # output comes from last module
-        m.d.comb += self.output.eq(next_level.output)
-        copy_part_ops = [self.out_part_ops[i].eq(next_level.out_part_ops[i])
-                                     for i in range(len(self.part_ops))]
-        m.d.comb += copy_part_ops
+        m.d.comb += self.o.eq(i)
 
         return m
 
@@ -1129,11 +1141,11 @@ class Mul8_16_32_64(Elaboratable):
                                expanded_part_pts,
                                self.part_ops)
 
-        out_part_ops = add_reduce.levels[-1].out_part_ops
-        out_part_pts = add_reduce.levels[-1]._reg_partition_points
+        out_part_ops = add_reduce.o.part_ops
+        out_part_pts = add_reduce.o.reg_partition_points
 
         m.submodules.add_reduce = add_reduce
-        m.d.comb += self._intermediate_output.eq(add_reduce.output)
+        m.d.comb += self._intermediate_output.eq(add_reduce.o.output)
         # create _output_64
         m.submodules.io64 = io64 = IntermediateOut(64, 128, 1)
         m.d.comb += io64.intermed.eq(self._intermediate_output)