Repository: gonum/gonum (go)

Issue:
mat: calling qr.Factorize leads to OOM for matrixes with many rows
### What are you trying to do?

After upgrading gonum to 0.15.0 calls to `(VecDense).SolveVec(a,b)` where `a.Dims() == (10k+, 3)`  leads to OOM.
The stack where OOM happens:
```
runtime.systemstack_switch()
        /nix/store/8yw3g52r95h7cv09lcrran92n212997b-go-1.22.0/share/go/src/runtime/asm_amd64.s:474 +0x8 fp=0xc000087230 sp=0xc000087220 pc=0x4788a8
runtime.(*mheap).alloc(0x3667b52000?, 0x1b33da9?, 0x80?)
        /nix/store/8yw3g52r95h7cv09lcrran92n212997b-go-1.22.0/share/go/src/runtime/mheap.go:958 +0x5b fp=0xc000087278 sp=0xc000087230 pc=0x42fc1b
runtime.(*mcache).allocLarge(0x100?, 0x3667b50b88, 0x1?)
        /nix/store/8yw3g52r95h7cv09lcrran92n212997b-go-1.22.0/share/go/src/runtime/mcache.go:234 +0x85 fp=0xc0000872c0 sp=0xc000087278 pc=0x41ce45
runtime.mallocgc(0x3667b50b88, 0xd67320, 0x1)
        /nix/store/8yw3g52r95h7cv09lcrran92n212997b-go-1.22.0/share/go/src/runtime/malloc.go:1165 +0x597 fp=0xc000087348 sp=0xc0000872c0 pc=0x4141d7
runtime.makeslice(0xc0000124b0?, 0x18?, 0xc0000874b0?)
        /nix/store/8yw3g52r95h7cv09lcrran92n212997b-go-1.22.0/share/go/src/runtime/slice.go:107 +0x49 fp=0xc000087370 sp=0xc000087348 pc=0x459989
gonum.org/v1/gonum/mat.NewDense(...)
        /home/ericwenn/go/pkg/mod/gonum.org/v1/gonum@v0.15.1-0.20240413203616-1b7d9ca04ac9/mat/dense.go:60
gonum.org/v1/gonum/mat.(*QR).updateQ(0xc000087630)
        /home/ericwenn/go/pkg/mod/gonum.org/v1/gonum@v0.15.1-0.20240413203616-1b7d9ca04ac9/mat/qr.go:107 +0x78 fp=0xc0000874c0 sp=0xc000087370 pc=0xb5fd38
gonum.org/v1/gonum/mat.(*QR).factorize(0xc000087630, {0x1070598, 0xc0004c6540}, 0x49)
        /home/ericwenn/go/pkg/mod/gonum.org/v1/gonum@v0.15.1-0.20240413203616-1b7d9ca04ac9/mat/qr.go:101 +0x316 fp=0xc0000875e0 sp=0xc0000874c0 pc=0xb5fc36
gonum.org/v1/gonum/mat.(*QR).Factorize(...)
        /home/ericwenn/go/pkg/mod/gonum.org/v1/gonum@v0.15.1-0.20240413203616-1b7d9ca04ac9/mat/qr.go:82
gonum.org/v1/gonum/mat.(*Dense).Solve(0xc0004c6700, {0x1070598, 0xc0004c6540}, {0x1070598, 0xc0004c6740})
        /home/ericwenn/go/pkg/mod/gonum.org/v1/gonum@v0.15.1-0.20240413203616-1b7d9ca04ac9/mat/solve.go:65 +0x245 fp=0xc000087700 sp=0xc0000875e0 pc=0xb63545
gonum.org/v1/gonum/mat.(*VecDense).SolveVec(0xc000110360, {0x1070598, 0xc0004c6540}, {0x10735c0, 0xc0001103c0})
        /home/ericwenn/go/pkg/mod/gonum.org/v1/gonum@v0.15.1-0.20240413203616-1b7d9ca04ac9/mat/solve.go:118 +0x6fa fp=0xc000087930 sp=0xc000087700 pc=0xb63e9a
``` 
I've traced this down to changes in `mat.QR.updateQ` [here](https://github.com/gonum/gonum/commit/45b74210d616dd61bb893f6e8e4af13879b1d6bb#diff-5bf82974a09725d659ece49d56eff0b81eafc1a476a618949ec743544e329aaeR104-R117).

What `updateQ` does is allocate a new `Dense` with dimensions `MxM`, which for our value on M would mean allocating 100M cell matrix. Based on my understanding of the code, the resulting matrix should have dimensions (`MxN`).


