Add TaggedCounter (#637)

coreblocks/func_blocks/fu/alu.py

@@ -3,6 +3,7 @@
 
 from transactron import *
 from transactron.lib import FIFO
+from transactron.lib.metrics import *
 
 from coreblocks.frontend.decoder.isa import Funct3, Funct7
 from coreblocks.frontend.decoder.optypes import OpType
@@ -219,9 +220,17 @@ def __init__(self, gen_params: GenParams, alu_fn=AluFn()):
         self.issue = Method(i=layouts.issue)
         self.accept = Method(o=layouts.accept)
 
+        self.perf_instr = TaggedCounter(
+            "backend.fu.alu.instr",
+            "Counts of instructions executed by the jumpbranch unit",
+            tags=AluFn.Fn,
+        )
+
     def elaborate(self, platform):
         m = TModule()
 
+        m.submodules += [self.perf_instr]
+
         m.submodules.alu = alu = Alu(self.gen_params, alu_fn=self.alu_fn)
         m.submodules.fifo = fifo = FIFO(self.gen_params.get(FuncUnitLayouts).accept, 2)
         m.submodules.decoder = decoder = self.alu_fn.get_decoder(self.gen_params)
@@ -238,6 +247,8 @@ def _(arg):
             m.d.comb += alu.in1.eq(arg.s1_val)
             m.d.comb += alu.in2.eq(Mux(arg.imm, arg.imm, arg.s2_val))
 
+            self.perf_instr.incr(m, decoder.decode_fn)
+
             fifo.write(m, rob_id=arg.rob_id, result=alu.out, rp_dst=arg.rp_dst, exception=0)
 
         return m

coreblocks/func_blocks/fu/jumpbranch.py

@@ -136,16 +136,22 @@ def __init__(self, gen_params: GenParams, jb_fn=JumpBranchFn()):
         self.dm = gen_params.get(DependencyManager)
         self.dm.add_dependency(BranchVerifyKey(), self.fifo_branch_resolved.read)
 
-        self.perf_jumps = HwCounter("backend.fu.jumpbranch.jumps", "Number of jump instructions issued")
-        self.perf_branches = HwCounter("backend.fu.jumpbranch.branches", "Number of branch instructions issued")
+        self.perf_instr = TaggedCounter(
+            "backend.fu.jumpbranch.instr",
+            "Counts of instructions executed by the jumpbranch unit",
+            tags=JumpBranchFn.Fn,
+        )
         self.perf_misaligned = HwCounter(
             "backend.fu.jumpbranch.misaligned", "Number of instructions with misaligned target address"
         )
 
     def elaborate(self, platform):
         m = TModule()
 
-        m.submodules += [self.perf_jumps, self.perf_branches, self.perf_misaligned]
+        m.submodules += [
+            self.perf_instr,
+            self.perf_misaligned,
+        ]
 
         m.submodules.jb = jb = JumpBranch(self.gen_params, fn=self.jb_fn)
         m.submodules.fifo_res = fifo_res = FIFO(self.gen_params.get(FuncUnitLayouts).accept, 2)
@@ -169,12 +175,10 @@ def _(arg):
             m.d.top_comb += jb.in_rvc.eq(arg.exec_fn.funct7)
 
             is_auipc = decoder.decode_fn == JumpBranchFn.Fn.AUIPC
-            is_jump = (decoder.decode_fn == JumpBranchFn.Fn.JAL) | (decoder.decode_fn == JumpBranchFn.Fn.JALR)
 
             jump_result = Mux(jb.taken, jb.jmp_addr, jb.reg_res)
 
-            self.perf_jumps.incr(m, cond=is_jump)
-            self.perf_branches.incr(m, cond=(~is_jump & ~is_auipc))
+            self.perf_instr.incr(m, decoder.decode_fn)
 
             exception = Signal()
             exception_report = self.dm.get_dependency(ExceptionReportKey())
@@ -216,7 +220,7 @@ def _(arg):
                 log.debug(
                     m,
                     True,
-                    "jumping from 0x{:08x} to 0x{:08x}; misprediction: {}",
+                    "branch resolved from 0x{:08x} to 0x{:08x}; misprediction: {}",
                     jb.in_pc,
                     jump_result,
                     misprediction,

test/transactron/test_metrics.py

@@ -1,6 +1,9 @@
 import json
 import random
 import queue
+from typing import Type
+from enum import IntFlag, IntEnum, auto, Enum
+
 from parameterized import parameterized_class
 
 from amaranth import *
@@ -138,6 +141,85 @@ def test_process():
             sim.add_sync_process(test_process)
 
 
+class OneHotEnum(IntFlag):
+    ADD = auto()
+    XOR = auto()
+    OR = auto()
+
+
+class PlainIntEnum(IntEnum):
+    TEST_1 = auto()
+    TEST_2 = auto()
+    TEST_3 = auto()
+
+
+class TaggedCounterCircuit(Elaboratable):
+    def __init__(self, tags: range | Type[Enum] | list[int]):
+        self.counter = TaggedCounter("counter", "", tags=tags)
+
+        self.cond = Signal()
+        self.tag = Signal(self.counter.tag_width)
+
+    def elaborate(self, platform):
+        m = TModule()
+
+        m.submodules.counter = self.counter
+
+        with Transaction().body(m):
+            self.counter.incr(m, self.tag, cond=self.cond)
+
+        return m
+
+
+class TestTaggedCounter(TestCaseWithSimulator):
+    def setUp(self) -> None:
+        random.seed(42)
+
+    def do_test_enum(self, tags: range | Type[Enum] | list[int], tag_values: list[int]):
+        m = TaggedCounterCircuit(tags)
+        DependencyContext.get().add_dependency(HwMetricsEnabledKey(), True)
+
+        counts: dict[int, int] = {}
+        for i in tag_values:
+            counts[i] = 0
+
+        def test_process():
+            for _ in range(200):
+                for i in tag_values:
+                    self.assertEqual(counts[i], (yield m.counter.counters[i].value))
+
+                tag = random.choice(list(tag_values))
+
+                yield m.cond.eq(1)
+                yield m.tag.eq(tag)
+                yield
+                yield m.cond.eq(0)
+                yield
+
+                counts[tag] += 1
+
+        with self.run_simulation(m) as sim:
+            sim.add_sync_process(test_process)
+
+    def test_one_hot_enum(self):
+        self.do_test_enum(OneHotEnum, [e.value for e in OneHotEnum])
+
+    def test_plain_int_enum(self):
+        self.do_test_enum(PlainIntEnum, [e.value for e in PlainIntEnum])
+
+    def test_negative_range(self):
+        r = range(-10, 15, 3)
+        self.do_test_enum(r, list(r))
+
+    def test_positive_range(self):
+        r = range(0, 30, 2)
+        self.do_test_enum(r, list(r))
+
+    def test_value_list(self):
+        values = [-2137, 2, 4, 8, 42]
+        self.do_test_enum(values, values)
+
+
 class ExpHistogramCircuit(Elaboratable):
     def __init__(self, bucket_cnt: int, sample_width: int):
         self.sample_width = sample_width

transactron/lib/metrics.py

@@ -1,14 +1,14 @@
 from dataclasses import dataclass, field
 from dataclasses_json import dataclass_json
-from typing import Optional
+from typing import Optional, Type
 from abc import ABC
+from enum import Enum
 
 from amaranth import *
-from amaranth.utils import bits_for
+from amaranth.utils import bits_for, ceil_log2, exact_log2
 
-from transactron.utils import ValueLike
+from transactron.utils import ValueLike, OneHotSwitchDynamic, SignalBundle
 from transactron import Method, def_method, TModule
-from transactron.utils import SignalBundle
 from transactron.lib import FIFO
 from transactron.utils.dependencies import ListKey, DependencyContext, SimpleKey
 
@@ -17,6 +17,7 @@
     "MetricModel",
     "HwMetric",
     "HwCounter",
+    "TaggedCounter",
     "HwExpHistogram",
     "LatencyMeasurer",
     "HardwareMetricsManager",
@@ -230,6 +231,127 @@ def incr(self, m: TModule, *, cond: ValueLike = C(1)):
             self._incr(m)
 
 
+class TaggedCounter(Elaboratable, HwMetric):
+    """Hardware Tagged Counter
+
+    Like HwCounter, but contains multiple counters, each with its own tag.
+    At a time a single counter can be increased and the value of the tag
+    can be provided dynamically. The type of the tag can be either an int
+    enum, a range or a list of integers (negative numbers are ok).
+
+    Internally, it detects if tag values can be one-hot encoded and if so,
+    it generates more optimized circuit.
+
+    Attributes
+    ----------
+    tag_width: int
+        The length of the signal holding a tag value.
+    one_hot: bool
+        Whether tag values can be one-hot encoded.
+    counters: dict[int, HwMetricRegisters]
+        Mapping from a tag value to a register holding a counter for that tag.
+    """
+
+    def __init__(
+        self,
+        fully_qualified_name: str,
+        description: str = "",
+        *,
+        tags: range | Type[Enum] | list[int],
+        registers_width: int = 32,
+    ):
+        """
+        Parameters
+        ----------
+        fully_qualified_name: str
+            The fully qualified name of the metric.
+        description: str
+            A human-readable description of the metric's functionality.
+        tags: range | Type[Enum] | list[int]
+            Tag values.
+        registers_width: int
+            Width of the underlying registers. Defaults to 32 bits.
+        """
+
+        super().__init__(fully_qualified_name, description)
+
+        if isinstance(tags, range) or isinstance(tags, list):
+            counters_meta = [(i, f"{i}") for i in tags]
+        else:
+            counters_meta = [(i.value, i.name) for i in tags]
+
+        values = [value for value, _ in counters_meta]
+        self.tag_width = max(bits_for(max(values)), bits_for(min(values)))
+
+        self.one_hot = True
+        negative_values = False
+        for value in values:
+            if value < 0:
+                self.one_hot = False
+                negative_values = True
+                break
+
+            log = ceil_log2(value)
+            if 2**log != value:
+                self.one_hot = False
+
+        self._incr = Method(i=[("tag", Shape(self.tag_width, signed=negative_values))])
+
+        self.counters: dict[int, HwMetricRegister] = {}
+        for tag_value, name in counters_meta:
+            value_str = ("1<<" + str(exact_log2(tag_value))) if self.one_hot else str(tag_value)
+            description = f"the counter for tag {name} (value={value_str})"
+
+            self.counters[tag_value] = HwMetricRegister(
+                name,
+                registers_width,
+                description,
+            )
+
+        self.add_registers(list(self.counters.values()))
+
+    def elaborate(self, platform):
+        if not self.metrics_enabled():
+            return TModule()
+
+        m = TModule()
+
+        @def_method(m, self._incr)
+        def _(tag):
+            if self.one_hot:
+                sorted_tags = sorted(list(self.counters.keys()))
+                for i in OneHotSwitchDynamic(m, tag):
+                    counter = self.counters[sorted_tags[i]]
+                    m.d.sync += counter.value.eq(counter.value + 1)
+            else:
+                for tag_value, counter in self.counters.items():
+                    with m.If(tag == tag_value):
+                        m.d.sync += counter.value.eq(counter.value + 1)
+
+        return m
+
+    def incr(self, m: TModule, tag: ValueLike, *, cond: ValueLike = C(1)):
+        """
+        Increases the counter of a given tag by 1.
+
+        Should be called in the body of either a transaction or a method.
+
+        Parameters
+        ----------
+        m: TModule
+            Transactron module
+        tag: ValueLike
+            The tag of the counter.
+        cond: ValueLike
+            When set to high, the counter will be increased. By default set to high.
+        """
+        if not self.metrics_enabled():
+            return
+
+        with m.If(cond):
+            self._incr(m, tag)
+
+
 class HwExpHistogram(Elaboratable, HwMetric):
     """Hardware Exponential Histogram