[query] Dispatch TypeCheck on a per-node tag

TypeCheck.checkSingleNode was one match of about 20.8 KB of bytecode.
HotSpot never compiles a method over 8000 bytes, since
DontCompileHugeMethods is on by default. So it always ran in the
interpreter.

Every IR node class now carries a tag, and TypeCheck switches on it in
two levels. checkSingleNode switches on the node's family and calls one
method per family, which switches on the tag. Both levels compile to
jump tables. Each case casts the node and reads its fields. The largest
method, checkStream, is 4.2 KB.

ir-gen puts each node in one family, named where the node is defined:
node(Const, "I32", ...). There are 13 families: Const, Core, Struct,
Stream, Collection, Agg, IO, NDArray, Random, Relational, Table, Matrix
and BlockMatrix. Relational, Table and Matrix hold the nodes a lowering
pass removes, and a node goes in the family of the first invariant that
forbids it. BlockMatrix is removed at the same point as Table, and takes
the nodes that read or make a block matrix but no Table. The other
families are buckets that keep each method small, and a node goes in
the family of what it produces.

A tag holds the family in the bits above Tag.FamilyShift and the node's
position in its family below. Within a family, nodes are ordered so the
cases TypeCheck and ChildEnv handle sit on consecutive tags, which keeps
each switch dense enough to compile to a jump table.

ir-gen writes Tag_gen.scala, with the Family and Tag value classes and
an Int constant for each family and node. The constants are untyped
final vals, so a @switch on them compiles to a jump table. The Table,
Matrix and BlockMatrix nodes are written by hand. ir-gen lists them
among the generated nodes as handWritten(Table, "TableRead") and
generates their tags too.

BaseIR takes the tag as a final constructor field, so reading it is a
field load. An overridden def tag was a virtual call HotSpot could not
inline, and made TypeCheck 11% slower on a 1000-field MakeStruct. Tag
stores a Short, which fits in padding most node classes already have.
203 of 217 classes keep their size. Weighted by node visits in the
Python test suite, the average node grows from 42.4 to 42.5 bytes.

The NDArrayIR trait goes. ir-gen marks those nodes TypedIR[TNDArray]
instead, and TypeCheck reads x.typ.elementType.

In IRPassBenchmark, against the previous commit, TypeCheck drops from
4.9 ms to 1.0 ms on a 3000-binding Block, from 5.3 ms to 0.58 ms on
2000 nested StreamMaps, from 6.6 ms to 0.92 ms on a balanced tree and
from 1.5 ms to 0.26 ms on a 1000-field MakeStruct. The gain comes from
getting each method under 8000 bytes. A version that only split the
instanceof chain into four methods ran within noise of the tag version
over 3 forks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: ehigham

hail/hail/ir-gen/src/Main.scala

@@ -11,7 +11,11 @@ case class Name(str: String) {
   override def toString: String = str
 }
 