### What version of Go and Gonum are you using?
<!--
Paste the output of `go version` and if you are installing Gonum from source, paste
the output of `(cd $(go env GOPATH)/src/gonum.org/v1/gonum && git rev-parse HEAD)`.
If you are using modules, also paste the output of `grep gonum.org/v1/gonum go.sum`,
executed in the root of your dependent module.
-->

```
$ go version
go version go1.22.0 linux/amd64

$ grep gonum.org/v1/gonum go.sum
gonum.org/v1/gonum v0.15.0 h1:2lYxjRbTYyxkJxlhC+LvJIx3SsANPdRybu1tGj9/OrQ=
gonum.org/v1/gonum v0.15.0/go.mod h1:xzZVBJBtS+Mz4q0Yl2LJTk+OxOg4jiXZ7qBoM0uISGo=
```

### Does this issue reproduce with the current master?

Yes


Pull request:
mat: calculate Q elements lazily when calling QR.At
Fixes #1968.
Builds upon #1970 initial proposal.
- mat: fix dst matrix shape check in QR.RTo
> When a matrix is very tall, calculating Q will currently allocate a large Q at the end of the factorisation, even if it is not going to be used, and a large Q matrix can lead to out of memory issues.
> 
> For this reason, Q is never eagerly computed unless explicitly required to by the user, with QR.ToQ.
To keep fulfilling the Matrix interface, the QR.At method will compute the requested element only, which only require computing a single row of Q.


Small typo found and fixed:
- mat: calculate Q elements lazily when calling QR.At
> Fixed a typo in row/column size check, the expected row size was incorrectly checked against the column size.

Please take a look.

<!--
Checklist:

- API changes have been discussed
- code is goformated correctly (goimports)
- packages with generated code have had code generation run
- tests pass locally
- linked to relevant issues

Please make sure your commit message summary line and pull request
title match the Go convention; a one-line summary of the change,
prefixed by the primary affected package that should complete the
sentence, "This change modifies Gonum to _____."
-->


Patch:
diff --git a/AUTHORS b/AUTHORS
index 791b6985..8e5896db 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -125,6 +125,7 @@ The University of Washington
 Thomas Berg <tomfuture@gmail.com>
 Tobin Harding <me@tobin.cc>
 Tom Payne <twpayne@gmail.com>
+Tristan Nicholls <tvk.nicholls@gmail.com>
 Valentin Deleplace <deleplace2015@gmail.com>
 Vincent Thiery <vjmthiery@gmail.com>
 Vladimír Chalupecký <vladimir.chalupecky@gmail.com>
diff --git a/CONTRIBUTORS b/CONTRIBUTORS
index a7e62e1e..e367595b 100644
--- a/CONTRIBUTORS
+++ b/CONTRIBUTORS
@@ -128,6 +128,7 @@ Tamir Hyman <hyman.tamir@gmail.com>
 Thomas Berg <tomfuture@gmail.com>
 Tobin Harding <me@tobin.cc>
 Tom Payne <twpayne@gmail.com>
+Tristan Nicholls <tvk.nicholls@gmail.com>
 Valentin Deleplace <deleplace2015@gmail.com>
 Vincent Thiery <vjmthiery@gmail.com>
 Vladimír Chalupecký <vladimir.chalupecky@gmail.com>
diff --git a/mat/qr.go b/mat/qr.go
index af99dbca..7f8fec8f 100644
--- a/mat/qr.go
+++ b/mat/qr.go
@@ -31,8 +31,13 @@ func (qr *QR) Dims() (r, c int) {
 	return qr.qr.Dims()
 }
 
