From 15d191937f05a8d8fc4fbfedc9fa888aa5dee80e Mon Sep 17 00:00:00 2001
From: Martijn Bastiaan <martijn@qbaylogic.com>
Date: Fri, 7 Feb 2025 16:47:41 +0100
Subject: [PATCH 01/12] Allow running `format` in non-project-root

---
 nix/bin/format | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nix/bin/format b/nix/bin/format
index 2f79e5586..c8d8eec4d 100755
--- a/nix/bin/format
+++ b/nix/bin/format
@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 ROOT=$(git rev-parse --show-toplevel)
+cd "${ROOT}" || exit 1
 
 echo "Formatting Cabal files.."
 "${ROOT}"/.github/scripts/cabal-gild.sh

From 44b5778a46048a38dd8ff0cc06ac090a80ada272 Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Thu, 30 Jan 2025 09:42:35 +0100
Subject: [PATCH 02/12] Add simple PE for switch demo

---
 bittide/bittide.cabal                         |   1 +
 .../Bittide/SwitchDemoProcessingElement.hs    | 143 ++++++++++++++++++
 2 files changed, 144 insertions(+)
 create mode 100644 bittide/src/Bittide/SwitchDemoProcessingElement.hs

diff --git a/bittide/bittide.cabal b/bittide/bittide.cabal
index 707989e22..877082310 100644
--- a/bittide/bittide.cabal
+++ b/bittide/bittide.cabal
@@ -150,6 +150,7 @@ library
     Bittide.ScatterGather
     Bittide.SharedTypes
     Bittide.Switch
+    Bittide.SwitchDemoProcessingElement
     Bittide.Transceiver
     Bittide.Transceiver.Cdc
     Bittide.Transceiver.Comma