-abstract class BaseIR {
+/** @param tag
+  *   identifies this node's class, for dispatch with `@switch`. See [[Tag]]. It is a field, not a
+  *   method, so reading it is a load rather than a virtual call.
+  */
+abstract class BaseIR(final val tag: Tag) {
   def typ: VType
 
   protected def childrenSeq: IndexedSeq[BaseIR]

hail/hail/src/is/hail/expr/ir/BaseIR.scala

@@ -43,7 +43,7 @@ object BlockMatrixIR {
     )
 }
 
-sealed abstract class BlockMatrixIR extends BaseIR {
+sealed abstract class BlockMatrixIR(tag: Tag) extends BaseIR(tag) {
   override def typ: BlockMatrixType
   override def deepCopy: BlockMatrixIR = super.deepCopy.asInstanceOf[BlockMatrixIR]
 
@@ -55,7 +55,8 @@ sealed abstract class BlockMatrixIR extends BaseIR {
   def blockCostIsLinear: Boolean
 }
 
-case class BlockMatrixRead(reader: BlockMatrixReader) extends BlockMatrixIR {
+case class BlockMatrixRead(reader: BlockMatrixReader)
+    extends BlockMatrixIR(new Tag(Tag.BlockMatrixRead)) {
   override lazy val typ: BlockMatrixType = reader.fullType
 
   lazy val childrenSeq: IndexedSeq[BaseIR] = ArraySeq.empty
@@ -261,7 +262,7 @@ case class BlockMatrixPersistReader(id: String, typ: BlockMatrixType) extends Bl
 }
 
 case class BlockMatrixMap(child: BlockMatrixIR, eltName: Name, f: IR, needsDense: Boolean)
-    extends BlockMatrixIR {
+    extends BlockMatrixIR(new Tag(Tag.BlockMatrixMap)) {
   override def typ: BlockMatrixType = child.typ
 
   lazy val childrenSeq: IndexedSeq[BaseIR] = ArraySeq(child, f)
@@ -398,7 +399,7 @@ case class BlockMatrixMap2(
   rightName: Name,
   f: IR,
   sparsityStrategy: SparsityStrategy,
-) extends BlockMatrixIR with Logging {
+) extends BlockMatrixIR(new Tag(Tag.BlockMatrixMap2)) with Logging {
   override lazy val typ: BlockMatrixType =
     left.typ.copy(sparsity = sparsityStrategy.mergeSparsity(left.typ.sparsity, right.typ.sparsity))
 
@@ -529,7 +530,7 @@ case class BlockMatrixMap2(
 }
 
 case class BlockMatrixDot(left: BlockMatrixIR, right: BlockMatrixIR)
-    extends BlockMatrixIR with Logging {
+    extends BlockMatrixIR(new Tag(Tag.BlockMatrixDot)) with Logging {
   override lazy val typ: BlockMatrixType = {
     val blockSize = left.typ.blockSize
     assert(left.typ.nCols == right.typ.nRows)
@@ -588,7 +589,7 @@ case class BlockMatrixBroadcast(
   inIndexExpr: IndexedSeq[Int],
   shape: IndexedSeq[Long],
   blockSize: Int,
-) extends BlockMatrixIR {
+) extends BlockMatrixIR(new Tag(Tag.BlockMatrixBroadcast)) {
 
   val blockCostIsLinear: Boolean = child.blockCostIsLinear
 
@@ -679,7 +680,7 @@ case class BlockMatrixBroadcast(
 case class BlockMatrixAgg(
   child: BlockMatrixIR,
   axesToSumOut: IndexedSeq[Int],
-) extends BlockMatrixIR {
+) extends BlockMatrixIR(new Tag(Tag.BlockMatrixAgg)) {
 
   val blockCostIsLinear: Boolean = child.blockCostIsLinear
 
@@ -729,7 +730,7 @@ case class BlockMatrixAgg(
 case class BlockMatrixFilter(
   child: BlockMatrixIR,
   indices: IndexedSeq[IndexedSeq[Long]],
-) extends BlockMatrixIR {
+) extends BlockMatrixIR(new Tag(Tag.BlockMatrixFilter)) {
 
   assert(indices.length == 2)
 
@@ -772,7 +773,8 @@ case class BlockMatrixFilter(
   }
 }
 
-case class BlockMatrixDensify(child: BlockMatrixIR) extends BlockMatrixIR {
+case class BlockMatrixDensify(child: BlockMatrixIR)
+    extends BlockMatrixIR(new Tag(Tag.BlockMatrixDensify)) {
   override lazy val typ: BlockMatrixType = child.typ.densify
 
   override def blockCostIsLinear: Boolean = child.blockCostIsLinear
@@ -897,7 +899,7 @@ case class PerBlockSparsifier(blocks: IndexedSeq[Int]) extends BlockMatrixSparsi
 case class BlockMatrixSparsify(
   child: BlockMatrixIR,
   sparsifier: BlockMatrixSparsifier,
-) extends BlockMatrixIR {
+) extends BlockMatrixIR(new Tag(Tag.BlockMatrixSparsify)) {
   override lazy val typ: BlockMatrixType =
     child.typ.copy(sparsity = sparsifier.definedBlocks(child.typ))
 
@@ -915,7 +917,7 @@ case class BlockMatrixSparsify(
 }
 
 case class BlockMatrixSlice(child: BlockMatrixIR, slices: IndexedSeq[IndexedSeq[Long]])
-    extends BlockMatrixIR {
+    extends BlockMatrixIR(new Tag(Tag.BlockMatrixSlice)) {
   assert(slices.length == 2)
   assert(slices.forall(_.length == 3))
 
@@ -986,7 +988,7 @@ case class ValueToBlockMatrix(
   child: IR,
   shape: IndexedSeq[Long],
   blockSize: Int,
-) extends BlockMatrixIR {
+) extends BlockMatrixIR(new Tag(Tag.ValueToBlockMatrix)) {
   assert(shape.length == 2)
 
   val blockCostIsLinear: Boolean = true
@@ -1041,7 +1043,7 @@ case class BlockMatrixRandom(
   gaussian: Boolean,
   shape: IndexedSeq[Long],
   blockSize: Int,
-) extends BlockMatrixIR {
+) extends BlockMatrixIR(new Tag(Tag.BlockMatrixRandom)) {
 
   assert(shape.length == 2)
 

hail/hail/src/is/hail/expr/ir/BlockMatrixIR.scala

@@ -133,8 +133,8 @@ object ChildEnv {
   }
 
   /** Walks Block `b`'s children (each binding's value, then the body), threading `b`'s env so the
-    * walk is linear in its bindings. `f` takes child `i`, `i` and child `i`'s env, and returns child
-    * `i`'s final form, whose type binding `i` binds.
+    * walk is linear in its bindings. `f` takes child `i`, `i` and child `i`'s env, and returns
+    * child `i`'s final form, whose type binding `i` binds.
     */
   def walkBlock[E[_]: Scopes, A](
     b: Block,

hail/hail/src/is/hail/expr/ir/ChildEnv.scala

@@ -30,7 +30,7 @@ import java.io.OutputStream
 import org.json4s.{DefaultFormats, Extraction, Formats, JValue, ShortTypeHints}
 import org.json4s.JsonAST.{JNothing, JString}
 
-abstract class IR extends BaseIR {
+abstract class IR(tag: Tag) extends BaseIR(tag) {
   private var _typ: Type = null
 
   override def typ: Type =
@@ -348,10 +348,6 @@ package defs {
     }
   }
 
-  trait NDArrayIR extends TypedIR[TNDArray] {
-    def elementTyp: Type = typ.elementType
-  }
-
   object GetFieldByIdx {
     def apply(s: IR, field: Int): IR =
       (s.typ: @unchecked) match {