[compiler] Refactor BlockMatrix sparsity representations in type and lowering (#15163)

## Change Description

The bulk of the changes are in the files `BlockMatrixType.scala`,
`MatrixSparsity.scala`, and `LowerBlockMatrixIR.scala`. I suggest
starting with those, with the summaries below.

I acknowledge this is a big change containing a lot of non-trivial
logic. Please do ask for explanations and clarifications.

### BlockMatrixType and MatrixSparsity
`BlockMatrixType` has a `sparsity`, of type `BlockMatrixSparsity`, which
is either dense or a set of present blocks. In this PR I've moved the
sparsity representation to `MatrixSparsity`, which is a generic encoding
of the sparsity pattern of a sparse matrix. For a `BlockMatrixType`, the
sparsity will be a `nBlockRows` by `nBlockCols` `MatrixSparsity`.

Besides moving `BlockMatrixSparsity` to `MatrixSparsity`, and removing
references to blocks, the encoding of the sparsity is also streamlined.
Before, `BlockMatrixSparsity` held an array of coordinates of present
blocks, in an arbitrary order. Methods that need the present blocks in a
particular order (usually column major) would need to sort them, and we
also built a `Set` of present blocks to handle the `isPresent` query.

Now, `MatrixSparsity` (in the sparse case) enforces that the array of
present blocks is always in column major order. This simplifies some of
the logic, and lets us handle `isPresent` with binary search (similarly
for unions/intersections of sparsity, which now take advantage of the
ordering).

In addition, `MatrixSparsity` now knows its dimensions, which
`BlockMatrixSparsity` did not.

I've also gotten rid of all the complicated methods on the CSC
(compressed sparse column) encoding (e.g. `transposeCSCSparsity`).
Instead, I do all the transformations on the simple coordinate list
encoding, and convert to CSC as late as possible.

### LowerBlockMatrixIR
The changes in this file are mostly about a redesign of the
`BMSContexts` class. `BMSContexts` bundles together a runtime array of
context values (similar to the contexts for a `TableStage`), with a
representation of the sparsity pattern, which conceptually maps each
context value to the coordinates of its block.

Before, the sparsity pattern was also encoded by runtime values. In the
sparse case, this used runtime arrays `rowPos` and `rowIdx` using a CSC
encoding. Working with this encoding created a lot of non-trivial
runtime logic. But sparsity is completely known at compile time, so all
this runtime logic was unnecessary.

There are still cases where we need to work with a CSC encoding of the
sparsity at runtime. These cases are now handled by the class
`DynamicBMSContexts`, which is essentially the same as the old
`BMSContexts`, but now has a minimal interface of three methods.
Importantly, these methods are only about consuming a block matrix, not
transforming to a new one (like transpose).

The new `BMSContexts` now pairs a statically known `MatrixSparsity` with
a `DynamicBMSContexts`. The big change from before is that all methods
that produce a new `BMSContexts` now handle all the sparsity logic at
compile time, and simply embed the new sparsity in the IR using
literals. This typically requires us to also embed an array of ints
mapping from old to new positions in the contexts array, which we use at
runtime to reorder the contexts as needed.

## Security Assessment

- This change cannot impact the Hail Batch instance as deployed by Broad
Institute in GCP

Author: patrick-schultz

hail/build.mill

@@ -113,6 +113,10 @@ trait HailModule extends ScalaModule with ScalafmtModule with ScalafixModule { o
     mvn"io.github.tanin47::scalafix-forbidden-symbol:1.0.0",
   )
 
+  override def depManagement: T[Seq[Dep]] = Task {
+    Seq(Deps.collection_compat)
+  }
+
   override def javacOptions: T[Seq[String]] = Seq(
     "-Xlint:all",
     "-Werror",
@@ -350,7 +354,7 @@ trait RootHailModule extends CrossScalaModule with HailModule { outer =>
     override def defaultTask(): String = "generate"
 
     override def mvnDeps = Seq(
-      mvn"com.lihaoyi::mainargs:0.6.2",
+      mvn"com.lihaoyi::mainargs:0.7.7",
       mvn"com.lihaoyi::os-lib:0.10.7",
       mvn"com.lihaoyi::sourcecode:0.4.2",
     )

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

@@ -8,7 +8,7 @@ import is.hail.expr.ir.defs.{MetadataWriter, Str, UUID4, WriteMetadata, WriteVal
 import is.hail.expr.ir.lowering.{BlockMatrixStage2, LowererUnsupportedOperation}
 import is.hail.io.{StreamBufferSpec, TypedCodecSpec}
 import is.hail.io.fs.FS
-import is.hail.linalg.{BlockMatrix, BlockMatrixMetadata}
+import is.hail.linalg.{BlockMatrix, BlockMatrixMetadata, MatrixSparsity}
 import is.hail.types.TypeWithRequiredness
 import is.hail.types.encoded.{EBlockMatrixNDArray, ENumpyBinaryNDArray, EType}
 import is.hail.types.virtual._
@@ -135,8 +135,11 @@ case class BlockMatrixNativeMetadataWriter(
     cb: EmitCodeBuilder,
     region: Value[Region],
   ): Unit = {
-    val metaHelper =
-      BMMetadataHelper(path, typ.blockSize, typ.nRows, typ.nCols, typ.linearizedDefinedBlocks)
+    val partIdxToBlockIdx = typ.sparsity match {
+      case _: MatrixSparsity.Dense => None
+      case x: MatrixSparsity.Sparse => Some(x.definedBlocksColMajorLinear)
+    }
+    val metaHelper = BMMetadataHelper(path, typ.blockSize, typ.nRows, typ.nCols, partIdxToBlockIdx)
 
     val pc = writeAnnotations.getOrFatal(cb, "write annotations can't be missing!").asIndexable
     val partFiles = cb.newLocal[Array[String]]("partFiles")

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

@@ -512,6 +512,8 @@ package defs {
     def <=(other: IR): IR = ApplyComparisonOp(LTEQ, self, other)
 
     def >=(other: IR): IR = ApplyComparisonOp(GTEQ, self, other)
+
+    def log(messages: AnyRef*): IR = logIR(self, messages: _*)
   }
 
   object ErrorIDs {

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

@@ -14,7 +14,7 @@ import is.hail.io.gen.{BgenWriter, ExportGen}
 import is.hail.io.index.StagedIndexWriter
 import is.hail.io.plink.{BitPacker, ExportPlink}
 import is.hail.io.vcf.{ExportVCF, TabixVCF}
-import is.hail.linalg.BlockMatrix
+import is.hail.linalg.{BlockMatrix, MatrixSparsity}
 import is.hail.rvd.{IndexSpec, RVDPartitioner, RVDSpecMaker}
 import is.hail.types._
 import is.hail.types.encoded.{EBaseStruct, EBlockMatrixNDArray, EType}
@@ -2340,7 +2340,7 @@ case class MatrixBlockMatrixWriter(
 
     val countColumnsIR = ArrayLen(GetField(ts.getGlobals(), colsFieldName))
     val numCols: Int = CompileAndEvaluate[Int](ctx, countColumnsIR)
-    val numBlockCols: Int = (numCols - 1) / blockSize + 1
+    val numBlockCols: Int = BlockMatrixType.numBlocks(numCols.toLong, blockSize)
     val lastBlockNumCols = (numCols - 1) % blockSize + 1
 
     val rowCountIR = ts.mapCollect("matrix_block_matrix_writer_partition_counts")(paritionIR =>
@@ -2353,7 +2353,7 @@ case class MatrixBlockMatrixWriter(
     val inputPartStops = inputPartStartsPlusLast.tail
 
     val numRows = inputPartStartsPlusLast.last
-    val numBlockRows: Int = (numRows.toInt - 1) / blockSize + 1
+    val numBlockRows: Int = BlockMatrixType.numBlocks(numRows, blockSize)
 
     // Zip contexts with partition starts and ends
     val zippedWithStarts = ts.mapContexts { oldContextsStream =>
@@ -2510,7 +2510,7 @@ case class MatrixBlockMatrixWriter(
       numRows,
       numCols.toLong,
       blockSize,
-      BlockMatrixSparsity.dense,
+      MatrixSparsity.dense(numBlockRows, numBlockCols),
     )
     RelationalWriter.scoped(path, overwrite, None)(WriteMetadata(
       flatPaths,

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

@@ -311,6 +311,54 @@ object ArrayFunctions extends RegistryFunctions {
       ToArray(flatMapIR(ToStream(a))(ToStream(_)))
     }
 
+    /* Construct an array of length `len`, with the values in `elts` copied to the positions in
+     * `indices` */
+    registerSCode3t(
+      "scatter",
+      Array(tv("T")),
+      TArray(tv("T")), // elts
+      TArray(TInt32), // indices
+      TInt32, // len
+      TArray(tv("T")),
+      (_, a, _, _) => PCanonicalArray(a.asInstanceOf[SContainer].elementType.storageType()).sType,
+    ) {
+      case (
+            er,
+            cb,
+            _,
+            rt: SIndexablePointer,
+            elts: SIndexableValue,
+            indices: SIndexableValue,
+            len: SInt32Value,
+            errorID,
+          ) =>
+        cb.if_(
+          elts.loadLength.cne(indices.loadLength),
+          cb._fatalWithError(errorID, "scatter: values and indices arrays have different lengths"),
+        )
+        cb.if_(
+          elts.loadLength > len.value,
+          cb._fatalWithError(errorID, "scatter: values array is larger than result length"),
+        )
+        val pt = rt.pType.asInstanceOf[PCanonicalArray]
+        val (push, finish) =
+          pt.constructFromIndicesUnsafe(cb, er.region, len.value, deepCopy = false)
+        indices.forEachDefined(cb) { case (cb, pos, idx: SInt32Value) =>
+          cb.if_(
+            idx.value < 0 || idx.value >= len.value,
+            cb._fatalWithError(
+              errorID,
+              "scatter: indices array contained index ",
+              idx.value.toS,
+              ", which is greater than result length ",
+              len.value.toS,
+            ),
+          )
+          push(cb, idx.value, elts.loadElement(cb, pos))
+        }
+        finish(cb)
+    }
+
     registerSCode4(
       "lowerBound",
       TArray(tv("T")),

hail/hail/src/is/hail/expr/ir/functions/ArrayFunctions.scala

@@ -815,6 +815,31 @@ abstract class RegistryFunctions {
       case (r, cb, _, rt, Array(a1, a2, a3), errorID) => impl(r, cb, rt, a1, a2, a3, errorID)
     }
 
+  def registerSCode3t(
+    name: String,
+    typeParams: Array[Type],
+    mt1: Type,
+    mt2: Type,
+    mt3: Type,
+    rt: Type,
+    pt: (Type, SType, SType, SType) => SType,
+  )(
+    impl: (
+      EmitRegion,
+      EmitCodeBuilder,
+      Seq[Type],
+      SType,
+      SValue,
+      SValue,
+      SValue,
+      Value[Int],
+    ) => SValue
+  ): Unit =
+    registerSCode(name, Array(mt1, mt2, mt3), rt, unwrappedApply(pt), typeParams) {
+      case (r, cb, typeParams, rt, Array(a1, a2, a3), errorID) =>
+        impl(r, cb, typeParams, rt, a1, a2, a3, errorID)
+    }
+
   def registerSCode4(
     name: String,
     mt1: Type,