-// At returns the element at row i, column j.
+// At returns the element at row i, column j. At will panic if the receiver
+// does not contain a successful factorization.
 func (qr *QR) At(i, j int) float64 {
+	if !qr.isValid() {
+		panic(badQR)
+	}
+
 	m, n := qr.Dims()
 	if uint(i) >= uint(m) {
 		panic(ErrRowAccess)
@@ -41,6 +46,20 @@ func (qr *QR) At(i, j int) float64 {
 		panic(ErrColAccess)
 	}
 
+	if qr.q == nil || qr.q.IsEmpty() {
+		// Calculate Qi, Q i-th row
+		qi := getFloat64s(m, true)
+		qr.qRowTo(i, qi)
+
+		// Compute QR(i,j)
+		var val float64
+		for k := 0; k <= j; k++ {
+			val += qi[k] * qr.qr.at(k, j)
+		}
+		putFloat64s(qi)
+		return val
+	}
+
 	var val float64
 	for k := 0; k <= j; k++ {
 		val += qr.q.at(i, k) * qr.qr.at(k, j)
@@ -48,6 +67,25 @@ func (qr *QR) At(i, j int) float64 {
 	return val
 }
 
+// qRowTo extracts the i-th row of the orthonormal matrix Q from a QR
+// decomposition.
+func (qr *QR) qRowTo(i int, dst []float64) {
+	c := blas64.General{
+		Rows:   1,
+		Cols:   len(dst),
+		Stride: len(dst),
+		Data:   dst,
+	}
+	c.Data[i] = 1 // C is the i-th unit vector
+
+	// Construct Qi from the elementary reflectors: Qi = C * (H(1) H(2) ... H(nTau))
+	work := []float64{0}
+	lapack64.Ormqr(blas.Right, blas.NoTrans, qr.qr.mat, qr.tau, c, work, -1)
+	work = getFloat64s(int(work[0]), false)
+	lapack64.Ormqr(blas.Right, blas.NoTrans, qr.qr.mat, qr.tau, c, work, len(work))
+	putFloat64s(work)
+}
+
 // T performs an implicit transpose by returning the receiver inside a
 // Transpose.
 func (qr *QR) T() Matrix {
@@ -98,7 +136,9 @@ func (qr *QR) factorize(a Matrix, norm lapack.MatrixNorm) {
 	lapack64.Geqrf(qr.qr.mat, qr.tau, work, len(work))
 	putFloat64s(work)
 	qr.updateCond(norm)
-	qr.updateQ()
+	if qr.q != nil {
+		qr.q.Reset()
+	}
 }
 
 func (qr *QR) updateQ() {
@@ -149,7 +189,7 @@ func (qr *QR) RTo(dst *Dense) {
 		dst.ReuseAs(r, c)
 	} else {
 		r2, c2 := dst.Dims()
-		if c != r2 || c != c2 {
+		if r != r2 || c != c2 {
 			panic(ErrShape)
 		}
 	}
@@ -192,6 +232,10 @@ func (qr *QR) QTo(dst *Dense) {
 			panic(ErrShape)
 		}
 	}
+
+	if qr.q == nil || qr.q.IsEmpty() {
+		qr.updateQ()
+	}
 	dst.Copy(qr.q)
 }
 


Test patch:
diff --git a/mat/qr_test.go b/mat/qr_test.go
index b71bee56..e9ec26aa 100644
--- a/mat/qr_test.go
+++ b/mat/qr_test.go
@@ -18,9 +18,11 @@ func TestQR(t *testing.T) {
 	rnd := rand.New(rand.NewSource(1))
 	for _, test := range []struct {
 		m, n int
+		big  bool
 	}{
-		{5, 5},
-		{10, 5},
+		{m: 5, n: 5},
+		{m: 10, n: 5},
+		{m: 1e5, n: 3, big: true}, // Test that very tall matrices do not OoM.
 	} {
 		m := test.m
 		n := test.n
@@ -35,6 +37,13 @@ func TestQR(t *testing.T) {
 
 		var qr QR
 		qr.Factorize(a)
+		if test.big {
+			_ = qr.At(0, 0)     // should not panic, even for big matrices
+			_ = qr.At(m-1, n-1) // should not panic, even for big matrices
+			// We cannot proceed past here for big matrices.
+			continue
+		}
+
 		var q, r Dense
 		qr.QTo(&q)
 
@@ -56,6 +65,23 @@ func TestQR(t *testing.T) {
 		if !EqualApprox(&got, &want, 1e-12) {
 			t.Errorf("QR does not equal original matrix. \nWant: %v\nGot: %v", want, got)
 		}
+
+		// Verify indirect QR.At()
+		got.Reset()
+		got.ReuseAs(m, n)
+		qr.q.Reset() // reset q matrix to force lazy computation
+		for i := 0; i < m; i++ {
+			for j := 0; j < n; j++ {
+				got.set(i, j, qr.At(i, j))
+			}
+		}
+
+		if !EqualApprox(a, &got, 1e-14) {
+			t.Errorf("m=%d,n=%d: A and QR (computed with QR.At()) are not equal", m, n)
+		}
+		if !EqualApprox(a.T(), got.T(), 1e-14) {
+			t.Errorf("m=%d,n=%d: Aᵀ and (QR)ᵀ (computed with QR.At()) are not equal", m, n)
+		}
 	}
 }