Add the programmable mux

bittide/bittide.cabal

@@ -223,6 +223,7 @@ library
     Bittide.ProcessingElement.ProgramStream
     Bittide.ProcessingElement.ReadElf
     Bittide.ProcessingElement.Util
+    Bittide.ProgrammableMux
     Bittide.ScatterGather
     Bittide.SharedTypes
     Bittide.Shutter
@@ -305,6 +306,7 @@ test-suite unittests
     Tests.GeneralPurposeProcessingElement.Calculator
     Tests.Handshake
     Tests.ProcessingElement.ReadElf
+    Tests.ProgrammableMux
     Tests.ScatterGather
     Tests.Shared
     Tests.Switch

bittide/src/Bittide/ProgrammableMux.hs

@@ -0,0 +1,85 @@
+-- SPDX-FileCopyrightText: 2024 Google LLC
+--
+-- SPDX-License-Identifier: Apache-2.0
+
+module Bittide.ProgrammableMux (programmableMux) where
+
+import Clash.Prelude
+import Protocols
+
+import Bittide.ElasticBuffer (stickyE)
+import Clash.Class.BitPackC (ByteOrder)
+import GHC.Stack (HasCallStack)
+import Protocols.MemoryMap (Access (ReadWrite, WriteOnly), Mm)
+import Protocols.MemoryMap.Registers.WishboneStandard (
+  RegisterConfig (access, description),
+  deviceWb,
+  registerConfig,
+  registerWbI,
+ )
+import Protocols.Wishbone (Wishbone, WishboneMode (Standard))
+
+{- | A mux which switches its output from input 'A' to input 'B' when the local counter is
+'cycle_to_switch' and the device is armed. Once switched the mux cannot switched back
+using the registers, but requires a reset.
+
+An intended use for this component is in the 'WireDemo'. In this demo, the management unit
+(input A) starts with control over the link, while the processing element (input B) is
+held in reset. When the management unit has finished its setup, it sets the
+'cycle_to_switch' and arms the mux. At the specified cycle, the mux switches to the
+processing element's links and deasserts its reset, allowing it to start writing to the
+link.
+-}
+programmableMux ::
+  forall dom addrW nBytes a.
+  ( HasCallStack
+  , HiddenClockResetEnable dom
+  , KnownNat addrW
+  , KnownNat nBytes
+  , 1 <= nBytes
+  , ?busByteOrder :: ByteOrder
+  , ?regByteOrder :: ByteOrder
+  ) =>
+  -- | Local counter
+  Signal dom (Unsigned 64) ->
+  Circuit
+    ( (ToConstBwd Mm, Wishbone dom 'Standard addrW nBytes)
+    , "A" ::: CSignal dom a
+    , "B" ::: CSignal dom a
+    )
+    ( "B_RESET" ::: CSignal dom Bool
+    , "OUT" ::: CSignal dom a
+    )
+programmableMux localCounter = circuit $ \(bus, a, b) -> do
+  [wbCycleToSwitch, wbArm] <- deviceWb "ProgrammableMux" -< bus
+
+  let
+    -- Trigger one cycle earlier to account for the delay from 'sticky'
+    trigger = stickyE hasClock hasReset $ arm .&&. (localCounter .==. (cycleToSwitch - 1))
+    bReset = not <$> trigger
+  Fwd linkOut <- muxC trigger -< (b, a)
+
+  (Fwd cycleToSwitch, _cycleToSwitchActivity) <-
+    registerWbI
+      (registerConfig "cycle_to_switch")
+        { access = ReadWrite
+        , description = "Clock cycle to switch from input 'A' to 'B'."
+        }
+      maxBound
+      -< (wbCycleToSwitch, Fwd (pure Nothing))
+
+  (Fwd arm, _armActivity) <-
+    registerWbI
+      (registerConfig "arm")
+        { access = WriteOnly
+        , description = "Arm the mux to switch on cycle_to_switch. Prevents atomicity issues."
+        }
+      False
+      -< (wbArm, Fwd (pure Nothing))
+
+  idC -< Fwd (bReset, linkOut)
+ where
+  muxC :: Signal dom Bool -> Circuit (CSignal dom a, CSignal dom a) (CSignal dom a)
+  muxC bool = Circuit go
+   where
+    go ((t, f), _) = (units, mux bool t f)

bittide/tests/Tests/ProgrammableMux.hs

@@ -0,0 +1,128 @@
+-- SPDX-FileCopyrightText: 2026 Google LLC
+--
+-- SPDX-License-Identifier: Apache-2.0
+{-# LANGUAGE OverloadedStrings #-}
+
+module Tests.ProgrammableMux where
+
+import Clash.Prelude
+import Protocols
+
+import Bittide.ProgrammableMux (programmableMux)
+import Bittide.SharedTypes (withByteOrderings)
+import Clash.Class.BitPackC (ByteOrder (..))
+import Clash.Class.BitPackC.Padding (packWordC)
+import Data.Maybe (fromJust)
+import Protocols.Hedgehog (defExpectOptions)
+import Protocols.MemoryMap
+import Protocols.Wishbone
+
+import Clash.Hedgehog.Sized.Unsigned (genUnsigned)
+import Hedgehog (Property)
+import Protocols.Wishbone.Standard.Hedgehog (WishboneMasterRequest (..))
+import Test.Tasty
+import Test.Tasty.Hedgehog (testProperty)
+import Test.Tasty.TH (testGroupGenerator)
+
+import qualified Data.List as L
+import qualified Data.Map as Map
+import qualified Data.String.Interpolate as Str
+import qualified Debug.Trace as Debug
+import qualified Hedgehog as H
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Range as Range
+import qualified Protocols.Wishbone.Standard.Hedgehog as Wb
+
+{- | Differentiate between the data from the management unit and processing element in the
+constructor. The Unsigned 64 is the cycle number at which the data is generated, so we can
+verify the exact cycle the programmable mux switches.
+-}
+data LinkData = Mu (Unsigned 64) | Pe (Unsigned 64)
+  deriving (Show, ShowX, Eq, Generic, NFDataX, BitPack)
+
+isMu :: LinkData -> Bool
+isMu (Mu _) = True
+isMu _ = False
+
+isPe :: LinkData -> Bool
+isPe (Pe _) = True
+isPe _ = False
+
+{- Lucas suspects a bug with the 'offset' in `maskWriteData', which now only triggers
+because we use a bus datawidth of 64 bits instead of the usual 32 bits.
+-}
+
+prop_ProgrammableMux :: Property
+prop_ProgrammableMux = H.property $ do
+  switchCycle <- H.forAll $ genUnsigned @64 (Range.linear 100 200)
+  busByteOrder <- H.forAll $ Gen.element [BigEndian, LittleEndian]
+  regByteOrder <- H.forAll $ Gen.element [BigEndian, LittleEndian]
+  let
+    dut =
+      withByteOrderings busByteOrder regByteOrder
+        $ withClockResetEnable @System clockGen resetGen enableGen
+        $ circuit
+        $ \wb -> do
+          let
+            counter = register 0 (counter + 1)
+            muLinks = fmap Mu counter
+            peLinks = fmap Pe counter
+          programmableMux @System @32 @8 counter -< (wb, Fwd muLinks, Fwd peLinks)
+
+    deviceName = "ProgrammableMux"
+    defs = (((getMMAny dut).deviceDefs) Map.! deviceName)
+
+    counterLoc = L.find (\loc -> loc.name.name == "cycle_to_switch") defs.registers
+    armLoc = L.find (\loc -> loc.name.name == "arm") defs.registers
+
+    counterAddr = fromIntegral (fromJust counterLoc).value.address `div` 8
+    armAddr = fromIntegral (fromJust armLoc).value.address `div` 8
+
+    switchCycleBv = pack $ packWordC @8 busByteOrder switchCycle
+    armBv = pack $ packWordC @8 busByteOrder True
+    requests = fmap (,0) [Write counterAddr maxBound switchCycleBv, Write armAddr maxBound armBv]
+
+    (resets, outLink) =
+      sampleC
+        def{timeoutAfter = simLength}
+        (unMemmap dut <| trace "ProgrammableMux" <| Wb.driveStandard defExpectOptions requests)
+    resetsBeforeSwitch = L.length $ L.takeWhile id resets
+    resetsAfterSwitch = L.length (L.takeWhile (== False) (L.dropWhile (== True) resets))
+
+    outLinkBeforeSwitch = L.length $ L.takeWhile isMu outLink
+    outLinkAfterSwitch = L.length $ L.takeWhile isPe (L.dropWhile isMu outLink)
+
+  H.footnote [Str.i|Asserted Cycles: #{resetsBeforeSwitch}|]
+  H.footnote [Str.i|Deasserted Cycles: #{resetsAfterSwitch}|]
+  let interestingResetCycles = L.take 10 (L.drop (fromIntegral switchCycle - 5) resets)
+  let interestingLinkCycles = L.take 10 (L.drop (fromIntegral switchCycle - 5) outLink)
+  H.footnote [Str.i|Interesting reset cycles (around the switch point): \n#{interestingResetCycles}|]
+  H.footnote [Str.i|Interesting link cycles (around the switch point): \n#{interestingLinkCycles}|]
+
+  -- Check that the reset is asserted before the switch cycle, and deasserted after. The
+  -- +-1 is to account for the single cycle the dut is in reset because of 'resetGen'.
+  resetsBeforeSwitch H.=== fromIntegral switchCycle + 1
+  resetsAfterSwitch H.=== simLength - fromIntegral switchCycle - 1
+
+  -- Check that the out link is from the MU before the switch cycle, and from the PE after.
+  outLinkBeforeSwitch H.=== fromIntegral switchCycle + 1
+  outLinkAfterSwitch H.=== simLength - fromIntegral switchCycle - 1
+ where
+  simLength = 300
+
+trace ::
+  (KnownDomain dom, KnownNat addrW, KnownNat nBytes) =>
+  String ->
+  Circuit (Wishbone dom mode addrW nBytes) (Wishbone dom mode addrW nBytes)
+trace msg =
+  Circuit
+    (unbundle . withClockResetEnable clockGen resetGen enableGen mealy go (0 :: Int) . bundle)
+ where
+  go cnt ~(m2s, s2m)
+    | m2s.busCycle = (cnt + 1, (s2m', m2s))
+    | otherwise = (cnt + 1, (s2m, m2s))
+   where
+    s2m' = Debug.trace [Str.i| Df.Trace #{msg} | #{cnt}: #{showX m2s}, #{showX s2m}|] s2m
+
+tests :: TestTree
+tests = $(testGroupGenerator)

bittide/tests/UnitTests.hs

@@ -24,6 +24,7 @@ import qualified Tests.ElasticBuffer
 import qualified Tests.GeneralPurposeProcessingElement.Calculator
 import qualified Tests.Handshake
 import qualified Tests.ProcessingElement.ReadElf
+import qualified Tests.ProgrammableMux
 import qualified Tests.ScatterGather
 import qualified Tests.Switch
 import qualified Tests.SwitchDemoProcessingElement
@@ -53,6 +54,7 @@ tests =
     , Tests.GeneralPurposeProcessingElement.Calculator.tests
     , Tests.Handshake.tests
     , Tests.ProcessingElement.ReadElf.tests
+    , Tests.ProgrammableMux.tests
     , Tests.ScatterGather.tests
     , Tests.Switch.tests
     , Tests.SwitchDemoProcessingElement.Calculator.tests