Extension to nearest power of two and 'test_lza' changes

coreblocks/func_blocks/fu/fpu/lza.py

@@ -1,6 +1,16 @@
 from amaranth import *
 from transactron import TModule, Method, def_method
 from coreblocks.func_blocks.fu.fpu.fpu_common import FPUParams
+from math import log2
+from transactron.utils.amaranth_ext import count_leading_zeros
+
+
+def nearestpow2(n):
+    a = int(log2(n))
+    if 2**a == n:
+        return n
+    else:
+        return 2 ** (a + 1)
 
 
 class LZAMethodLayout:
@@ -73,10 +83,14 @@ def elaborate(self, platform):
 
         @def_method(m, self.predict_request)
         def _(sig_a, sig_b, carry):
+            f_size = nearestpow2(self.lza_params.sig_width)
+            filler_size = f_size - self.lza_params.sig_width
+            lower_ones = Const((2**filler_size) - 1, f_size)
+
             t = Signal(self.lza_params.sig_width + 1)
             g = Signal(self.lza_params.sig_width + 1)
             z = Signal(self.lza_params.sig_width + 1)
-            f = Signal(self.lza_params.sig_width)
+            f = Signal(f_size)
             shift_amount = Signal(range(self.lza_params.sig_width))
             is_zero = Signal(1)
 
@@ -88,12 +102,11 @@ def _(sig_a, sig_b, carry):
                 m.d.av_comb += z[0].eq(1)
 
             for i in reversed(range(1, self.lza_params.sig_width + 1)):
-                m.d.av_comb += f[i - 1].eq((t[i] ^ z[i - 1]))
+                m.d.av_comb += f[i + filler_size - 1].eq((t[i] ^ z[i - 1]))
 
             m.d.av_comb += shift_amount.eq(0)
-            for i in reversed(range(self.lza_params.sig_width)):
-                with m.If(f[self.lza_params.sig_width - i - 1]):
-                    m.d.av_comb += shift_amount.eq(i)
+            m.d.av_comp += f.eq(f | lower_ones)
+            m.d.av_comb += shift_amount.eq(count_leading_zeros(f))
 
             m.d.av_comb += is_zero.eq((carry & t[1 : self.lza_params.sig_width].all()))
 

test/func_blocks/fu/test_lza.py

@@ -7,6 +7,18 @@
 from amaranth import *
 
 
+def clz(sig_a, sig_b, carry, sig_width):
+    zeros = 0
+    msb_bit_mask = 1 << (sig_width - 1)
+    sum = sig_a + sig_b + carry
+    while 1:
+        if not (sum & msb_bit_mask):
+            zeros += 1
+            sum = sum << 1
+        else:
+            return zeros
+
+
 class TestLZA(TestCaseWithSimulator):
     class LZAModuleTest(Elaboratable):
         def __init__(self, params: FPUParams):
@@ -38,26 +50,17 @@ def test_manual(self):
         help_values = TestLZA.HelpValues(params)
         lza = TestLZA.LZAModuleTest(params)
 
-        def clz(sig_a, sig_b, carry):
-            zeros = 0
-            msb_bit_mask = 1 << (params.sig_width - 1)
-            sum = sig_a + sig_b + carry
-            while 1:
-                if not (sum & msb_bit_mask):
-                    zeros += 1
-                    sum = sum << 1
-                else:
-                    return zeros
-
-        def random_test():
+        def random_test(seed, iters):
             xor_mask = (2**params.sig_width) - 1
-            sig_a = randint(1 << (params.sig_width - 1), (2**params.sig_width) - 1)
-            sig_b = randint(1 << (params.sig_width - 1), sig_a)
-            sig_b = (sig_b ^ xor_mask) | (1 << params.sig_width)
-            resp = yield from lza.predict_request_adapter.call({"sig_a": sig_a, "sig_b": sig_b, "carry": 0})
-            pred_lz = resp["shift_amount"]
-            true_lz = clz(sig_a, sig_b, 0)
-            assert pred_lz == true_lz or (pred_lz + 1) == true_lz
+            random.seed(seed)
+            for _ in range(iters):
+                sig_a = randint(1 << (params.sig_width - 1), (2**params.sig_width) - 1)
+                sig_b = randint(1 << (params.sig_width - 1), sig_a)
+                sig_b = (sig_b ^ xor_mask) | (1 << params.sig_width)
+                resp = yield from lza.predict_request_adapter.call({"sig_a": sig_a, "sig_b": sig_b, "carry": 0})
+                pred_lz = resp["shift_amount"]
+                true_lz = clz(sig_a, sig_b, 0, params.sig_width)
+                assert pred_lz == true_lz or (pred_lz + 1) == true_lz
 
         def lza_test():
             test_cases = [
@@ -143,7 +146,7 @@ def lza_test():
 
         def test_process():
             yield from lza_test()
-            yield from random_test()
+            yield from random_test(2024, 20)
 
         with self.run_simulation(lza) as sim:
             sim.add_process(test_process)