ReduceArg ops returns the arg index and its corresponding value.

Author: shtaxxx

tests/extension/thread_/stream_reduce_arg_max/thread_stream_reduce_arg_max.py

@@ -26,24 +26,24 @@ def mkLed():
     strm = vthread.Stream(m, 'mystream', clk, rst)
     a = strm.source('a')
     size = strm.constant('size')
-    max, max_valid = strm.ReduceArgMaxValid(a, size)
-    strm.sink(max, 'max', when=max_valid, when_name='max_valid')
+    index, _max, argmax_valid = strm.ReduceArgMaxValid(a, size)
+    strm.sink(index, 'index', when=argmax_valid, when_name='argmax_valid')
 
     def comp_stream(size, offset):
         strm.set_source('a', ram_a, offset, size)
         strm.set_constant('size', size)
-        strm.set_sink('max', ram_b, offset, 1)
+        strm.set_sink('index', ram_b, offset, 1)
         strm.run()
         strm.join()
 
     def comp_sequential(size, offset):
         index = 0
-        max = 0
+        _max = 0
         for i in range(size):
             a = ram_a.read(i + offset)
-            if a > max:
+            if a > _max:
                 index = i
-                max = a
+                _max = a
         ram_b.write(offset, index)
 
     def check(size, offset_stream, offset_seq):

veriloggen/stream/stypes.py

@@ -2610,6 +2610,7 @@ def __init__(self, right, size=None, initval=0,
                  enable=None, reset=None, width=32, signed=True):
         _Accumulator.__init__(self, right, size, initval,
                               enable, reset, width, signed)
+        self.graph_label = 'ReduceAdd'
 
 
 class ReduceSub(_Accumulator):
@@ -2619,6 +2620,7 @@ def __init__(self, right, size=None, initval=0,
                  enable=None, reset=None, width=32, signed=True):
         _Accumulator.__init__(self, right, size, initval,
                               enable, reset, width, signed)
+        self.graph_label = 'ReduceSub'
 
 
 class ReduceMul(_Accumulator):
@@ -2629,6 +2631,7 @@ def __init__(self, right, size=None, initval=0,
                  enable=None, reset=None, width=32, signed=True):
         _Accumulator.__init__(self, right, size, initval,
                               enable, reset, width, signed)
+        self.graph_label = 'ReduceMul'
 
 
 class ReduceDiv(_Accumulator):
@@ -2640,6 +2643,7 @@ def __init__(self, right, size=None, initval=0,
         raise NotImplementedError()
         _Accumulator.__init__(self, right, size, initval,
                               enable, reset, width, signed)
+        self.graph_label = 'ReduceDiv'
 
 
 class ReduceMax(_Accumulator):
@@ -2649,6 +2653,7 @@ def __init__(self, right, size=None, initval=0,
                  enable=None, reset=None, width=32, signed=True):
         _Accumulator.__init__(self, right, size, initval,
                               enable, reset, width, signed)
+        self.graph_label = 'ReduceMax'
 
 
 class ReduceMin(_Accumulator):
@@ -2658,6 +2663,7 @@ def __init__(self, right, size=None, initval=0,
                  enable=None, reset=None, width=32, signed=True):
         _Accumulator.__init__(self, right, size, initval,
                               enable, reset, width, signed)
+        self.graph_label = 'ReduceMin'
 
 
 class ReduceCustom(_Accumulator):
@@ -3444,45 +3450,49 @@ def enable(self):
 def ReduceArgMax(right, size=None, initval=0,
                  enable=None, reset=None, width=32, signed=True):
 
-    reduce_max = ReduceMax(right, size, initval,
-                           enable, reset, width, signed)
+    _max = ReduceMax(right, size, initval,
+                     enable, reset, width, signed)
     counter = Counter(size, control=right, enable=enable, reset=reset)
-    update = NotEq(reduce_max, reduce_max.prev(1))
+    update = NotEq(_max, _max.prev(1))
     update.latency = 0
-    return Predicate(counter, update)
+    index = Predicate(counter, update)
+    return index, _max
 
 
 def ReduceArgMin(right, size=None, initval=0,
                  enable=None, reset=None, width=32, signed=True):
 
-    reduce_min = ReduceMin(right, size, initval,
-                           enable, reset, width, signed)
+    _min = ReduceMin(right, size, initval,
+                     enable, reset, width, signed)
     counter = Counter(size, control=right, enable=enable, reset=reset)
-    update = NotEq(reduce_min, reduce_min.prev(1))
+    update = NotEq(_min, reduce_min.prev(1))
     update.latency = 0
-    return Predicate(counter, update)
+    index = Predicate(counter, update)
+    return index, _min
 
 
 def ReduceArgMaxValid(right, size=None, initval=0,
                       enable=None, reset=None, width=32, signed=True):
 
-    reduce_max, valid = ReduceMaxValid(right, size, initval,
-                                       enable, reset, width, signed)
+    _max, valid = ReduceMaxValid(right, size, initval,
+                                 enable, reset, width, signed)
     counter = Counter(size, control=right, enable=enable, reset=reset)
-    update = NotEq(reduce_max, reduce_max.prev(1))
+    update = NotEq(_max, _max.prev(1))
     update.latency = 0
-    return Predicate(counter, update), valid
+    index = Predicate(counter, update)
+    return index, _max, valid
 
 
 def ReduceArgMinValid(right, size=None, initval=0,
                       enable=None, reset=None, width=32, signed=True):
 
-    reduce_min, valid = ReduceMinValid(right, size, initval,
-                                       enable, reset, width, signed)
+    _min, valid = ReduceMinValid(right, size, initval,
+                                 enable, reset, width, signed)
     counter = Counter(size, control=right, enable=enable, reset=reset)
-    update = NotEq(reduce_min, reduce_min.prev(1))
+    update = NotEq(_min, _min.prev(1))
     update.latency = 0
-    return Predicate(counter, update), valid
+    index = Predicate(counter, update)
+    return index, _min, valid
 
 
 def make_condition(*cond, **kwargs):