Specialize Index for literal indices with LitIndex

This saves 16-bytes per Vec element.
chipsalliance · Jul 10, 2024 · e4f121e · e4f121e
1 parent 94b133e
commit e4f121e
Show file tree

Hide file tree

Showing 9 changed files with 32 additions and 14 deletions.
diff --git a/core/src/main/scala/chisel3/experimental/hierarchy/core/Lookupable.scala b/core/src/main/scala/chisel3/experimental/hierarchy/core/Lookupable.scala
@@ -9,7 +9,7 @@ import scala.annotation.implicitNotFound
 import scala.collection.mutable.HashMap
 import chisel3._
 import chisel3.experimental.dataview.{isView, reify, reifyIdentityView}
-import chisel3.internal.firrtl.ir.{Arg, ILit, Index, ModuleIO, Slot, ULit}
+import chisel3.internal.firrtl.ir.{Arg, ILit, Index, LitIndex, ModuleIO, Slot, ULit}
 import chisel3.internal.{throwException, Builder, ViewParent}
 import chisel3.internal.binding.{AggregateViewBinding, ChildBinding, CrossModuleBinding, ViewBinding, ViewWriteability}
 
@@ -128,7 +128,7 @@ object Lookupable {
     def unrollCoordinates(res: List[Arg], d: Data): (List[Arg], Data) = d.binding.get match {
       case ChildBinding(parent) =>
         d.getRef match {
-          case arg @ (_: Slot | _: Index | _: ModuleIO) => unrollCoordinates(arg :: res, parent)
+          case arg @ (_: Slot | _: Index | _: LitIndex | _: ModuleIO) => unrollCoordinates(arg :: res, parent)
           case other => err(s"unrollCoordinates failed for '$arg'! Unexpected arg '$other'")
         }
       case _ => (res, d)
@@ -139,6 +139,7 @@ object Lookupable {
         else {
           val next = (coor.head, d) match {
             case (Slot(_, name), rec: Record) => rec._elements(name)
+            case (LitIndex(_, n), vec: Vec[_]) => vec.apply(n)
             case (Index(_, ILit(n)), vec: Vec[_]) => vec.apply(n.toInt)
             case (ModuleIO(_, name), rec: Record) => rec._elements(name)
             case (arg, _) => err(s"Unexpected Arg '$arg' applied to '$d'! Root was '$start'.")

diff --git a/core/src/main/scala/chisel3/internal/Builder.scala b/core/src/main/scala/chisel3/internal/Builder.scala
@@ -265,8 +265,11 @@ private[chisel3] trait HasId extends chisel3.InstanceId {
     else setRef(OpaqueSlot(parent))
   }
 
-  private[chisel3] def setRef(parent: Node, index: Int):  Unit = setRef(Index(parent, ILit(index)))
-  private[chisel3] def setRef(parent: Node, index: UInt): Unit = setRef(Index(parent, index.ref))
+  private[chisel3] def setRef(parent: Node, index: Int): Unit = setRef(LitIndex(parent, index))
+  private[chisel3] def setRef(parent: Node, index: UInt): Unit = index.litOption match {
+    case Some(lit) if lit.isValidInt => setRef(LitIndex(parent, lit.intValue))
+    case _                           => setRef(Index(parent, index.ref))
+  }
   private[chisel3] def getRef:       Arg = _ref.get
   private[chisel3] def getOptionRef: Option[Arg] = _ref
 
@@ -651,6 +654,7 @@ private[chisel3] object Builder extends LazyLogging {
         case Slot(_, field)       => Some(field) // Record
         case OpaqueSlot(_)        => None // OpaqueSlots don't contribute to the name
         case Index(_, ILit(n))    => Some(n.toString) // Vec static indexing
+        case LitIndex(_, n)       => Some(n.toString) // Vec static indexing
         case Index(_, ULit(n, _)) => Some(n.toString) // Vec lit indexing
         case Index(_, _: Node) => None // Vec dynamic indexing
         case ModuleIO(_, n) => Some(n) // BlackBox port

diff --git a/core/src/main/scala/chisel3/internal/firrtl/Converter.scala b/core/src/main/scala/chisel3/internal/firrtl/Converter.scala
@@ -71,6 +71,8 @@ private[chisel3] object Converter {
       fir.SubField(convert(imm, ctx, info), name, fir.UnknownType)
     case OpaqueSlot(imm) =>
       convert(imm, ctx, info)
+    case LitIndex(imm, idx) =>
+      fir.SubIndex(convert(imm, ctx, info), idx, fir.UnknownType)
     case Index(imm, ILit(idx)) =>
       fir.SubIndex(convert(imm, ctx, info), castToInt(idx, "Index"), fir.UnknownType)
     case Index(imm, value) =>

diff --git a/core/src/main/scala/chisel3/internal/firrtl/IR.scala b/core/src/main/scala/chisel3/internal/firrtl/IR.scala
@@ -97,10 +97,11 @@ private[chisel3] object ir {
     def earlyLocalName(id: HasId, includeRoot: Boolean): String = id.getOptionRef match {
       case Some(Index(Node(imm), Node(value))) =>
         s"${earlyLocalName(imm, includeRoot)}[${earlyLocalName(value, includeRoot)}]"
-      case Some(Index(Node(imm), arg)) => s"${earlyLocalName(imm, includeRoot)}[${arg.localName}]"
-      case Some(Slot(Node(imm), name)) => s"${earlyLocalName(imm, includeRoot)}.$name"
-      case Some(OpaqueSlot(Node(imm))) => s"${earlyLocalName(imm, includeRoot)}"
-      case Some(arg) if includeRoot    => arg.name
+      case Some(LitIndex(Node(imm), idx)) => s"${earlyLocalName(imm, includeRoot)}[$idx]"
+      case Some(Index(Node(imm), arg))    => s"${earlyLocalName(imm, includeRoot)}[${arg.localName}]"
+      case Some(Slot(Node(imm), name))    => s"${earlyLocalName(imm, includeRoot)}.$name"
+      case Some(OpaqueSlot(Node(imm)))    => s"${earlyLocalName(imm, includeRoot)}"
+      case Some(arg) if includeRoot       => arg.name
       case None if includeRoot =>
         id match {
           case data: Data          => data._computeName(Some("?")).get
@@ -252,6 +253,13 @@ private[chisel3] object ir {
     override def localName: String = s"${imm.localName}[${value.localName}]"
   }
 
+  // Like index above, except the index is a literal, used for elements of Vecs
+  case class LitIndex(imm: Arg, value: Int) extends Arg {
+    def name: String = s"[$value]"
+    override def contextualName(ctx: Component): String = s"${imm.contextualName(ctx)}[$value]"
+    override def localName: String = s"${imm.localName}[$value]"
+  }
+
   sealed trait ProbeDetails { this: Arg =>
     val probe: Arg
     override def name: String = s"$probe"

diff --git a/panamaconverter/src/PanamaCIRCTConverter.scala b/panamaconverter/src/PanamaCIRCTConverter.scala
@@ -370,6 +370,7 @@ class PanamaCIRCTConverter(val circt: PanamaCIRCT, fos: Option[FirtoolOptions],
               binding.parent match {
                 case vec: Vec[_] =>
                   data.getRef match {
+                    case LitIndex(_, index)    => Reference.SubIndex(index, tpe)
                     case Index(_, ILit(index)) => Reference.SubIndex(index.toInt, tpe)
                     case Index(_, dynamicIndex) =>
                       val index = referTo(dynamicIndex, srcInfo)

diff --git a/src/main/scala/chisel3/aop/Select.scala b/src/main/scala/chisel3/aop/Select.scala
@@ -489,8 +489,9 @@ object Select {
   private def getEffected(a: Arg): Seq[Data] = a match {
     case Node(id: Data) => DataMirror.collectAllMembers(id)
     case Slot(imm: Node, name) => Seq(imm.id.asInstanceOf[Record].elements(name))
-    case Index(imm, _) => getEffected(imm)
-    case _             => throw new InternalErrorException("Match error: a=$a")
+    case Index(imm, _)    => getEffected(imm)
+    case LitIndex(imm, _) => getEffected(imm)
+    case _                => throw new InternalErrorException("Match error: a=$a")
   }
 
   // Given an arg, return the corresponding id. Don't use on a loc of a connect.

diff --git a/src/main/scala/chisel3/util/SRAM.scala b/src/main/scala/chisel3/util/SRAM.scala
@@ -7,7 +7,7 @@ import chisel3.internal.binding.{FirrtlMemTypeBinding, SramPortBinding}
 import chisel3.internal.plugin.autoNameRecursively
 import chisel3.experimental.SourceInfo
 import chisel3.internal.sourceinfo.{MemTransform, SourceInfoTransform}
-import chisel3.internal.firrtl.ir.{Arg, FirrtlMemory, ILit, Index, Node, Ref, Slot}
+import chisel3.internal.firrtl.ir.{Arg, FirrtlMemory, LitIndex, Node, Ref, Slot}
 import chisel3.util.experimental.loadMemoryFromFileInline
 import firrtl.annotations.MemoryLoadFileType
 import scala.language.reflectiveCalls
@@ -585,7 +585,7 @@ object SRAM {
         case v: Vec[_] =>
           v.elementsIterator.zipWithIndex.foreach {
             case (data, idx) =>
-              assignElementMask(data, writeMask, Index(arg, ILit(idx)))
+              assignElementMask(data, writeMask, LitIndex(arg, idx))
           }
       }
     }
@@ -599,7 +599,7 @@ object SRAM {
     ): Unit = {
       writeData.zip(writeMask).zipWithIndex.foreach {
         case ((elem, mask), idx) =>
-          assignElementMask(elem, mask, Index(arg, ILit(idx)))
+          assignElementMask(elem, mask, LitIndex(arg, idx))
       }
     }
   }

diff --git a/src/test/scala/chiselTests/DataPrint.scala b/src/test/scala/chiselTests/DataPrint.scala
@@ -62,6 +62,7 @@ class DataPrintSpec extends ChiselFlatSpec with Matchers {
     }))
     port.vec(idx).toString should be("BoundDataModule.port.vec[idx]: IO[UInt<8>]")
     port.vec(jdx.value).toString should be("BoundDataModule.port.vec[jdx.value]: IO[UInt<8>]")
+    port.vec(3.U).toString should be("BoundDataModule.port.vec[3]: IO[UInt<8>]")
 
     class InnerModule extends Module {
       val io = IO(Output(new Bundle {

diff --git a/src/test/scala/chiselTests/Vec.scala b/src/test/scala/chiselTests/Vec.scala
@@ -425,7 +425,7 @@ class VecSpec extends ChiselPropSpec with Utils {
       val out = IO(Output(UInt(8.W)))
       out := vec(10.U)
     }))
-    chirrtl should include("""connect out, vec[UInt<2>(0h2)]""")
+    chirrtl should include("""connect out, vec[2]""")
     log should include("Dynamic index with width 4 is too wide for Vec of size 4 (expected index width 2)")
   }