Enrich ShiftRegister with SyncReadMem-based implementation. (#2891)

Co-authored-by: Jiuyang Liu <[email protected]>
Co-authored-by: Megan Wachs <[email protected]>
Co-authored-by: Jack Koenig <[email protected]>

integration-tests/src/test/scala/chiselTest/ShiftRegisterMemSpec.scala

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package chiselTests
+
+import chisel3._
+import chisel3.testers.BasicTester
+import chisel3.util.{Counter, ShiftRegister}
+import org.scalacheck.{Gen, Shrink}
+
+class ShiftMemTester(n: Int, dp_mem: Boolean) extends BasicTester {
+  val (cntVal, done) = Counter(true.B, n)
+  val start = 23.U
+  val sr = ShiftRegister.mem(cntVal + start, n, true.B, dp_mem, Some("simple_sr"))
+  when(RegNext(done)) {
+    assert(sr === start)
+    stop()
+  }
+}
+
+class ShiftRegisterMemSpec extends ChiselPropSpec {
+
+  implicit val nonNegIntShrinker: Shrink[Int] = Shrink.shrinkIntegral[Int].suchThat(_ >= 0)
+
+  property("ShiftRegister with dual-port SRAM should shift") {
+    forAll(Gen.choose(0, 4)) { (shift: Int) => assertTesterPasses { new ShiftMemTester(shift, true) } }
+  }
+
+  property("ShiftRegister with single-port SRAM should shift") {
+    forAll(Gen.choose(0, 6).suchThat(_ % 2 == 0)) { (shift: Int) =>
+      assertTesterPasses { new ShiftMemTester(shift, false) }
+    }
+  }
+}
+

macros/src/main/scala/chisel3/internal/sourceinfo/SourceInfoTransform.scala

@@ -325,6 +325,10 @@ class SourceInfoTransform(val c: Context) extends AutoSourceTransform {
   def inNResetDataEnArg(in: c.Tree, n: c.Tree, resetData: c.Tree, en: c.Tree): c.Tree = {
     q"$thisObj.$doFuncTerm($in, $n, $resetData, $en)($implicitSourceInfo)"
   }
+
+  def inNEnUseDualPortSramNameArg(in: c.Tree, n: c.Tree, en: c.Tree, useDualPortSram: c.Tree, name: c.Tree): c.Tree = {
+    q"$thisObj.$doFuncTerm($in, $n, $en, $useDualPortSram, $name)($implicitSourceInfo)"
+  }
 }
 
 // Workaround for https://github.com/sbt/sbt/issues/3966

src/main/scala/chisel3/util/Reg.scala

@@ -120,6 +120,94 @@ object ShiftRegister {
     implicit sourceInfo: SourceInfo
   ): T =
     ShiftRegisters(in, n, resetData, en).lastOption.getOrElse(in)
+
+  /** Returns the n-cycle delayed version of the input signal (SyncReadMem-based ShiftRegister implementation).
+    *
+    * @param in input to delay
+    * @param n  number of cycles to delay
+    * @param en enable the shift
+    * @param useDualPortSram dual port or single port SRAM based implementation
+    * @param name name of SyncReadMem object
+    */
+  def mem[T <: Data](in: T, n: Int, en: Bool, useDualPortSram: Boolean, name: Option[String]): T =
+    macro SourceInfoTransform.inNEnUseDualPortSramNameArg
+
+  /** @group SourceInfoTransformMacro */
+  def do_mem[T <: Data](
+    in:              T,
+    n:               Int,
+    en:              Bool,
+    useDualPortSram: Boolean,
+    name:            Option[String]
+  )(
+    implicit sourceInfo: SourceInfo
+  ): T = _apply_impl_mem(in, n, en, useDualPortSram, name)
+
+  private def _apply_impl_mem[T <: Data](
+    in:              T,
+    n:               Int,
+    en:              Bool = true.B,
+    useDualPortSram: Boolean = false,
+    name:            Option[String] = None
+  )(
+    implicit sourceInfo: SourceInfo
+  ): T = {
+    if (n == 0) {
+      in
+    } else if (n == 1) {
+      val out = RegEnable(in, en)
+      out
+    } else if (useDualPortSram) {
+      val mem = SyncReadMem(n, in.cloneType)
+      if (name != None) {
+        mem.suggestName(name.get)
+      }
+      val raddr = Counter(en, n)._1
+      val out = mem.read(raddr, en)
+
+      val waddr = RegEnable(raddr, (n - 1).U, en)
+      when(en) {
+        mem.write(waddr, in)
+      }
+      out
+    } else {
+      require(n % 2 == 0, "Odd shift register length with single-port SRAMs is not supported")
+
+      val out_sp0 = Wire(in.cloneType)
+      out_sp0 := DontCare
+
+      val out_sp1 = Wire(in.cloneType)
+      out_sp1 := DontCare
+
+      val mem_sp0 = SyncReadMem(n / 2, in.cloneType)
+      val mem_sp1 = SyncReadMem(n / 2, in.cloneType)
+
+      if (name != None) {
+        mem_sp0.suggestName(name.get + "_0")
+        mem_sp1.suggestName(name.get + "_1")
+      }
+
+      val index_counter = Counter(en, n)._1
+      val raddr_sp0 = index_counter >> 1.U
+      val raddr_sp1 = RegEnable(raddr_sp0, (n / 2 - 1).U, en)
+
+      val wen_sp0 = index_counter(0)
+      val wen_sp1 = WireDefault(false.B)
+      wen_sp1 := ~wen_sp0
+
+      when(en) {
+        val rdwrPort = mem_sp0(raddr_sp0)
+        when(wen_sp0) { rdwrPort := in }.otherwise { out_sp0 := rdwrPort }
+      }
+
+      when(en) {
+        val rdwrPort = mem_sp1(raddr_sp1)
+        when(wen_sp1) { rdwrPort := in }.otherwise { out_sp1 := rdwrPort }
+      }
+      val out = Mux(~wen_sp1, out_sp0, out_sp1)
+      out
+    }
+  }
 }
 
 object ShiftRegisters {

src/test/scala/chiselTests/util/RegSpec.scala

@@ -83,6 +83,19 @@ class ShiftRegisterSpec extends AnyFlatSpec with Matchers {
     (chirrtl should not).include("Reg.scala")
   }
 
+  it should "have source locators when passed in, n, en, useDualPortSram, name" in {
+    class MyModule extends Module {
+      val in = IO(Input(Bool()))
+      val out = IO(Output(Bool()))
+      out := ShiftRegister.mem(in, 2, true.B, false, Some("sr"))
+    }
+    val chirrtl = ChiselStage.emitCHIRRTL(new MyModule)
+    val reset = """reset .*RegSpec.scala""".r
+    (chirrtl should include).regex(reset)
+    val update = """out_r.* in .*RegSpec.scala""".r
+    (chirrtl should include).regex(update)
+    (chirrtl should not).include("Reg.scala")
+  }
 }
 
 class ShiftRegistersSpec extends AnyFlatSpec with Matchers {