diff --git a/bittide/src/Bittide/SwitchDemoProcessingElement.hs b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
new file mode 100644
index 000000000..94cd879ec
--- /dev/null
+++ b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
@@ -0,0 +1,143 @@
+-- SPDX-FileCopyrightText: 2025 Google LLC
+--
+-- SPDX-License-Identifier: Apache-2.0
+{-# LANGUAGE NumericUnderscores #-}
+
+module Bittide.SwitchDemoProcessingElement where
+
+import Clash.Prelude
+
+import Data.Maybe (fromMaybe)
+import Data.Tuple (swap)
+import GHC.Stack (HasCallStack)
+
+import Protocols
+import Protocols.Wishbone
+
+import Bittide.SharedTypes (Bytes)
+import Bittide.Wishbone (wbToVec)
+import Clash.Sized.Vector.ToTuple (vecToTuple)
+
+{- | Multiplying by 3 should always fit, though if n~1, the output type is `Index 3`
+which doesn't fit the 3 we're multiplying by hence yielding an undefined. This
+function works around that.
+-}
+zeroExtendTimesThree :: forall n. (1 <= n, KnownNat n) => Index n -> Index (n * 3)
+zeroExtendTimesThree = truncateB . mul (3 :: Index 4)
+
+-- | Simple processing element used for the Bittide switch demo.
+switchDemoPe ::
+  forall bufferSize dom.
+  ( HasCallStack
+  , HiddenClockResetEnable dom
+  , 1 <= bufferSize
+  ) =>
+  -- | Size of buffer in number of "tri-cycles". That is, we always store 3 64-bit words:
+  -- DNA (32 msbs), DNA (64 lsbs), local clock cycle counter.
+  SNat bufferSize ->
+  -- | Local clock cycle counter
+  Signal dom (Unsigned 64) ->
+  -- | Incoming crossbar link
+  Signal dom (BitVector 64) ->
+  -- | Device DNA
+  Signal dom (Maybe (BitVector 96)) ->
+  -- | When to read from the crossbar link
+  Signal dom (Unsigned 64) ->
+  -- | How many tri-cycles to read from the crossbar link
+  Signal dom (Index bufferSize) ->
+  -- | When to write to the crossbar link
+  Signal dom (Unsigned 64) ->
+  -- | How many tri-cycles to write to the crossbar link. Includes writing \"own\" data.
+  Signal dom (Index (bufferSize + 1)) ->
+  ( -- \| Outgoing crossbar link
+    Signal dom (BitVector 64)
+  , -- \| Buffer output
+    Signal dom (Vec (bufferSize * 3) (BitVector 64))
+  )
+switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart writeCycles =
+  (linkOut, buffer)
+ where
+  readCyclesExtended = zeroExtendTimesThree <$> readCycles
+  writeCyclesExtended = zeroExtendTimesThree <$> writeCycles
+
+  localData :: Signal dom (Vec 3 (BitVector 64))
+  localData = bundle ((pack <$> localCounter) :> unbundle dnaVec)
+   where
+    dnaVec :: Signal dom (Vec 2 (BitVector 64))
+    dnaVec = bitCoerce . zeroExtend <$> dnaLocked
+    dnaLocked = fromMaybe 0xBAAB_BAAB_BAAB_BAAB_BAAB_BAAB <$> maybeDna
+
+  linkOut = stateToLinkOutput <$> peState <*> buffer <*> localData
+
+  stateToLinkOutput ::
+    SimplePeState bufferSize ->
+    Vec (bufferSize * 3) (BitVector 64) ->
+    Vec 3 (BitVector 64) ->
+    BitVector 64
+  stateToLinkOutput state buf locData =
+    case state of
+      Write i
+        | i <= 2 -> locData !! i
+        | otherwise -> buf !! (i - 3)
+      _ -> 0xAAAA_BBBB_AAAA_BBBB
+
+  -- \| The buffer stores all the incoming bittide data. For the Bittide Switch demo,
+  -- each FPGA sends its DNA and local clock cycle counter, along with all received data.
+  -- The last FPGA will therefore receive all DNAs and local clock cycle counters.
+  buffer :: (HasCallStack) => Signal dom (Vec (bufferSize * 3) (BitVector 64))
+  buffer = bundle $ regEn <$> initVec <*> enableVec <*> linkInVec
+   where
+    initVec = iterateI (+ 1) 0xABBA_ABBA_ABBA_0000
+    linkInVec = repeat linkIn
+
+    enableVec :: (HasCallStack) => Vec (bufferSize * 3) (Signal dom Bool)
+    enableVec = unbundle $ go <$> peState
+     where
+      go :: (HasCallStack) => SimplePeState bufferSize -> Vec (bufferSize * 3) Bool
+      go (Read x) = (== x) <$> indicesI
+      go _ = repeat False
+
+  prevPeState = register Idle peState
+
+  peState =
+    update
+      <$> localCounter
+      <*> readStart
+      <*> readCyclesExtended
+      <*> writeStart
+      <*> writeCyclesExtended
+      <*> prevPeState
+   where
+    update ::
+      -- \| Local clock cycle counter
+      Unsigned 64 ->
+      -- \| When to read from the crossbar link
+      Unsigned 64 ->
+      -- \| How many cycles to read from the crossbar link
+      Index (bufferSize * 3) ->
+      -- \| When to write to the crossbar link
+      Unsigned 64 ->
+      -- \| How many cycles to write to the crossbar link
+      Index ((bufferSize + 1) * 3) ->
+      SimplePeState bufferSize ->
+      SimplePeState bufferSize
+    update cntr rs rc ws wc state =
+      case state of
+        Idle -> nextState
+        Read x
+          | x >= rc - 1 -> nextState
+          | otherwise -> Read (satSucc SatBound x)
+        Write x
+          | x >= wc - 1 -> nextState
+          | otherwise -> Write (satSucc SatBound x)
+     where
+      nextState
+        | cntr == ws && wc > 0 = Write 0
+        | cntr == rs && rc > 0 = Read 0
+        | otherwise = Idle
+
+data SimplePeState bufferSize
+  = Idle
+  | Read (Index (bufferSize * 3))
+  | Write (Index ((bufferSize + 1) * 3))
+  deriving (Generic, NFDataX, Eq, Show)

From 40e31fcd0f3ded49953ef0eff8bf2a3d8d4654a8 Mon Sep 17 00:00:00 2001
From: Martijn Bastiaan <martijn@qbaylogic.com>
Date: Fri, 7 Feb 2025 16:32:13 +0100
Subject: [PATCH 03/12] Add tests for `switchDemoPe`

---
 bittide/bittide.cabal                         |   1 +
 .../Tests/SwitchDemoProcessingElement.hs      | 184 ++++++++++++++++++
 bittide/tests/UnitTests.hs                    |   2 +
 3 files changed, 187 insertions(+)
 create mode 100644 bittide/tests/Tests/SwitchDemoProcessingElement.hs

diff --git a/bittide/bittide.cabal b/bittide/bittide.cabal
index 877082310..45b2883fc 100644
--- a/bittide/bittide.cabal
+++ b/bittide/bittide.cabal
@@ -222,6 +222,7 @@ test-suite unittests
     Tests.Shared
     Tests.StabilityChecker
     Tests.Switch
+    Tests.SwitchDemoProcessingElement
     Tests.Transceiver
     Tests.Transceiver.Prbs
     Tests.Transceiver.WordAlign
diff --git a/bittide/tests/Tests/SwitchDemoProcessingElement.hs b/bittide/tests/Tests/SwitchDemoProcessingElement.hs
new file mode 100644
index 000000000..27f4306c5
--- /dev/null
+++ b/bittide/tests/Tests/SwitchDemoProcessingElement.hs
@@ -0,0 +1,184 @@
+-- SPDX-FileCopyrightText: 2025 Google LLC
+--
+-- SPDX-License-Identifier: Apache-2.0
+{-# LANGUAGE NumericUnderscores #-}
+{-# LANGUAGE OverloadedStrings #-}
+
+module Tests.SwitchDemoProcessingElement where
+
+import Clash.Prelude hiding (someNatVal, withSomeSNat)
+
+import Test.Tasty
+import Test.Tasty.HUnit
+import Test.Tasty.Hedgehog
+
+import Control.Monad (forM_)
+import Data.Proxy (Proxy (..))
+import GHC.TypeNats (someNatVal)
+
+import Bittide.SwitchDemoProcessingElement
+
+import qualified Hedgehog as H
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Range as Range
+
+import qualified Clash.Explicit.Prelude as E
+import qualified Data.List as L
+
+import Clash.Explicit.Reset (noReset)
+import Clash.Hedgehog.Sized.BitVector (genDefinedBitVector)
+import Clash.Hedgehog.Sized.Index (genIndex)
+import Clash.Hedgehog.Sized.Unsigned (genUnsigned)
+import Hedgehog ((===))
+
+main :: IO ()
+main = defaultMain tests
+
+tests :: TestTree
+tests =
+  testGroup
+    "SwitchDemoProcessingElement"
+    [ testPropertyNamed "prop_readThenWrite" "prop_readThenWrite" prop_readThenWrite
+    , testCase "case_zeroExtendTimesThree" case_zeroExtendTimesThree
+    ]
+
+withSomeSNat :: Natural -> (forall (n :: Nat). SNat n -> r) -> r
+withSomeSNat n f = case someNatVal n of
+  SomeNat (_ :: Proxy n) -> f (SNat @n)
+
+linearLength :: (Integral a) => a -> a -> Range.Range a
+linearLength start len = Range.linear start (start + len)
+
+singletonInt :: (Integral a) => a -> Range.Range Int
+singletonInt = Range.singleton . fromIntegral
+
+-- | Exhaustive test for 'zeroExtendTimesThree' for n ~ [1..64]
+case_zeroExtendTimesThree :: Assertion
+case_zeroExtendTimesThree =
+  forM_ [0 .. 63] $ \nMinusOne ->
+    withSomeSNat nMinusOne $ \(succSNat -> (SNat :: SNat n)) -> do
+      forM_ [(0 :: Index n) ..] $ \i -> do
+        let
+          actual = fromIntegral (zeroExtendTimesThree @n i)
+          expected = fromIntegral @_ @Integer i * 3
+        actual @?= expected
+
+prop_readThenWrite :: H.Property
+prop_readThenWrite = H.property $ do
+  bufferSizeMinusOne <- H.forAll $ Gen.integral (Range.linear 0 10)
+  withSomeSNat bufferSizeMinusOne $ \(succSNat -> bufferSizeSNat@(SNat :: SNat bufferSize)) -> do
+    nReadTriCycles <-
+      H.forAll
+        $ Gen.frequency
+          [ (30, Gen.constant 0)
+          , (70, genIndex Range.linearBounded)
+          ]
+    nWriteTriCycles <-
+      if nReadTriCycles == 0
+        then H.forAll $ genIndex (Range.linear 1 maxBound)
+        else H.forAll $ genIndex Range.linearBounded
+
+    let
+      cyclesPerReadWrite = 3 :: Unsigned 64
+      nReadCycles = cyclesPerReadWrite * fromIntegral nReadTriCycles
+      nWriteCycles = cyclesPerReadWrite * fromIntegral nWriteTriCycles
+      maxIdle1 = 100
+      maxIdle2 = 100
+
+    -- Notice that the PE needs a single clock cycle in its idle state to function
+    -- correctly. Hence, we always start reading a minimum at clockStart+1.
+    readData <- H.forAll $ Gen.list (singletonInt nReadCycles) genDefinedBitVector
+    clockStart <- H.forAll $ genUnsigned @_ @64 (Range.linear 0 100)
+    readStart <-
+      H.forAll
+        $ Gen.frequency
+          [ (30, Gen.constant clockStart)
+          , (70, genUnsigned @_ @64 (linearLength clockStart maxIdle1))
+          ]
+    let readEnd = readStart + fromIntegral nReadCycles
+    writeStart <-
+      H.forAll
+        $ Gen.frequency
+          [ (30, Gen.constant readEnd)
+          , (70, genUnsigned @_ @64 (linearLength readEnd maxIdle2))
+          ]
+    deviceDna <- H.forAll genDefinedBitVector
+
+    let immediateRead = nReadCycles > 0 && readStart == clockStart
+    H.cover 5 "Read in the very first cycle we're allowed to" immediateRead
+
+    let immediateWrite = nWriteCycles > 0 && writeStart == clockStart
+    H.cover 5 "Write in the very first cycle we're allowed to" immediateWrite
+
+    let isBackToBack = nReadCycles > 0 && nWriteCycles > 0 && readEnd == writeStart
+    H.cover 5 "Back-to-back read/write" isBackToBack
+
+    let
+      idle1length = readStart - clockStart
+      idle2length = writeStart - readEnd
+      idle1in = L.replicate (fromIntegral idle1length) 0
+      crossBarIn = fromList (idle1in <> readData <> L.repeat 0)
+
+      out =
+        E.sample
+          $ bundle
+          $ withClockResetEnable @System clockGen noReset enableGen
+          $ switchDemoPe
+            bufferSizeSNat
+            (fromList [clockStart ..])
+            crossBarIn
+            (pure (Just deviceDna))
+            (pure readStart)
+            (pure nReadTriCycles)
+            (pure writeStart)
+            (pure nWriteTriCycles)
+
+      (idle1out, rest0) = L.splitAt (fromIntegral idle1length) out
+      (readOut, rest1) = L.splitAt (fromIntegral nReadCycles) rest0
+      (idle2out, rest2) = L.splitAt (fromIntegral idle2length) rest1
+      (writeOut, rest3) = L.splitAt (fromIntegral nWriteCycles) rest2
+
+      (idle1outs, _idle1buffers) = L.unzip idle1out
+      (readOuts, _readBuffers) = L.unzip readOut
+      (idle2outs, _idle2buffers) = L.unzip idle2out
+      (writeOuts, _writeBuffers) = L.unzip writeOut
+      (_restOuts, restBuffers) = L.unzip rest3
+
+    H.footnote $ "idle1in: " <> show idle1in
+    H.footnote $ "idle1length: " <> show idle1length
+    H.footnote $ "idle2length: " <> show idle2length
+    H.footnote $ "nReadCycles: " <> show nReadCycles
+    H.footnote $ "nWriteCycles: " <> show nWriteCycles
+    H.footnote $ "readData: " <> show readData
+    H.footnote $ "readEnd: " <> show readEnd
+    H.footnote $ "readStart: " <> show readStart
+    H.footnote $ "writeStart: " <> show writeStart
+    H.footnote $ "clockStart: " <> show clockStart
+    H.footnote $ "idle2outs: " <> showX idle2outs
+    H.footnote $ "bufferSizeSNat: " <> show bufferSizeSNat
+
+    -- Check that at the end of the simulation the buffer is what we expect it
+    -- to be. The buffer should be equal to the data we send to the PE. We don't
+    -- care about data we don't write, hence we truncate (L.take) the buffer to
+    -- match the number of read cycles.
+    case restBuffers of
+      [] -> error "Unexpected end of output"
+      (buffer : _) -> do
+        H.footnote $ "buffer: " <> show buffer
+        L.take (fromIntegral nReadCycles) (toList buffer) === readData
+
+    -- Check that idle value is written at correct times
+    let idleValue = 0xAAAA_BBBB_AAAA_BBBB
+    idle1outs === L.replicate (L.length idle1outs) idleValue
+    readOuts === L.replicate (L.length readOuts) idleValue
+    idle2outs === L.replicate (L.length idle2outs) idleValue
+
+    -- Check that the right data is written the crossbar link at the time we
+    -- expect it to.
+    let
+      -- Note we can always write one tri-cycle more than we read, since internal
+      -- data comes first.
+      relevantOutCycles = fromIntegral (min nWriteCycles (nReadCycles + cyclesPerReadWrite))
+      deviceDnaVec = bitCoerce @_ @(Vec 2 (BitVector 64)) (zeroExtend deviceDna)
+      expectedOutData = toList (pack writeStart :> deviceDnaVec) <> readData
+    L.take relevantOutCycles writeOuts === L.take relevantOutCycles expectedOutData
diff --git a/bittide/tests/UnitTests.hs b/bittide/tests/UnitTests.hs
index e4fa57950..40f7587b7 100644
--- a/bittide/tests/UnitTests.hs
+++ b/bittide/tests/UnitTests.hs
@@ -21,6 +21,7 @@ import qualified Tests.ProcessingElement.ReadElf
 import qualified Tests.ScatterGather
 import qualified Tests.StabilityChecker
 import qualified Tests.Switch
+import qualified Tests.SwitchDemoProcessingElement
 import qualified Tests.Transceiver
 import qualified Tests.Transceiver.Prbs
 import qualified Tests.Transceiver.WordAlign
@@ -36,6 +37,7 @@ tests =
     , Tests.DoubleBufferedRam.tests
     , Tests.ElasticBuffer.tests
     , Tests.ProcessingElement.ReadElf.tests
+    , Tests.SwitchDemoProcessingElement.tests
     , Tests.ScatterGather.tests
     , Tests.StabilityChecker.tests
     , Tests.Switch.tests

From 1fe40a30f69ea0a9eeecc32b722ce34d114ac2ed Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Mon, 10 Feb 2025 11:50:20 +0100
Subject: [PATCH 04/12] Add Wishbone circuit wrapper for switchDemoPe

---
 .../Bittide/SwitchDemoProcessingElement.hs    | 70 +++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/bittide/src/Bittide/SwitchDemoProcessingElement.hs b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
index 94cd879ec..a161a10a1 100644
--- a/bittide/src/Bittide/SwitchDemoProcessingElement.hs
+++ b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
@@ -141,3 +141,73 @@ data SimplePeState bufferSize
   | Read (Index (bufferSize * 3))
   | Write (Index ((bufferSize + 1) * 3))
   deriving (Generic, NFDataX, Eq, Show)
+
+{- | Wishbone circuit wrapper for `switchDemoPe`.
+
+Buffer uses 64-bit words internally, but WB interface is 32-bit.
+
+The register layout is as follows (lsbs in first 32-bit word, msbs in second):
+- Address 0-1: read start
+- Address 2-3: read cycles
+- Address 4-5: write start
+- Address 6-7: write cycles
+- Address 8-.: buffer (bufferSize*3*2)
+-}
+switchDemoPeWb ::
+  forall bufferSize dom addrW.
+  ( HiddenClockResetEnable dom
+  , KnownNat addrW
+  , 1 <= bufferSize
+  ) =>
+  SNat bufferSize ->
+  -- | Local clock cycle counter
+  Signal dom (Unsigned 64) ->
+  -- | Device DNA
+  Signal dom (Maybe (BitVector 96)) ->
+  Circuit
+    ( Wishbone dom 'Standard addrW (Bytes 4)
+    , CSignal dom (BitVector 64)
+    )
+    (CSignal dom (BitVector 64))
+switchDemoPeWb SNat localCounter maybeDna = Circuit go
+ where
+  go ((wbM2S, linkIn), _) = ((wbS2M, pure ()), linkOut)
+   where
+    readVec :: Vec (8 + bufferSize * 3 * 2) (Signal dom (BitVector 32))
+    readVec =
+      dflipflop
+        <$> ( unbundle (bitCoerce . map swapWords <$> writableRegs)
+                ++ unbundle (bitCoerce . map swapWords <$> buffer)
+            )
+
+    (linkOut, buffer) =
+      switchDemoPe
+        (SNat @bufferSize)
+        localCounter
+        linkIn
+        maybeDna
+        readStart
+        readCycles
+        writeStart
+        writeCycles
+
+    readStart = unpack <$> rs
+    readCycles = bitCoerce . resize <$> rc
+    writeStart = unpack <$> ws
+    writeCycles = bitCoerce . resize <$> wc
+
+    -- Swap the two words of a 64-bit Bitvector to match the word order of
+    -- the Vexriscv. This allows the CPU to read the two words as one 64-bit value.
+    swapWords :: BitVector 64 -> BitVector 64
+    swapWords = bitCoerce . (swap @(BitVector 32) @(BitVector 32)) . bitCoerce
+
+    rs, rc, ws, wc :: Signal dom (BitVector 64)
+    (rs, rc, ws, wc) = unbundle $ vecToTuple <$> writableRegs
+
+    writableRegs :: Signal dom (Vec 4 (BitVector 64))
+    writableRegs =
+      (fmap (map swapWords . bitCoerce) . bundle . map (regMaybe maxBound) . unbundle)
+        $ take d8
+        <$> writeVec
+
+    (writeVec, wbS2M) = unbundle $ wbToVec <$> bundle readVec <*> wbM2S

From b5e5439dc2f9f5d7757e014a642b9a412ae0606d Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Tue, 11 Feb 2025 11:13:37 +0100
Subject: [PATCH 05/12] Add Rust firmware for switch demo PE

---
 firmware-support/bittide-sys/src/lib.rs       |  1 +
 .../bittide-sys/src/switch_demo_pe.rs         | 76 +++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 firmware-support/bittide-sys/src/switch_demo_pe.rs

diff --git a/firmware-support/bittide-sys/src/lib.rs b/firmware-support/bittide-sys/src/lib.rs
index 8bdc6db2c..35bf93cce 100644
--- a/firmware-support/bittide-sys/src/lib.rs
+++ b/firmware-support/bittide-sys/src/lib.rs
@@ -14,6 +14,7 @@ pub mod mac;
 pub mod program_stream;
 pub mod scatter_unit;
 pub mod smoltcp;
+pub mod switch_demo_pe;
 pub mod time;
 pub mod uart;
 pub mod ugn;
diff --git a/firmware-support/bittide-sys/src/switch_demo_pe.rs b/firmware-support/bittide-sys/src/switch_demo_pe.rs
new file mode 100644
index 000000000..6e264a617
--- /dev/null
+++ b/firmware-support/bittide-sys/src/switch_demo_pe.rs
@@ -0,0 +1,76 @@
+// SPDX-FileCopyrightText: 2025 Google LLC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use ufmt::derive::uDebug;
+
+#[repr(C)]
+#[derive(uDebug, PartialEq, Eq, Copy, Clone)]
+pub struct NodeData {
+    pub local_counter: u64,
+    pub dna: u128,
+}
+
+pub struct SwitchDemoProcessingElement<const BUFFER_SIZE: usize> {
+    base_addr: *const u64,
+}
+
+impl<const BUFFER_SIZE: usize> SwitchDemoProcessingElement<BUFFER_SIZE> {
+    const READ_START: usize = 0;
+    const READ_CYCLES: usize = 1;
+    const WRITE_START: usize = 2;
+    const WRITE_CYCLES: usize = 3;
+    const BUFFER: usize = 4;
+
+    /// Create a new [`SwitchDemoProcessingElement`] instance given a base
+    /// address. The `BUFFER_SIZE` is the number of [`NodeData`] elements in its
+    /// internal buffer.
+    ///
+    /// # Safety
+    ///
+    /// The `base_addr` pointer must be a valid pointer that is backed by
+    /// a memory mapped switch demo processing element. The `BUFFER_SIZE` should
+    /// match the `bufferSize` of the associated `swtichDemoPeWb` device.
+    pub unsafe fn new(base_addr: *const ()) -> Self {
+        let addr = base_addr as *const u64;
+        Self { base_addr: addr }
+    }
+
+    pub fn set_read(&self, read_start: u64, read_cycles: u64) {
+        // SAFETY: This is safe since this function can only be called
+        // after construction, which is only valid with valid addresses.
+        unsafe {
+            self.base_addr
+                .add(Self::READ_START)
+                .cast_mut()
+                .write_volatile(read_start);
+            self.base_addr
+                .add(Self::READ_CYCLES)
+                .cast_mut()
+                .write_volatile(read_cycles);
+        }
+    }
+
+    pub fn set_write(&self, write_start: u64, write_cycles: u64) {
+        // SAFETY: This is safe since this function can only be called
+        // after construction, which is only valid with valid addresses.
+        unsafe {
+            self.base_addr
+                .add(Self::WRITE_START)
+                .cast_mut()
+                .write_volatile(write_start);
+            self.base_addr
+                .add(Self::WRITE_CYCLES)
+                .cast_mut()
+                .write_volatile(write_cycles);
+        }
+    }
+
+    pub fn buffer(&self) -> impl Iterator<Item = NodeData> + '_ {
+        // SAFETY: This is safe since this function can only be called
+        // after construction, which is only valid with valid addresses.
+        (Self::BUFFER..Self::BUFFER + BUFFER_SIZE * 3)
+            .step_by(3)
+            .map(|i| unsafe { self.base_addr.add(i).cast::<NodeData>().read_volatile() })
+    }
+}

From 035d9b244cdb344666515c847535e3987505e717 Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Tue, 11 Feb 2025 11:14:30 +0100
Subject: [PATCH 06/12] Add test for `switchDemoPeWb`

---
 bittide-instances/bittide-instances.cabal     |   1 +
 .../Wishbone/SwitchDemoProcessingElement.hs   | 107 ++++++++++++++++++
 bittide-instances/tests/unittests.hs          |   2 +
 firmware-binaries/Cargo.lock                  |   9 ++
 firmware-binaries/Cargo.toml                  |   1 +
 .../test-cases/switch_demo_pe_test/Cargo.toml |  17 +++
 .../test-cases/switch_demo_pe_test/build.rs   |  23 ++++
 .../test-cases/switch_demo_pe_test/memory.x   |  18 +++
 .../switch_demo_pe_test/src/main.rs           |  92 +++++++++++++++
 .../bittide-sys/src/switch_demo_pe.rs         |  27 +++++
 10 files changed, 297 insertions(+)
 create mode 100644 bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
 create mode 100644 firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml
 create mode 100644 firmware-binaries/test-cases/switch_demo_pe_test/build.rs
 create mode 100644 firmware-binaries/test-cases/switch_demo_pe_test/memory.x
 create mode 100644 firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs

diff --git a/bittide-instances/bittide-instances.cabal b/bittide-instances/bittide-instances.cabal
index f4debd0e5..82288139d 100644
--- a/bittide-instances/bittide-instances.cabal
+++ b/bittide-instances/bittide-instances.cabal
@@ -212,6 +212,7 @@ test-suite unittests
     Wishbone.CaptureUgn
     Wishbone.DnaPortE2
     Wishbone.ScatterGather
+    Wishbone.SwitchDemoProcessingElement
     Wishbone.Time
     Wishbone.Watchdog
 
diff --git a/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
new file mode 100644
index 000000000..6b32944ec
--- /dev/null
+++ b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
@@ -0,0 +1,107 @@
+-- SPDX-FileCopyrightText: 2025 Google LLC
+--
+-- SPDX-License-Identifier: Apache-2.0
+{-# LANGUAGE NumericUnderscores #-}
+{-# OPTIONS_GHC -fplugin=Protocols.Plugin #-}
+
+module Wishbone.SwitchDemoProcessingElement where
+
+import Clash.Explicit.Prelude
+import Clash.Prelude (HiddenClockResetEnable, withClockResetEnable)
+import qualified Prelude as P
+
+import Data.Char (chr)
+import Data.Maybe (mapMaybe)
+import Project.FilePath
+import Protocols
+import Protocols.Idle
+import System.FilePath ((</>))
+import System.IO.Unsafe (unsafePerformIO)
+import Test.Tasty
+import Test.Tasty.HUnit
+import Test.Tasty.TH
+import VexRiscv (DumpVcd (NoDumpVcd))
+
+import Bittide.DoubleBufferedRam
+import Bittide.ProcessingElement
+import Bittide.ProcessingElement.Util
+import Bittide.SharedTypes
+import Bittide.SwitchDemoProcessingElement
+import Bittide.Wishbone
+
+import qualified Protocols.Df as Df
+
+sim :: IO ()
+sim = putStr simResult
+
+simResult :: String
+simResult = chr . fromIntegral <$> mapMaybe Df.dataToMaybe uartStream
+ where
+  uartStream =
+    sampleC def{timeoutAfter = 100_000}
+      $ withClockResetEnable clk reset enable
+      $ dut @System localCounter dnaA dnaB
+
+  clk = clockGen
+  reset = resetGen
+  enable = enableGen
+  localCounter = register clk reset enable 0 (localCounter + 1)
+  dnaA = pure 0xAAAA_0123_4567_89AB_CDEF_0001
+  dnaB = pure 0xBBBB_0123_4567_89AB_CDEF_0001
+
+case_switch_demo_pe_test :: Assertion
+case_switch_demo_pe_test = assertBool msg (receivedString == expectedString)
+ where
+  msg =
+    "Received string "
+      <> receivedString
+      <> " not equal to expected string "
+      <> expectedString
+  expectedString = "Hello world!"
+  receivedString = (P.head . lines) simResult
+
+{- | A simulation-only design containing two `switchDemoPeWb`s connected to a single
+VexRiscV. The VexRiscV runs the `switch_demo_pe_test` binary from `firmware-binaries`.
+-}
+dut ::
+  forall dom.
+  ( HiddenClockResetEnable dom
+  , 1 <= DomainPeriod dom
+  ) =>
+  -- | Local clock cycle counter
+  Signal dom (Unsigned 64) ->
+  -- | Fake DNA (used to identify the different PEs)
+  Signal dom (BitVector 96) ->
+  -- | Fake DNA (used to identify the different PEs)
+  Signal dom (BitVector 96) ->
+  Circuit () (Df dom (BitVector 8))
+dut localCounter dnaA dnaB = circuit $ do
+  (uartRx, jtagIdle) <- idleSource -< ()
+  [uartBus, timeBus, peBusA, peBusB] <- processingElement NoDumpVcd peConfig -< jtagIdle
+  (uartTx, _uartStatus) <- uartInterfaceWb d16 d2 uartSim -< (uartBus, uartRx)
+  timeWb -< timeBus
+  linkAB <- switchDemoPeWb d2 localCounter (Just <$> dnaA) -< (peBusA, linkBA)
+  linkBA <- switchDemoPeWb d2 localCounter (Just <$> dnaB) -< (peBusB, linkAB)
+  idC -< uartTx
+ where
+  memMap = 0b000 :> 0b001 :> 0b010 :> 0b011 :> 0b100 :> 0b101 :> Nil
+  peConfig = unsafePerformIO $ do
+    root <- findParentContaining "cabal.project"
+    let
+      elfDir = root </> firmwareBinariesDir "riscv32imc" Release
+      elfPath = elfDir </> "switch_demo_pe_test"
+    (iMem, dMem) <- vecsFromElf @DMemWords @IMemWords BigEndian elfPath Nothing
+    pure
+      PeConfig
+        { memMapConfig = memMap
+        , initI = Reloadable (Vec iMem)
+        , initD = Reloadable (Vec dMem)
+        , iBusTimeout = d0 -- No timeouts on the instruction bus
+        , dBusTimeout = d0 -- No timeouts on the data bus
+        }
+
+type DMemWords = DivRU (32 * 1024) 4
+type IMemWords = DivRU (32 * 1024) 4
+
+tests :: TestTree
+tests = $(testGroupGenerator)
diff --git a/bittide-instances/tests/unittests.hs b/bittide-instances/tests/unittests.hs
index e2e8065b4..d14beb531 100644
--- a/bittide-instances/tests/unittests.hs
+++ b/bittide-instances/tests/unittests.hs
@@ -14,6 +14,7 @@ import qualified Wishbone.Axi as Axi
 import qualified Wishbone.CaptureUgn as CaptureUgn
 import qualified Wishbone.DnaPortE2 as DnaPortE2
 import qualified Wishbone.ScatterGather as ScatterGather
+import qualified Wishbone.SwitchDemoProcessingElement as SwitchDemoProcessingElement
 import qualified Wishbone.Time as Time
 import qualified Wishbone.Watchdog as Watchdog
 
@@ -24,6 +25,7 @@ tests =
     [ CaptureUgn.tests
     , ClockControlWb.tests
     , ScatterGather.tests
+    , SwitchDemoProcessingElement.tests
     , DnaPortE2.tests
     , Ord.tests
     , Time.tests
diff --git a/firmware-binaries/Cargo.lock b/firmware-binaries/Cargo.lock
index c177fad9e..fb5f47e59 100644
--- a/firmware-binaries/Cargo.lock
+++ b/firmware-binaries/Cargo.lock
@@ -481,6 +481,15 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
 
+[[package]]
+name = "switch_demo_pe_test"
+version = "0.1.0"
+dependencies = [
+ "bittide-sys",
+ "riscv-rt",
+ "ufmt",
+]
+
 [[package]]
 name = "syn"
 version = "1.0.109"
diff --git a/firmware-binaries/Cargo.toml b/firmware-binaries/Cargo.toml
index 88a2af2ab..dab63f5e1 100644
--- a/firmware-binaries/Cargo.toml
+++ b/firmware-binaries/Cargo.toml
@@ -21,6 +21,7 @@ members = [
   "test-cases/clock-control-wb",
   "test-cases/dna_port_e2_test",
   "test-cases/scatter_gather_test",
+  "test-cases/switch_demo_pe_test",
   "test-cases/time_self_test",
   "test-cases/watchdog_test",
 
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml b/firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml
new file mode 100644
index 000000000..ac14d05dd
--- /dev/null
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml
@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: 2024 Google LLC
+#
+# SPDX-License-Identifier: CC0-1.0
+
+[package]
+name = "switch_demo_pe_test"
+version = "0.1.0"
+edition = "2021"
+license = "Apache-2.0"
+authors = ["Google LLC"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+riscv-rt = "0.11.0"
+bittide-sys = { path = "../../../firmware-support/bittide-sys" }
+ufmt = "0.2.0"
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/build.rs b/firmware-binaries/test-cases/switch_demo_pe_test/build.rs
new file mode 100644
index 000000000..cd062a257
--- /dev/null
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/build.rs
@@ -0,0 +1,23 @@
+// SPDX-FileCopyrightText: 2024 Google LLC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use std::env;
+use std::fs;
+use std::path::Path;
+
+/// Put the linker script somewhere the linker can find it.
+fn main() {
+    let out_dir = env::var("OUT_DIR").expect("No out dir");
+    let dest_path = Path::new(&out_dir).join("memory.x");
+    fs::write(dest_path, include_bytes!("memory.x")).expect("Could not write file");
+
+    if env::var("CARGO_CFG_TARGET_ARCH").unwrap() == "riscv32" {
+        println!("cargo:rustc-link-arg=-Tmemory.x");
+        println!("cargo:rustc-link-arg=-Tlink.x"); // linker script from riscv-rt
+    }
+    println!("cargo:rustc-link-search={out_dir}");
+
+    println!("cargo:rerun-if-changed=memory.x");
+    println!("cargo:rerun-if-changed=build.rs");
+}
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/memory.x b/firmware-binaries/test-cases/switch_demo_pe_test/memory.x
new file mode 100644
index 000000000..34a0597f8
--- /dev/null
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/memory.x
@@ -0,0 +1,18 @@
+/*
+SPDX-FileCopyrightText: 2024 Google LLC
+
+SPDX-License-Identifier: CC0-1.0
+*/
+
+MEMORY
+{
+  IMEM : ORIGIN = 0x80000000, LENGTH = 32K
+  DMEM : ORIGIN = 0x20000000, LENGTH = 32K
+}
+
+REGION_ALIAS("REGION_TEXT", IMEM);
+REGION_ALIAS("REGION_RODATA", DMEM);
+REGION_ALIAS("REGION_DATA", DMEM);
+REGION_ALIAS("REGION_BSS", DMEM);
+REGION_ALIAS("REGION_HEAP", DMEM);
+REGION_ALIAS("REGION_STACK", DMEM);
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
new file mode 100644
index 000000000..8414ce0cd
--- /dev/null
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
@@ -0,0 +1,92 @@
+// SPDX-FileCopyrightText: 2025 Google LLC
+//
+// SPDX-License-Identifier: Apache-2.0
+#![no_std]
+#![cfg_attr(not(test), no_main)]
+
+use bittide_sys::switch_demo_pe::SwitchDemoProcessingElement;
+use bittide_sys::time::{Clock, Duration};
+use bittide_sys::uart::Uart;
+use core::fmt::Write;
+#[cfg(not(test))]
+use riscv_rt::entry;
+
+const UART_ADDR: *const () = (0b010 << 29) as *const ();
+const CLOCK_ADDR: *const () = (0b011 << 29) as *const ();
+const SWITCH_PE_A: *const () = (0b100 << 29) as *const ();
+const SWITCH_PE_B: *const () = (0b101 << 29) as *const ();
+
+// Size of buffer in number of "tri-cycles". That is, we always store 3 64-bit words:
+// local clock cycle counter, DNA (64 lsbs), DNA (32 msbs, zero-extended).
+// Should match `bufferSize` of the associated `switchDemoPeWb` device.
+const BUFFER_SIZE: usize = 2;
+
+#[cfg_attr(not(test), entry)]
+fn main() -> ! {
+    // Initialize peripherals.
+    let mut uart = unsafe { Uart::new(UART_ADDR) };
+    let mut clock = unsafe { Clock::new(CLOCK_ADDR) };
+    let switch_pe_a: SwitchDemoProcessingElement<BUFFER_SIZE> =
+        unsafe { SwitchDemoProcessingElement::new(SWITCH_PE_A) };
+    let switch_pe_b: SwitchDemoProcessingElement<BUFFER_SIZE> =
+        unsafe { SwitchDemoProcessingElement::new(SWITCH_PE_B) };
+
+    let first_transfer_start = 0x4000;
+    let second_transfer_start = 0x4100;
+
+    // A only writes its own data
+    switch_pe_a.set_write(first_transfer_start, 1);
+    // B reads data from A
+    switch_pe_b.set_read(first_transfer_start, 1);
+    // B writes its own data and data received from A
+    switch_pe_b.set_write(second_transfer_start, 2);
+    // A reads all data from B
+    switch_pe_a.set_read(second_transfer_start, 2);
+
+    clock.wait(Duration::from_micros(200));
+
+    // Write the buffer of A over UART
+    write!(uart, "Buffer A: [").unwrap();
+    switch_pe_a.buffer_u64().enumerate().for_each(|(i, nd)| {
+        let sep = if i + 1 < BUFFER_SIZE * 3 { ", " } else { "" };
+        write!(uart, "0x{:X}{sep}", nd).unwrap();
+    });
+    writeln!(uart, "]").unwrap();
+
+    // Write the buffer of B over UART
+    write!(uart, "Buffer B: [").unwrap();
+    switch_pe_b.buffer_u64().enumerate().for_each(|(i, nd)| {
+        let sep = if i + 1 < BUFFER_SIZE * 3 { ", " } else { "" };
+        write!(uart, "0x{:X}{sep}", nd).unwrap();
+    });
+    writeln!(uart, "]").unwrap();
+
+    // Write the buffer of A over UART
+    write!(uart, "Buffer A: [").unwrap();
+    switch_pe_a.buffer().enumerate().for_each(|(i, nd)| {
+        let sep = if i + 1 < BUFFER_SIZE { ", " } else { "" };
+        write!(uart, "(0x{:X}, 0x{:X}){sep}", nd.local_counter, nd.dna).unwrap();
+    });
+    writeln!(uart, "]").unwrap();
+
+    // Write the buffer of B over UART
+    write!(uart, "Buffer B: [").unwrap();
+    switch_pe_b.buffer().enumerate().for_each(|(i, nd)| {
+        let sep = if i + 1 < BUFFER_SIZE { ", " } else { "" };
+        write!(uart, "(0x{:X}, 0x{:X}){sep}", nd.local_counter, nd.dna).unwrap();
+    });
+    writeln!(uart, "]").unwrap();
+
+    loop {
+        continue;
+    }
+}
+
+#[panic_handler]
+fn panic_handler(info: &core::panic::PanicInfo) -> ! {
+    let mut uart = unsafe { Uart::new(UART_ADDR) };
+    writeln!(uart, "Panicked! #{info}").unwrap();
+    loop {
+        continue;
+    }
+}
diff --git a/firmware-support/bittide-sys/src/switch_demo_pe.rs b/firmware-support/bittide-sys/src/switch_demo_pe.rs
index 6e264a617..e4d65c412 100644
--- a/firmware-support/bittide-sys/src/switch_demo_pe.rs
+++ b/firmware-support/bittide-sys/src/switch_demo_pe.rs
@@ -66,6 +66,26 @@ impl<const BUFFER_SIZE: usize> SwitchDemoProcessingElement<BUFFER_SIZE> {
         }
     }
 
+    pub fn get_read(&self) -> (u64, u64) {
+        // SAFETY: This is safe since this function can only be called
+        // after construction, which is only valid with valid addresses.
+        unsafe {
+            let read_start = self.base_addr.add(Self::READ_START).read_volatile();
+            let read_cycles = self.base_addr.add(Self::READ_CYCLES).read_volatile();
+            (read_start, read_cycles)
+        }
+    }
+
+    pub fn get_write(&self) -> (u64, u64) {
+        // SAFETY: This is safe since this function can only be called
+        // after construction, which is only valid with valid addresses.
+        unsafe {
+            let write_start = self.base_addr.add(Self::WRITE_START).read_volatile();
+            let write_cycles = self.base_addr.add(Self::WRITE_CYCLES).read_volatile();
+            (write_start, write_cycles)
+        }
+    }
+
     pub fn buffer(&self) -> impl Iterator<Item = NodeData> + '_ {
         // SAFETY: This is safe since this function can only be called
         // after construction, which is only valid with valid addresses.
@@ -73,4 +93,11 @@ impl<const BUFFER_SIZE: usize> SwitchDemoProcessingElement<BUFFER_SIZE> {
             .step_by(3)
             .map(|i| unsafe { self.base_addr.add(i).cast::<NodeData>().read_volatile() })
     }
+
+    pub fn buffer_u64(&self) -> impl Iterator<Item = u64> + '_ {
+        // SAFETY: This is safe since this function can only be called
+        // after construction, which is only valid with valid addresses.
+        (Self::BUFFER..Self::BUFFER + (BUFFER_SIZE * 3))
+            .map(|i| unsafe { self.base_addr.add(i).read_volatile() })
+    }
 }

From 50c0b1573d832a9d2eaed98abb8729bca01aaf2a Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Mon, 17 Feb 2025 12:02:31 +0100
Subject: [PATCH 07/12] Increment size `readCycles` by one to match user input

The incremented size of `readCycles` better matches the intended number
of tri-cycles the user writes in Rust, since this makes no sense to be
zero-indexed.
---
 .../src/Bittide/SwitchDemoProcessingElement.hs | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/bittide/src/Bittide/SwitchDemoProcessingElement.hs b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
index a161a10a1..f23354211 100644
--- a/bittide/src/Bittide/SwitchDemoProcessingElement.hs
+++ b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
@@ -44,7 +44,7 @@ switchDemoPe ::
   -- | When to read from the crossbar link
   Signal dom (Unsigned 64) ->
   -- | How many tri-cycles to read from the crossbar link
-  Signal dom (Index bufferSize) ->
+  Signal dom (Index (bufferSize + 1)) ->
   -- | When to write to the crossbar link
   Signal dom (Unsigned 64) ->
   -- | How many tri-cycles to write to the crossbar link. Includes writing \"own\" data.
@@ -57,7 +57,7 @@ switchDemoPe ::
 switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart writeCycles =
   (linkOut, buffer)
  where
-  readCyclesExtended = zeroExtendTimesThree <$> readCycles
+  readCyclesExtended = checkedResize . zeroExtendTimesThree <$> readCycles
   writeCyclesExtended = zeroExtendTimesThree <$> writeCycles
 
   localData :: Signal dom (Vec 3 (BitVector 64))
@@ -94,7 +94,7 @@ switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart w
     enableVec = unbundle $ go <$> peState
      where
       go :: (HasCallStack) => SimplePeState bufferSize -> Vec (bufferSize * 3) Bool
-      go (Read x) = (== x) <$> indicesI
+      go (Read x) = (== checkedResize x) <$> indicesI
       go _ = repeat False
 
   prevPeState = register Idle peState
@@ -114,7 +114,7 @@ switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart w
       -- \| When to read from the crossbar link
       Unsigned 64 ->
       -- \| How many cycles to read from the crossbar link
-      Index (bufferSize * 3) ->
+      Index (bufferSize * 3 + 1) ->
       -- \| When to write to the crossbar link
       Unsigned 64 ->
       -- \| How many cycles to write to the crossbar link
@@ -138,7 +138,7 @@ switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart w
 
 data SimplePeState bufferSize
   = Idle
-  | Read (Index (bufferSize * 3))
+  | Read (Index (bufferSize * 3 + 1))
   | Write (Index ((bufferSize + 1) * 3))
   deriving (Generic, NFDataX, Eq, Show)
 
@@ -192,9 +192,13 @@ switchDemoPeWb SNat localCounter maybeDna = Circuit go
         writeCycles
 
     readStart = unpack <$> rs
-    readCycles = bitCoerce . resize <$> rc
+    readCycles = checkedResize . bvToIndex <$> rc
     writeStart = unpack <$> ws
-    writeCycles = bitCoerce . resize <$> wc
+    writeCycles = checkedResize . bvToIndex <$> wc
+
+    -- \| Unpack a BitVector to an Index of the same size
+    bvToIndex :: (KnownNat n) => BitVector n -> Index (2 ^ n)
+    bvToIndex = unpack
 
     -- Swap the two words of a 64-bit Bitvector to match the word order of
     -- the Vexriscv. This allows the CPU to read the two words as one 64-bit value.

From 3a1afdb8668516310021738866ae4b9760c99561 Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Mon, 17 Feb 2025 14:42:12 +0100
Subject: [PATCH 08/12] Change word order of DNA to match C representation

---
 .../Wishbone/SwitchDemoProcessingElement.hs   | 17 ++++++---
 .../Bittide/SwitchDemoProcessingElement.hs    |  6 ++--
 .../Tests/SwitchDemoProcessingElement.hs      |  2 +-
 .../switch_demo_pe_test/src/main.rs           | 35 +++++++++----------
 .../bittide-sys/src/switch_demo_pe.rs         | 11 ++----
 5 files changed, 37 insertions(+), 34 deletions(-)

diff --git a/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
index 6b32944ec..256992b3f 100644
--- a/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
+++ b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
@@ -8,7 +8,6 @@ module Wishbone.SwitchDemoProcessingElement where
 
 import Clash.Explicit.Prelude
 import Clash.Prelude (HiddenClockResetEnable, withClockResetEnable)
-import qualified Prelude as P
 
 import Data.Char (chr)
 import Data.Maybe (mapMaybe)
@@ -31,12 +30,17 @@ import Bittide.Wishbone
 
 import qualified Protocols.Df as Df
 
+takeWhileInclusive :: (a -> Bool) -> [a] -> [a]
+takeWhileInclusive _ [] = []
+takeWhileInclusive p (x : xs) = x : if p x then takeWhileInclusive p xs else []
+
 sim :: IO ()
 sim = putStr simResult
 
 simResult :: String
-simResult = chr . fromIntegral <$> mapMaybe Df.dataToMaybe uartStream
+simResult = unlines . takeWhileInclusive (/= "Finished") . lines $ uartString
  where
+  uartString = chr . fromIntegral <$> mapMaybe Df.dataToMaybe uartStream
   uartStream =
     sampleC def{timeoutAfter = 100_000}
       $ withClockResetEnable clk reset enable
@@ -57,8 +61,13 @@ case_switch_demo_pe_test = assertBool msg (receivedString == expectedString)
       <> receivedString
       <> " not equal to expected string "
       <> expectedString
-  expectedString = "Hello world!"
-  receivedString = (P.head . lines) simResult
+  receivedString = simResult
+  expectedString =
+    unlines
+      [ "Buffer A: [(0x4100, 0xBBBB0123456789ABCDEF0001), (0x4000, 0xAAAA0123456789ABCDEF0001)]"
+      , "Buffer B: [(0x4000, 0xAAAA0123456789ABCDEF0001), (0xABBAABBAABBA0003, 0xABBA0005ABBAABBAABBA0004)]"
+      , "Finished"
+      ]
 
 {- | A simulation-only design containing two `switchDemoPeWb`s connected to a single
 VexRiscV. The VexRiscV runs the `switch_demo_pe_test` binary from `firmware-binaries`.
diff --git a/bittide/src/Bittide/SwitchDemoProcessingElement.hs b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
index f23354211..5743cd244 100644
--- a/bittide/src/Bittide/SwitchDemoProcessingElement.hs
+++ b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
@@ -33,7 +33,7 @@ switchDemoPe ::
   , 1 <= bufferSize
   ) =>
   -- | Size of buffer in number of "tri-cycles". That is, we always store 3 64-bit words:
-  -- DNA (32 msbs), DNA (64 lsbs), local clock cycle counter.
+  -- local clock cycle counter, DNA (64 lsbs), DNA (32 msbs, zero-extended).
   SNat bufferSize ->
   -- | Local clock cycle counter
   Signal dom (Unsigned 64) ->
@@ -64,7 +64,7 @@ switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart w
   localData = bundle ((pack <$> localCounter) :> unbundle dnaVec)
    where
     dnaVec :: Signal dom (Vec 2 (BitVector 64))
-    dnaVec = bitCoerce . zeroExtend <$> dnaLocked
+    dnaVec = reverse . bitCoerce . zeroExtend <$> dnaLocked
     dnaLocked = fromMaybe 0xBAAB_BAAB_BAAB_BAAB_BAAB_BAAB <$> maybeDna
 
   linkOut = stateToLinkOutput <$> peState <*> buffer <*> localData
@@ -200,7 +200,7 @@ switchDemoPeWb SNat localCounter maybeDna = Circuit go
     bvToIndex :: (KnownNat n) => BitVector n -> Index (2 ^ n)
     bvToIndex = unpack
 
-    -- Swap the two words of a 64-bit Bitvector to match the word order of
+    -- \| Swap the two words of a 64-bit Bitvector to match the word order of
     -- the Vexriscv. This allows the CPU to read the two words as one 64-bit value.
     swapWords :: BitVector 64 -> BitVector 64
     swapWords = bitCoerce . (swap @(BitVector 32) @(BitVector 32)) . bitCoerce
diff --git a/bittide/tests/Tests/SwitchDemoProcessingElement.hs b/bittide/tests/Tests/SwitchDemoProcessingElement.hs
index 27f4306c5..a73eaa4f3 100644
--- a/bittide/tests/Tests/SwitchDemoProcessingElement.hs
+++ b/bittide/tests/Tests/SwitchDemoProcessingElement.hs
@@ -179,6 +179,6 @@ prop_readThenWrite = H.property $ do
       -- Note we can always write one tri-cycle more than we read, since internal
       -- data comes first.
       relevantOutCycles = fromIntegral (min nWriteCycles (nReadCycles + cyclesPerReadWrite))
-      deviceDnaVec = bitCoerce @_ @(Vec 2 (BitVector 64)) (zeroExtend deviceDna)
+      deviceDnaVec = reverse $ bitCoerce @_ @(Vec 2 (BitVector 64)) (zeroExtend deviceDna)
       expectedOutData = toList (pack writeStart :> deviceDnaVec) <> readData
     L.take relevantOutCycles writeOuts === L.take relevantOutCycles expectedOutData
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
index 8414ce0cd..f98e79fc0 100644
--- a/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
@@ -4,6 +4,7 @@
 #![no_std]
 #![cfg_attr(not(test), no_main)]
 
+use bittide_sys::dna_port_e2::dna_to_u128;
 use bittide_sys::switch_demo_pe::SwitchDemoProcessingElement;
 use bittide_sys::time::{Clock, Duration};
 use bittide_sys::uart::Uart;
@@ -45,27 +46,17 @@ fn main() -> ! {
 
     clock.wait(Duration::from_micros(200));
 
-    // Write the buffer of A over UART
-    write!(uart, "Buffer A: [").unwrap();
-    switch_pe_a.buffer_u64().enumerate().for_each(|(i, nd)| {
-        let sep = if i + 1 < BUFFER_SIZE * 3 { ", " } else { "" };
-        write!(uart, "0x{:X}{sep}", nd).unwrap();
-    });
-    writeln!(uart, "]").unwrap();
-
-    // Write the buffer of B over UART
-    write!(uart, "Buffer B: [").unwrap();
-    switch_pe_b.buffer_u64().enumerate().for_each(|(i, nd)| {
-        let sep = if i + 1 < BUFFER_SIZE * 3 { ", " } else { "" };
-        write!(uart, "0x{:X}{sep}", nd).unwrap();
-    });
-    writeln!(uart, "]").unwrap();
-
     // Write the buffer of A over UART
     write!(uart, "Buffer A: [").unwrap();
     switch_pe_a.buffer().enumerate().for_each(|(i, nd)| {
         let sep = if i + 1 < BUFFER_SIZE { ", " } else { "" };
-        write!(uart, "(0x{:X}, 0x{:X}){sep}", nd.local_counter, nd.dna).unwrap();
+        write!(
+            uart,
+            "(0x{:X}, 0x{:X}){sep}",
+            nd.local_counter,
+            dna_to_u128(nd.dna)
+        )
+        .unwrap();
     });
     writeln!(uart, "]").unwrap();
 
@@ -73,10 +64,18 @@ fn main() -> ! {
     write!(uart, "Buffer B: [").unwrap();
     switch_pe_b.buffer().enumerate().for_each(|(i, nd)| {
         let sep = if i + 1 < BUFFER_SIZE { ", " } else { "" };
-        write!(uart, "(0x{:X}, 0x{:X}){sep}", nd.local_counter, nd.dna).unwrap();
+        write!(
+            uart,
+            "(0x{:X}, 0x{:X}){sep}",
+            nd.local_counter,
+            dna_to_u128(nd.dna)
+        )
+        .unwrap();
     });
     writeln!(uart, "]").unwrap();
 
+    writeln!(uart, "Finished").unwrap();
+
     loop {
         continue;
     }
diff --git a/firmware-support/bittide-sys/src/switch_demo_pe.rs b/firmware-support/bittide-sys/src/switch_demo_pe.rs
index e4d65c412..eea4f4b76 100644
--- a/firmware-support/bittide-sys/src/switch_demo_pe.rs
+++ b/firmware-support/bittide-sys/src/switch_demo_pe.rs
@@ -4,11 +4,13 @@
 
 use ufmt::derive::uDebug;
 
+use crate::dna_port_e2::DnaValue;
+
 #[repr(C)]
 #[derive(uDebug, PartialEq, Eq, Copy, Clone)]
 pub struct NodeData {
     pub local_counter: u64,
-    pub dna: u128,
+    pub dna: DnaValue,
 }
 
 pub struct SwitchDemoProcessingElement<const BUFFER_SIZE: usize> {
@@ -93,11 +95,4 @@ impl<const BUFFER_SIZE: usize> SwitchDemoProcessingElement<BUFFER_SIZE> {
             .step_by(3)
             .map(|i| unsafe { self.base_addr.add(i).cast::<NodeData>().read_volatile() })
     }
-
-    pub fn buffer_u64(&self) -> impl Iterator<Item = u64> + '_ {
-        // SAFETY: This is safe since this function can only be called
-        // after construction, which is only valid with valid addresses.
-        (Self::BUFFER..Self::BUFFER + (BUFFER_SIZE * 3))
-            .map(|i| unsafe { self.base_addr.add(i).read_volatile() })
-    }
 }

From f11d57e709e14a85b42d4e9a93a18ae38076ad1c Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Tue, 18 Feb 2025 10:25:56 +0100
Subject: [PATCH 09/12] Make local counter wishbone accessible for debugging

---
 .../Bittide/SwitchDemoProcessingElement.hs    | 14 ++++++-----
 .../switch_demo_pe_test/src/main.rs           | 24 +++++++++++++++++++
 .../bittide-sys/src/switch_demo_pe.rs         |  9 ++++++-
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/bittide/src/Bittide/SwitchDemoProcessingElement.hs b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
index 5743cd244..dca737849 100644
--- a/bittide/src/Bittide/SwitchDemoProcessingElement.hs
+++ b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
@@ -147,11 +147,12 @@ data SimplePeState bufferSize
 Buffer uses 64-bit words internally, but WB interface is 32-bit.
 
 The register layout is as follows (lsbs in first 32-bit word, msbs in second):
-- Address 0-1: read start
-- Address 2-3: read cycles
-- Address 4-5: write start
-- Address 6-7: write cycles
-- Address 8-.: buffer (bufferSize*3*2)
+- Address  0 - 1: read start
+- Address  2 - 3: read cycles
+- Address  4 - 5: write start
+- Address  6 - 7: write cycles
+- Address  8 - 9: local clock cycle counter
+- Address 10 - .: buffer (bufferSize*3*2)
 -}
 switchDemoPeWb ::
   forall bufferSize dom addrW.
@@ -173,10 +174,11 @@ switchDemoPeWb SNat localCounter maybeDna = Circuit go
  where
   go ((wbM2S, linkIn), _) = ((wbS2M, pure ()), linkOut)
    where
-    readVec :: Vec (8 + bufferSize * 3 * 2) (Signal dom (BitVector 32))
+    readVec :: Vec (8 + bufferSize * 3 * 2 + 2) (Signal dom (BitVector 32))
     readVec =
       dflipflop
         <$> ( unbundle (bitCoerce . map swapWords <$> writableRegs)
+                ++ unbundle (bitCoerce . map swapWords . bitCoerce <$> localCounter)
                 ++ unbundle (bitCoerce . map swapWords <$> buffer)
             )
 
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
index f98e79fc0..9d5284ca4 100644
--- a/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
@@ -8,7 +8,10 @@ use bittide_sys::dna_port_e2::dna_to_u128;
 use bittide_sys::switch_demo_pe::SwitchDemoProcessingElement;
 use bittide_sys::time::{Clock, Duration};
 use bittide_sys::uart::Uart;
+
 use core::fmt::Write;
+use ufmt::uwriteln;
+
 #[cfg(not(test))]
 use riscv_rt::entry;
 
@@ -21,6 +24,7 @@ const SWITCH_PE_B: *const () = (0b101 << 29) as *const ();
 // local clock cycle counter, DNA (64 lsbs), DNA (32 msbs, zero-extended).
 // Should match `bufferSize` of the associated `switchDemoPeWb` device.
 const BUFFER_SIZE: usize = 2;
+const DEBUG_MODE: bool = false;
 
 #[cfg_attr(not(test), entry)]
 fn main() -> ! {
@@ -32,6 +36,10 @@ fn main() -> ! {
     let switch_pe_b: SwitchDemoProcessingElement<BUFFER_SIZE> =
         unsafe { SwitchDemoProcessingElement::new(SWITCH_PE_B) };
 
+    if DEBUG_MODE {
+        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
+    }
+
     let first_transfer_start = 0x4000;
     let second_transfer_start = 0x4100;
 
@@ -46,6 +54,14 @@ fn main() -> ! {
 
     clock.wait(Duration::from_micros(200));
 
+    if DEBUG_MODE {
+        let (rs_a, rc_a) = switch_pe_a.get_read();
+        let (rs_b, rc_b) = switch_pe_b.get_read();
+        uwriteln!(uart, "A: readStart: 0x{:X}, readCycles: 0x{:X}", rs_a, rc_a).unwrap();
+        uwriteln!(uart, "B: readStart: 0x{:X}, readCycles: 0x{:X}", rs_b, rc_b).unwrap();
+        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
+    }
+
     // Write the buffer of A over UART
     write!(uart, "Buffer A: [").unwrap();
     switch_pe_a.buffer().enumerate().for_each(|(i, nd)| {
@@ -60,6 +76,10 @@ fn main() -> ! {
     });
     writeln!(uart, "]").unwrap();
 
+    if DEBUG_MODE {
+        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
+    }
+
     // Write the buffer of B over UART
     write!(uart, "Buffer B: [").unwrap();
     switch_pe_b.buffer().enumerate().for_each(|(i, nd)| {
@@ -74,6 +94,10 @@ fn main() -> ! {
     });
     writeln!(uart, "]").unwrap();
 
+    if DEBUG_MODE {
+        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
+    }
+
     writeln!(uart, "Finished").unwrap();
 
     loop {
diff --git a/firmware-support/bittide-sys/src/switch_demo_pe.rs b/firmware-support/bittide-sys/src/switch_demo_pe.rs
index eea4f4b76..f32dfad32 100644
--- a/firmware-support/bittide-sys/src/switch_demo_pe.rs
+++ b/firmware-support/bittide-sys/src/switch_demo_pe.rs
@@ -22,7 +22,8 @@ impl<const BUFFER_SIZE: usize> SwitchDemoProcessingElement<BUFFER_SIZE> {
     const READ_CYCLES: usize = 1;
     const WRITE_START: usize = 2;
     const WRITE_CYCLES: usize = 3;
-    const BUFFER: usize = 4;
+    const COUNTER: usize = 4;
+    const BUFFER: usize = 5;
 
     /// Create a new [`SwitchDemoProcessingElement`] instance given a base
     /// address. The `BUFFER_SIZE` is the number of [`NodeData`] elements in its
@@ -88,6 +89,12 @@ impl<const BUFFER_SIZE: usize> SwitchDemoProcessingElement<BUFFER_SIZE> {
         }
     }
 
+    pub fn get_counter(&self) -> u64 {
+        // SAFETY: This is safe since this function can only be called
+        // after construction, which is only valid with valid addresses.
+        unsafe { self.base_addr.add(Self::COUNTER).read_volatile() }
+    }
+
     pub fn buffer(&self) -> impl Iterator<Item = NodeData> + '_ {
         // SAFETY: This is safe since this function can only be called
         // after construction, which is only valid with valid addresses.

From d94b8536cdb1ea911a3ef8e4536ab2a33895270d Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Wed, 19 Feb 2025 11:50:14 +0100
Subject: [PATCH 10/12] Increase logging level for debug build

Having a lower max debug level for a debug build than a release build makes no
sense. The maximum level for a debug build is now 'trace' (the highest) and for
a release build it is 'info'.
---
 firmware-binaries/examples/smoltcp_client/Cargo.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firmware-binaries/examples/smoltcp_client/Cargo.toml b/firmware-binaries/examples/smoltcp_client/Cargo.toml
index d658ed34c..a4f881e38 100644
--- a/firmware-binaries/examples/smoltcp_client/Cargo.toml
+++ b/firmware-binaries/examples/smoltcp_client/Cargo.toml
@@ -28,7 +28,7 @@ default-features = false
 
 [dependencies.log]
 version = "0.4.21"
-features = ["max_level_off", "release_max_level_info"]
+features = ["max_level_trace", "release_max_level_info"]
 
 [dependencies.smoltcp]
 git = "https://github.com/smoltcp-rs/smoltcp.git"

From 9a3376fd688fbfebf2ad25b3b5acb27d727a4824 Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Wed, 19 Feb 2025 12:05:40 +0100
Subject: [PATCH 11/12] Use `log` instead of manual debugging option

The 'features' of the log module set in the respective Cargo.toml are global.
Adding other features in two Cargo.toml files in the same module gives a
duplicate features error.

Using the log module is slower than the previous method of a global constant
boolean `DEBUG_MODE` with if-statements around `uwrite` calls. Therefore, I had
to increase the values for when to start reading/writing in the switch demo PE
test.
---
 .../Wishbone/SwitchDemoProcessingElement.hs   | 10 +++--
 firmware-binaries/Cargo.lock                  |  1 +
 .../test-cases/switch_demo_pe_test/Cargo.toml |  6 ++-
 .../switch_demo_pe_test/src/main.rs           | 45 ++++++++++---------
 4 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
index 256992b3f..03cb1a151 100644
--- a/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
+++ b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
@@ -10,6 +10,7 @@ import Clash.Explicit.Prelude
 import Clash.Prelude (HiddenClockResetEnable, withClockResetEnable)
 
 import Data.Char (chr)
+import Data.List (isPrefixOf)
 import Data.Maybe (mapMaybe)
 import Project.FilePath
 import Protocols
@@ -42,7 +43,7 @@ simResult = unlines . takeWhileInclusive (/= "Finished") . lines $ uartString
  where
   uartString = chr . fromIntegral <$> mapMaybe Df.dataToMaybe uartStream
   uartStream =
-    sampleC def{timeoutAfter = 100_000}
+    sampleC def{timeoutAfter = 200_000}
       $ withClockResetEnable clk reset enable
       $ dut @System localCounter dnaA dnaB
 
@@ -61,11 +62,12 @@ case_switch_demo_pe_test = assertBool msg (receivedString == expectedString)
       <> receivedString
       <> " not equal to expected string "
       <> expectedString
-  receivedString = simResult
+  -- Filter the 'debugging' prints, which are prefixed with 'INFO'
+  receivedString = unlines . filter (not . isPrefixOf "INFO") . lines $ simResult
   expectedString =
     unlines
-      [ "Buffer A: [(0x4100, 0xBBBB0123456789ABCDEF0001), (0x4000, 0xAAAA0123456789ABCDEF0001)]"
-      , "Buffer B: [(0x4000, 0xAAAA0123456789ABCDEF0001), (0xABBAABBAABBA0003, 0xABBA0005ABBAABBAABBA0004)]"
+      [ "Buffer A: [(0x10100, 0xBBBB0123456789ABCDEF0001), (0x10000, 0xAAAA0123456789ABCDEF0001)]"
+      , "Buffer B: [(0x10000, 0xAAAA0123456789ABCDEF0001), (0xABBAABBAABBA0003, 0xABBA0005ABBAABBAABBA0004)]"
       , "Finished"
       ]
 
diff --git a/firmware-binaries/Cargo.lock b/firmware-binaries/Cargo.lock
index fb5f47e59..dec71424f 100644
--- a/firmware-binaries/Cargo.lock
+++ b/firmware-binaries/Cargo.lock
@@ -486,6 +486,7 @@ name = "switch_demo_pe_test"
 version = "0.1.0"
 dependencies = [
  "bittide-sys",
+ "log",
  "riscv-rt",
  "ufmt",
 ]
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml b/firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml
index ac14d05dd..6ae3da80c 100644
--- a/firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/Cargo.toml
@@ -12,6 +12,10 @@ authors = ["Google LLC"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-riscv-rt = "0.11.0"
 bittide-sys = { path = "../../../firmware-support/bittide-sys" }
+riscv-rt = "0.11.0"
 ufmt = "0.2.0"
+
+[dependencies.log]
+version = "0.4.21"
+features = ["max_level_trace", "release_max_level_info"]
diff --git a/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
index 9d5284ca4..60faf33e3 100644
--- a/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
+++ b/firmware-binaries/test-cases/switch_demo_pe_test/src/main.rs
@@ -7,10 +7,11 @@
 use bittide_sys::dna_port_e2::dna_to_u128;
 use bittide_sys::switch_demo_pe::SwitchDemoProcessingElement;
 use bittide_sys::time::{Clock, Duration};
+use bittide_sys::uart::log::LOGGER;
 use bittide_sys::uart::Uart;
 
 use core::fmt::Write;
-use ufmt::uwriteln;
+use log::{info, LevelFilter};
 
 #[cfg(not(test))]
 use riscv_rt::entry;
@@ -24,8 +25,9 @@ const SWITCH_PE_B: *const () = (0b101 << 29) as *const ();
 // local clock cycle counter, DNA (64 lsbs), DNA (32 msbs, zero-extended).
 // Should match `bufferSize` of the associated `switchDemoPeWb` device.
 const BUFFER_SIZE: usize = 2;
-const DEBUG_MODE: bool = false;
 
+// See https://github.com/bittide/bittide-hardware/issues/681
+#[allow(static_mut_refs)]
 #[cfg_attr(not(test), entry)]
 fn main() -> ! {
     // Initialize peripherals.
@@ -36,12 +38,21 @@ fn main() -> ! {
     let switch_pe_b: SwitchDemoProcessingElement<BUFFER_SIZE> =
         unsafe { SwitchDemoProcessingElement::new(SWITCH_PE_B) };
 
-    if DEBUG_MODE {
-        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
+    unsafe {
+        LOGGER.set_logger(uart.clone());
+        LOGGER.set_clock(clock.clone());
+        LOGGER.display_source = LevelFilter::Info;
+        log::set_logger_racy(&LOGGER).ok();
+        // The 'max_level' is actually the current debug level. Note that the
+        // unittest uses a release build, which has 'max_level_info', which sets
+        // the actual maximum level.
+        log::set_max_level_racy(LevelFilter::Info);
     }
 
-    let first_transfer_start = 0x4000;
-    let second_transfer_start = 0x4100;
+    info!("Local counter: 0x{:X}", switch_pe_a.get_counter());
+
+    let first_transfer_start = 0x10000;
+    let second_transfer_start = 0x10100;
 
     // A only writes its own data
     switch_pe_a.set_write(first_transfer_start, 1);
@@ -52,15 +63,13 @@ fn main() -> ! {
     // A reads all data from B
     switch_pe_a.set_read(second_transfer_start, 2);
 
-    clock.wait(Duration::from_micros(200));
+    clock.wait(Duration::from_micros(600));
 
-    if DEBUG_MODE {
-        let (rs_a, rc_a) = switch_pe_a.get_read();
-        let (rs_b, rc_b) = switch_pe_b.get_read();
-        uwriteln!(uart, "A: readStart: 0x{:X}, readCycles: 0x{:X}", rs_a, rc_a).unwrap();
-        uwriteln!(uart, "B: readStart: 0x{:X}, readCycles: 0x{:X}", rs_b, rc_b).unwrap();
-        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
-    }
+    let (rs_a, rc_a) = switch_pe_a.get_read();
+    let (rs_b, rc_b) = switch_pe_b.get_read();
+    info!("A: readStart: 0x{:X}, readCycles: 0x{:X}", rs_a, rc_a);
+    info!("B: readStart: 0x{:X}, readCycles: 0x{:X}", rs_b, rc_b);
+    info!("Local counter: 0x{:X}", switch_pe_a.get_counter());
 
     // Write the buffer of A over UART
     write!(uart, "Buffer A: [").unwrap();
@@ -76,9 +85,7 @@ fn main() -> ! {
     });
     writeln!(uart, "]").unwrap();
 
-    if DEBUG_MODE {
-        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
-    }
+    info!("Local counter: 0x{:X}", switch_pe_a.get_counter());
 
     // Write the buffer of B over UART
     write!(uart, "Buffer B: [").unwrap();
@@ -94,9 +101,7 @@ fn main() -> ! {
     });
     writeln!(uart, "]").unwrap();
 
-    if DEBUG_MODE {
-        uwriteln!(uart, "Local counter: 0x{:X}", switch_pe_a.get_counter()).unwrap();
-    }
+    info!("Local counter: 0x{:X}", switch_pe_a.get_counter());
 
     writeln!(uart, "Finished").unwrap();
 

From 7165f891767d1debe4866456d87a9ba724c1c09d Mon Sep 17 00:00:00 2001
From: Hidde Moll <hidde@qbaylogic.com>
Date: Wed, 19 Feb 2025 14:12:38 +0100
Subject: [PATCH 12/12] Make `readDnaPortE2Wb` also return the DNA

By also returning the DNA, instead of only having it wishbone accessible, we can
also use the device DNA in hardware.
---
 .../src/Bittide/Instances/Pnr/Ethernet.hs     |  4 ++--
 bittide-instances/tests/Wishbone/DnaPortE2.hs |  2 +-
 .../Wishbone/SwitchDemoProcessingElement.hs   |  9 ++++++--
 .../Bittide/SwitchDemoProcessingElement.hs    | 22 +++++++++----------
 bittide/src/Bittide/Wishbone.hs               |  6 ++---
 .../Tests/SwitchDemoProcessingElement.hs      |  2 +-
 6 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/bittide-instances/src/Bittide/Instances/Pnr/Ethernet.hs b/bittide-instances/src/Bittide/Instances/Pnr/Ethernet.hs
index 03dd12aa2..81ac92223 100644
--- a/bittide-instances/src/Bittide/Instances/Pnr/Ethernet.hs
+++ b/bittide-instances/src/Bittide/Instances/Pnr/Ethernet.hs
@@ -101,7 +101,7 @@ vexRiscGmii SNat sysClk sysRst rxClk rxRst txClk txRst fwd =
           [uartBus, timeBus, wbAxiRx, wbAxiTx, dnaWb, gpioWb, macWb] <- pe -< jtag
           (uartRx, _uartStatus) <- uart -< (uartBus, uartTx)
           time -< timeBus
-          dna -< dnaWb
+          _dna <- dnaC -< dnaWb
           macStatIf -< (macWb, macStatus)
           gpioDf <- idleSource -< ()
           gpioOut <- gpio -< (gpioWb, gpioDf)
@@ -116,7 +116,7 @@ vexRiscGmii SNat sysClk sysRst rxClk rxRst txClk txRst fwd =
       (fwd, (pure (), pure (), pure ()))
  where
   time = wcre timeWb
-  dna = wcre readDnaPortE2Wb simDna2
+  dnaC = wcre readDnaPortE2Wb simDna2
   mac =
     ethMac1GFifoC
       (SNat @1500)
diff --git a/bittide-instances/tests/Wishbone/DnaPortE2.hs b/bittide-instances/tests/Wishbone/DnaPortE2.hs
index 7c50e178e..ecab1bdca 100644
--- a/bittide-instances/tests/Wishbone/DnaPortE2.hs
+++ b/bittide-instances/tests/Wishbone/DnaPortE2.hs
@@ -63,7 +63,7 @@ dut = circuit $ \_unit -> do
   (uartRx, jtag) <- idleSource -< ()
   [uartBus, dnaWb] <- processingElement @dom NoDumpVcd peConfig -< jtag
   (uartTx, _uartStatus) <- uartInterfaceWb d2 d2 uartSim -< (uartBus, uartRx)
-  readDnaPortE2Wb simDna2 -< dnaWb
+  _dna <- readDnaPortE2Wb simDna2 -< dnaWb
   idC -< uartTx
  where
   (iMem, dMem) =
diff --git a/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
index 03cb1a151..0f2b2ee49 100644
--- a/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
+++ b/bittide-instances/tests/Wishbone/SwitchDemoProcessingElement.hs
@@ -91,10 +91,15 @@ dut localCounter dnaA dnaB = circuit $ do
   [uartBus, timeBus, peBusA, peBusB] <- processingElement NoDumpVcd peConfig -< jtagIdle
   (uartTx, _uartStatus) <- uartInterfaceWb d16 d2 uartSim -< (uartBus, uartRx)
   timeWb -< timeBus
-  linkAB <- switchDemoPeWb d2 localCounter (Just <$> dnaA) -< (peBusA, linkBA)
-  linkBA <- switchDemoPeWb d2 localCounter (Just <$> dnaB) -< (peBusB, linkAB)
+  linkAB <- switchDemoPeWb d2 localCounter -< (peBusA, dnaAC, linkBA)
+  linkBA <- switchDemoPeWb d2 localCounter -< (peBusB, dnaBC, linkAB)
+  dnaAC <- signalToCSignal dnaA -< ()
+  dnaBC <- signalToCSignal dnaB -< ()
   idC -< uartTx
  where
+  signalToCSignal :: Signal dom a -> Circuit () (CSignal dom a)
+  signalToCSignal = Circuit . const . ((),)
+
   memMap = 0b000 :> 0b001 :> 0b010 :> 0b011 :> 0b100 :> 0b101 :> Nil
   peConfig = unsafePerformIO $ do
     root <- findParentContaining "cabal.project"
diff --git a/bittide/src/Bittide/SwitchDemoProcessingElement.hs b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
index dca737849..5022072b4 100644
--- a/bittide/src/Bittide/SwitchDemoProcessingElement.hs
+++ b/bittide/src/Bittide/SwitchDemoProcessingElement.hs
@@ -7,7 +7,6 @@ module Bittide.SwitchDemoProcessingElement where
 
 import Clash.Prelude
 
-import Data.Maybe (fromMaybe)
 import Data.Tuple (swap)
 import GHC.Stack (HasCallStack)
 
@@ -40,7 +39,7 @@ switchDemoPe ::
   -- | Incoming crossbar link
   Signal dom (BitVector 64) ->
   -- | Device DNA
-  Signal dom (Maybe (BitVector 96)) ->
+  Signal dom (BitVector 96) ->
   -- | When to read from the crossbar link
   Signal dom (Unsigned 64) ->
   -- | How many tri-cycles to read from the crossbar link
@@ -54,7 +53,7 @@ switchDemoPe ::
   , -- \| Buffer output
     Signal dom (Vec (bufferSize * 3) (BitVector 64))
   )
-switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart writeCycles =
+switchDemoPe SNat localCounter linkIn dna readStart readCycles writeStart writeCycles =
   (linkOut, buffer)
  where
   readCyclesExtended = checkedResize . zeroExtendTimesThree <$> readCycles
@@ -64,8 +63,7 @@ switchDemoPe SNat localCounter linkIn maybeDna readStart readCycles writeStart w
   localData = bundle ((pack <$> localCounter) :> unbundle dnaVec)
    where
     dnaVec :: Signal dom (Vec 2 (BitVector 64))
-    dnaVec = reverse . bitCoerce . zeroExtend <$> dnaLocked
-    dnaLocked = fromMaybe 0xBAAB_BAAB_BAAB_BAAB_BAAB_BAAB <$> maybeDna
+    dnaVec = reverse . bitCoerce . zeroExtend <$> dna
 
   linkOut = stateToLinkOutput <$> peState <*> buffer <*> localData
 
@@ -163,16 +161,18 @@ switchDemoPeWb ::
   SNat bufferSize ->
   -- | Local clock cycle counter
   Signal dom (Unsigned 64) ->
-  -- | Device DNA
-  Signal dom (Maybe (BitVector 96)) ->
   Circuit
     ( Wishbone dom 'Standard addrW (Bytes 4)
-    , CSignal dom (BitVector 64)
+    , -- \| Device DNA
+      CSignal dom (BitVector 96)
+    , -- \| Incoming crossbar link
+      CSignal dom (BitVector 64)
     )
+    -- \| Outgoing crossbar link
     (CSignal dom (BitVector 64))
-switchDemoPeWb SNat localCounter maybeDna = Circuit go
+switchDemoPeWb SNat localCounter = Circuit go
  where
-  go ((wbM2S, linkIn), _) = ((wbS2M, pure ()), linkOut)
+  go ((wbM2S, dna, linkIn), _) = ((wbS2M, pure (), pure ()), linkOut)
    where
     readVec :: Vec (8 + bufferSize * 3 * 2 + 2) (Signal dom (BitVector 32))
     readVec =
@@ -187,7 +187,7 @@ switchDemoPeWb SNat localCounter maybeDna = Circuit go
         (SNat @bufferSize)
         localCounter
         linkIn
-        maybeDna
+        dna
         readStart
         readCycles
         writeStart
diff --git a/bittide/src/Bittide/Wishbone.hs b/bittide/src/Bittide/Wishbone.hs
index 7bc20df34..d5d197b4e 100644
--- a/bittide/src/Bittide/Wishbone.hs
+++ b/bittide/src/Bittide/Wishbone.hs
@@ -629,11 +629,11 @@ readDnaPortE2Wb ::
   ) =>
   -- | Simulation DNA value
   BitVector 96 ->
-  Circuit (Wishbone dom 'Standard addrW (Bytes nBytes)) ()
+  Circuit (Wishbone dom 'Standard addrW (Bytes nBytes)) (CSignal dom (BitVector 96))
 readDnaPortE2Wb simDna = circuit $ \wb -> do
   dnaDf <- dnaCircuit -< ()
-  _dna <- reg -< (wb, dnaDf)
-  idC -< ()
+  dna <- reg -< (wb, dnaDf)
+  idC -< dna
  where
   maybeDna = readDnaPortE2 hasClock hasReset hasEnable simDna
   regRst =
diff --git a/bittide/tests/Tests/SwitchDemoProcessingElement.hs b/bittide/tests/Tests/SwitchDemoProcessingElement.hs
index a73eaa4f3..1e095852f 100644
--- a/bittide/tests/Tests/SwitchDemoProcessingElement.hs
+++ b/bittide/tests/Tests/SwitchDemoProcessingElement.hs
@@ -127,7 +127,7 @@ prop_readThenWrite = H.property $ do
             bufferSizeSNat
             (fromList [clockStart ..])
             crossBarIn
-            (pure (Just deviceDna))
+            (pure deviceDna)
             (pure readStart)
             (pure nReadTriCycles)
             (pure writeStart)