Make MatrixNormal sampling non-allocating

[projekter]

Currently, sampling from MatrixNormal always allocates multiple temporaries: one of the matrix X that contains the normally-sampled values, one for the matrix product A * X, and one for the matrix product (A * X) * B. However, given that A and B are triangular (and square) matrices, all products can be done in-place and the sample can directly be stored to the output. So the method doesn't need to allocate at all, provided the Cholesky factors are available without allocation.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.36%. Comparing base (fb714e1) to head (bf9f75f).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #2012   +/-   ##
=======================================
  Coverage   86.36%   86.36%           
=======================================
  Files         146      146           
  Lines        8786     8788    +2     
=======================================
+ Hits         7588     7590    +2     
  Misses       1198     1198           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[devmotion]

Looks good to me. Could you add a test for it, to ensure that it doesn't regress silently?

[projekter]

I added a test for this case. While I think that at least some, if not all, matrixvariate distributions can be made to sample without allocations, the tests would probably still need to be different for each one. For example, MatrixNormal will only avoid sampling if the U and V matrices already have the proper format - if the Cholesky factor is stored with the wrong triangle, it will copy upon retrieval, due to the way LinearAlgebra implements the retrieval of Cholesky factors. (Actually, this should probably be a motivation to bring it into the correct format upon creation, so that the sampling doesn't have to do it all over...)

[projekter]

So I tried to figure out why CI is failing on 1.10. On my machine with latest 1.10, it goes through, so I bisected and it turns out that it fails on 1.10.4 but passes on 1.10.5. Finding out what actually triggered the change was not very easy and I'm still not quite sure.
This is on 1.10.4. First, I allocate as in the test all my variables.

julia> ((d, out) -> @allocated(rand!(d, out)))(noallocD, output)
32

This is why the test fails. So let's do it step by step to see where it happens.

julia> ((d, out) -> @allocated(randn!(out)))(noallocD, output)
0

julia> ((d, out) -> @allocated(cholesky(d.U)))(noallocD, output)
0

Ok, so both the randn! line as well as getting the Cholesky factor are allocationless.

julia> A = cholesky(noallocD.U);
julia> ((A) -> @allocated(begin A.L; nothing end))(A)
16

So here's the culprit: Accessing the L field. Let's go to cholesky.jl:516 (getproperty).

julia> f = getfield(A, :factors);
julia> ((f) -> @allocated(LowerTriangular(f)))(f)
0

Now that's tricky. All it has to do is to construct the LowerTriangular, and this is actually allocationless. The getfields are as well (yes, I even checked this obvious thing). Still, there's an allocation. What I could guess is that getproperty is that the constant propagation in getproperty fails and therefore the output type is not well defined, hence has to be passed on the heap. This should not happen, though...
So let's check @code_native ((A) -> begin @inline getproperty(A, :L); nothing end)(A). Indeed, the allocation that we can see doesn't happen because of the return type, but it comes from checksquare. And guess what, checksquare was modified in 1.10.5. The string in the exception was turned into a lazy one. Let's try it:

julia> function LinearAlgebra.checksquare(A)
           m,n = size(A)
           m == n || throw(DimensionMismatch(lazy"matrix is not square: dimensions are $(size(A))"))
           m
       end;
julia> ((A) -> @allocated(begin A.L; nothing end))(A)
0

So problem found. Of course, what this doesn't explain is why the allocation of this exception string plays a role at all, as it can be clearly seen in the @code_native that the allocations are not even called if the dimensions do match. And the method was also called in LowerTriangular, whose direct call didn't trigger an allocation. So it seems to be a weird interplay between the multiple return types of getproperty, suboptimal constant propagation and a string allocation that never happes... In other words: I don't know.

The question is, what shall we do: Bump the target version (at least of the tests) to 1.10.5?

[devmotion]

Thank you!