Reject aliased c in matmulInPlace!, add regression tests

vecxt/src/dimMatCheck.scala

@@ -14,11 +14,15 @@ object sameDimMatCheck:
     if !(a.cols == b.cols && a.rows == b.rows) then throw MatrixDimensionMismatch(a.rows, a.cols, b.rows, b.cols)
 end sameDimMatCheck
 
-/** Validates the output matrix `c` of a `matmulInPlace!` call: it must be shaped exactly `(m.rows, b.cols)` and dense
-  * column-major. `matmulInPlace!` hardcodes `ldc = m.rows` and always writes (and, when `beta != 0`, reads) `c`
-  * assuming that layout, so a wrongly-shaped or non-dense-column-major `c` would otherwise be corrupted or misread
-  * silently instead of failing loudly. `matmul`/`@@` always build a conforming `c` themselves, so this only bites
-  * direct callers of the in-place API.
+/** Validates the output matrix `c` of a `matmulInPlace!` call: it must be shaped exactly `(m.rows, b.cols)`, dense
+  * column-major, and backed by an array distinct from `m`'s and `b`'s. `matmulInPlace!` hardcodes `ldc = m.rows` and
+  * always writes (and, when `beta != 0`, reads) `c` assuming that layout, so a wrongly-shaped or non-dense-column-major
+  * `c` would otherwise be corrupted or misread silently instead of failing loudly. The aliasing check exists because
+  * BLAS `dgemm`/`sgemm` assume `c` does not overlap `a`/`b`: if `c.raw` is the same backing array as `m.raw` or
+  * `b.raw`, dgemm can read a not-yet-fully-read element of `a`/`b` after it's already been overwritten via the aliased
+  * `c`, silently corrupting the result (a's and b's own arrays may safely be the same as each other, e.g. `m @@ m`,
+  * since dgemm only ever reads those two, never writes them). `matmul`/`@@` always build a conforming, freshly
+  * allocated `c` themselves, so this only bites direct callers of the in-place API.
   */
 object matmulOutputCheck:
   inline def apply(m: Matrix[?], b: Matrix[?], c: Matrix[?]): Unit =
@@ -29,6 +33,16 @@ object matmulOutputCheck:
         s"matmulInPlace! requires a dense column-major output matrix `c`, but got layout: ${c.layoutString}"
       )
     end if
+    if c.raw.asInstanceOf[AnyRef] eq m.raw.asInstanceOf[AnyRef] then
+      throw MatrixAliasingException(
+        "matmulInPlace! requires `c` to be backed by a different array than `m` - writing into c while reading m from the same array would corrupt the result"
+      )
+    end if
+    if c.raw.asInstanceOf[AnyRef] eq b.raw.asInstanceOf[AnyRef] then
+      throw MatrixAliasingException(
+        "matmulInPlace! requires `c` to be backed by a different array than `b` - writing into c while reading b from the same array would corrupt the result"
+      )
+    end if
   end apply
 end matmulOutputCheck
 
@@ -136,3 +150,5 @@ case class InvalidMatrix(cols: Int, rows: Int, data: Int)
     )
 
 case class UnsupportedLayoutException(message: String) extends Exception(message)
+
+case class MatrixAliasingException(message: String) extends Exception(message)

vecxt/test/src/matMulLayoutChecks.test.scala

@@ -323,6 +323,42 @@ class DifferentMemoryLayoutTests extends FunSuite:
     }
   }
 
+  // `dgemm`/`sgemm` assume `c` does not overlap `a`/`b`: it's read progressively while `c` is written, so if `c`
+  // shares a backing array with `m` or `b`, dgemm can read an element of `m`/`b` after it's already been
+  // overwritten via the aliased `c`, silently corrupting the result. `m`/`b` themselves may safely share an array
+  // (e.g. `m @@ m`), since dgemm only ever reads those two - it never writes them.
+
+  test("matmulInPlace! throws when c aliases m's backing array") {
+    val m = Matrix.fromRows(Array(1.0, 2.0), Array(3.0, 4.0)) // 2x2
+    val b = Matrix.fromRows(Array(5.0, 6.0), Array(7.0, 8.0)) // 2x2
+
+    intercept[MatrixAliasingException] {
+      m.`matmulInPlace!`(b, m, 1.0, 0.0)
+    }
+  }
+
+  test("matmulInPlace! throws when c aliases b's backing array") {
+    val m = Matrix.fromRows(Array(1.0, 2.0), Array(3.0, 4.0)) // 2x2
+    val b = Matrix.fromRows(Array(5.0, 6.0), Array(7.0, 8.0)) // 2x2
+    // Same backing array as `b`, wrapped in a distinct (but identically shaped/laid out) Matrix instance - not the
+    // same object as `b`, so this specifically exercises the array-identity check rather than reference equality
+    // on the Matrix wrapper.
+    val cAliasingB = Matrix[Double](b.raw, b.rows, b.cols)
+
+    intercept[MatrixAliasingException] {
+      m.`matmulInPlace!`(b, cAliasingB, 1.0, 0.0)
+    }
+  }
+
+  test("matmulInPlace! allows m and b to alias each other") {
+    val m = Matrix.fromRows(Array(1.0, 2.0), Array(3.0, 4.0)) // 2x2
+    val out = Matrix.zeros[Double]((2, 2))
+
+    m.`matmulInPlace!`(m, out, 1.0, 0.0)
+
+    assertMatrixEquals(out, m @@ m)
+  }
+
   test("scalars in matmul, non-square") {
     val m = Matrix.fromRows(Array(1.0, 2.0, 3.0), Array(4.0, 5.0, 6.0)) // 2x3
     val b = Matrix.fromRows(Array(1.0, 0.0), Array(0.0, 1.0), Array(1.0, 1.0)) // 3x2