From ab130e170a268e326ecfe8441d102e87a054e9e5 Mon Sep 17 00:00:00 2001
From: Alessandro Comodi <acomodi@antmicro.com>
Date: Mon, 6 Sep 2021 17:04:42 +0200
Subject: [PATCH] lpddr5: add write leveling support

Signed-off-by: Alessandro Comodi <acomodi@antmicro.com>
---
 litedram/phy/lpddr5/basephy.py | 49 +++++++++++++++++++++++++++++++---
 test/test_lpddr5.py            | 32 ++++++++++++++++++++++
 2 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/litedram/phy/lpddr5/basephy.py b/litedram/phy/lpddr5/basephy.py
index 166a5fb..556bea2 100644
--- a/litedram/phy/lpddr5/basephy.py
+++ b/litedram/phy/lpddr5/basephy.py
@@ -361,8 +361,52 @@ class LPDDR5PHY(Module, AutoCSR):
 
         wck_out = {2: wck_pattern[::2], 4: wck_pattern}[wck_ck_ratio]
         assert len(wck_out) == len(self.out.wck[0]), (len(wck_out), len(self.out.wck))
-        for byte in range(databits//8):
-            self.comb += self.out.wck[byte].eq(wck_out)
+
+        # WCK2CK leveling --------------------------------------------------------------------------
+
+        # Strobe needs to be high for tWCKTGGL which is 8 tWCK periods
+        # NOTE: WCK2CK leveling always happens at 2:1 WCK2CK ratio.
+        twcktggl = 4
+        wckl_strobe_dly = TappedDelayLine(
+            signal = self._wlevel_strobe.re,
+            ntaps  = twcktggl
+        )
+        self.submodules += wckl_strobe_dly
+
+        wckl_strobe_en = Signal()
+        self.comb += [
+            wckl_strobe_en.eq(reduce(or_, wckl_strobe_dly.taps[0:twcktggl]))
+        ]
+
+        wckl_pattern = Signal(8)
+        self.comb += [
+            wckl_pattern.eq(bitpattern(patterns["disabled"])),
+            If(wckl_strobe_en,
+                wckl_pattern.eq(bitpattern(patterns["toggle"])),
+            )
+        ]
+
+        wckl_out = {2: wckl_pattern[::2], 4: wckl_pattern}[wck_ck_ratio]
+
+        wck_pattern_selected = Signal(2*wck_ck_ratio)
+        self.comb += [
+            If(self._wlevel_en.storage,
+                wck_pattern_selected.eq(wckl_out)
+            ).Else(
+                wck_pattern_selected.eq(wck_out)
+            )
+        ]
+
+        for byte in range(self.databits//8):
+            # output
+            self.submodules += BitSlip(
+                dw     = 2*wck_ck_ratio,
+                cycles = bitslip_cycles,
+                rst    = self.get_rst(byte, self._wdly_dq_bitslip_rst.re),
+                slp    = self.get_inc(byte, self._wdly_dq_bitslip.re),
+                i      = wck_pattern_selected,
+                o      = self.out.wck[byte],
+            )
 
         # Write Control Path -----------------------------------------------------------------------
         wrtap = write_latency - 1
@@ -439,7 +483,6 @@ class LPDDR5PHY(Module, AutoCSR):
             self.dfi.p0.rddata_valid.eq(rddata_converter.source.valid),
         ]
 
-
         # -2 is to take into account the serialization time for wck
         read_wck_latency = cmd_latency + cl + burst_ck_cycles - 2
         write_wck_latency = cmd_latency + cwl + burst_ck_cycles - 2
diff --git a/test/test_lpddr5.py b/test/test_lpddr5.py
index b89cbb0..885d8c2 100644
--- a/test/test_lpddr5.py
+++ b/test/test_lpddr5.py
@@ -682,6 +682,38 @@ class LPDDR5Tests(unittest.TestCase):
                     },
                 )
 
+    def test_lpddr5_wck_leveling(self):
+
+        # Test that correct WCK sequence is generated during WCK sync before burst write for WCK:CK=4:1
+        for wck_ck_ratio in [2, 4]:
+            with self.subTest(wck_ck_ratio=wck_ck_ratio):
+                phy = LPDDR5SimPHY(sys_clk_freq=50e6, wck_ck_ratio=wck_ck_ratio)
+
+                def write_leveling(pads):
+                    for _ in range(4):
+                        yield from phy._wlevel_en.write(1)
+                        yield from phy._wlevel_strobe.write(1)
+                        for i in range(4):
+                            yield
+                        yield from phy._wlevel_en.write(0)
+
+
+                self.run_test(phy,
+                    dfi_sequence = {},
+                    pad_checkers = {
+                        f"sys4x_270": {
+                            "wck0": "0000" + \
+                                    "0000" * 3 + "1010" * 4 + \
+                                    "0000" * 3 + "1010" * 4 + \
+                                    "0000" * 3 + "1010" * 4 + \
+                                    "0000" * 3 + "1010" * 4,
+                        },
+                    },
+                    pad_generators = {
+                        "sys": write_leveling,
+                    }
+                )
+
 
 class VerilatorLPDDR5Tests(unittest.TestCase):
     def check_logs(self, logs, allowed):