diff --git a/vecxt/src/dimMatCheck.scala b/vecxt/src/dimMatCheck.scala index 47490715..ccd65312 100644 --- a/vecxt/src/dimMatCheck.scala +++ b/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) diff --git a/vecxt/test/src/matMulLayoutChecks.test.scala b/vecxt/test/src/matMulLayoutChecks.test.scala index 481f9264..ba166ee7 100644 --- a/vecxt/test/src/matMulLayoutChecks.test.scala +++ b/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