rand methods

[cscherrer]

rand methods often look like this:

rand(d::T) = rand(Random.GLOBAL_RNG, d)

That's great for univariate distributions. But if sampletype(d) <: Array, there's going to have to be allocation. Ok fine, so maybe we have a Random.rand! method, and rand calls that for array types. So for example,

Random.rand!(rng::AbstractRNG, d::ProductMeasure, x::AbstractArray)

Now, say we have a distribution that's built from other distributions, which may themselves be built from other distributions, etc, and we want to draw a sample and insert it into an array without doing any allocations. One way to do this that seems natural to me is to make only rand! calls, and specify where each piece goes. Ideally, we can have things set up to be able to make a guarantee that rand! never allocates.

But this gets tricky. What if a sampletype along the way is immutable, say a Float64 or maybe a Tuple{Int, Char}? We can still have rand! put things into an array, but now we need a way to say where to put them.

We can't just pass things like x[2,3] as a destination for rand!, because it will be eagerly expanded before rand! is called.

I'm thinking we could have something like

struct ArraySlot{A,I}
    arr::A
    i::I
end

Then

function rand!(rng::AbstractRNG, d::AbstractMeasure, x::ArraySlot)
    x.arr[x.i...] = rand(rng, d)
end

What do you think?

[DilumAluthge]

I would recommend that the thing being mutated is the first argument.

That is:

function rand!(x::ArraySlot, rng::AbstractRNG, d::AbstractMeasure)

[DilumAluthge]

Everything else looks good to me!

[cscherrer]

I would recommend that the thing being mutated is the first argument.

Yeah, I'd think so too, first or maybe second. But then there's this:

help?> rand!
search: rand! randn rand transcode macroexpand @macroexpand1 @macroexpand CartesianIndex

  rand!([rng=GLOBAL_RNG], A, [S=eltype(A)])

  Populate the array A with random values. If S is specified (S can be a type or a
  collection, cf. rand for details), the values are picked randomly from S. This is
  equivalent to copyto!(A, rand(rng, S, size(A))) but without allocating a new array.

  Examples
  ≡≡≡≡≡≡≡≡≡≡

  julia> rng = MersenneTwister(1234);
  
  julia> rand!(rng, zeros(5))
  5-element Vector{Float64}:
   0.5908446386657102
   0.7667970365022592
   0.5662374165061859
   0.4600853424625171
   0.7940257103317943

  ───────────────────────────────────────────────────────────────────────────────────────────

  rand!([rng::AbstractRNG,] s::Sampleable, A::AbstractArray)

  Generate one or multiple samples from s to a pre-allocated array A. A should be in the form
  as specified above. The rules are summarized as below:

    •  When s is univariate, A can be an array of arbitrary shape. Each element of A will be
       overriden by one sample.

    •  When s is multivariate, A can be a vector to store one sample, or a matrix with each
       column for a sample.

    •  When s is matrix-variate, A can be a matrix to store one sample, or an array of
       matrices with each element for a sample matrix.

  ───────────────────────────────────────────────────────────────────────────────────────────

  rand!(rng::AbstractRNG, ::UnivariateDistribution, ::AbstractArray)

  Sample a univariate distribution and store the results in the provided array.

  ───────────────────────────────────────────────────────────────────────────────────────────

  rand!([rng::AbstractRNG,] d::MultivariateDistribution, x::AbstractArray)

  Draw samples and output them to a pre-allocated array x. Here, x can be either a vector of
  length dim(d) or a matrix with dim(d) rows.

  ───────────────────────────────────────────────────────────────────────────────────────────

  rand!(::AbstractRNG, ::Sampleable, ::AbstractArray)

  Samples in-place from the sampler and stores the result in the provided array.

  ───────────────────────────────────────────────────────────────────────────────────────────

  rand!(d::Union{UnivariateMixture, MultivariateMixture}, r::AbstactArray)

  Draw multiple samples from d and write them to r.

[DilumAluthge]

Huh. Good catch. We should be consistent with what other people have done previously.

[phipsgabler]

The style guide does not mention RNG arguments, but perhaps it can be argued that they are similar enough to IO streams.

[cscherrer]

@DilumAluthge Do you know of any work for a nice interface to avoid allocation, like this is trying to do? I'd think that comes up a lot.

[cscherrer]

A little more detail here, in case it's helpful.

First, every rand (without the bang) method can dispatch to

rand(::AbstractRNG, ::Type, ::AbstractMeasure)

Then this can allocate as needed and call rand!. So for example, something like

function Base.rand(rng::AbstractRNG, ::Type{A}, d::ProductMeasure) where {T, N, A <: AbstractArray{T,N}}
    dims = size(d.data)
    x = A(undef, dims)
    return @inbounds rand!(rng, d, x)
end

Then if we want to preallocate, we should be able to call rand! directly, and have a guarantee that there will be no allocation.

[chriselrod]

Perhaps of interest is this unit test from StrideArrays.jl:

@noinline function dostuff(A; B = I)
    StrideArrays.@avx for i in eachindex(A)
        A[i] += 1
    end
    C = A * B
    s = zero(eltype(C))
    StrideArrays.@avx for i in eachindex(C)
        s += C[i]
    end
    s
end
function gc_preserve_test()
    A = @StrideArray rand(8,8);
    B = @StrideArray rand(8,8);
    @gc_preserve dostuff(A, B = B)
end

@test iszero(@allocated gc_preserve_test())

It randomly generates two arrays, passes them to a function that isn't inlined, which then mutates one of them, multiplies them, and sums the result -- and none of this allocates.

Similarly, this discourse comment. I added make_dynamic from that example to that StrideArrays.jl PR, which I'll merge into master soon. Note that StrideArrays.jl is replacing PaddedMatrices.jl, since PaddedMatrices.jl hasn't actually been about padded matrices for quite a while now.

That's rather different than your ArraySlots idea, which is probably a better fit to your actual problem.

[cscherrer]

Thanks @chriselrod . I'm having trouble seeing how this doesn't allocate. I guess StrideArrays must do their own memory management? I think I need to learn more about the approach you take for these.

I think for MeasureTheory we want to find reasonable defaults but also have flexibility. So for example, we might have a ProductMeasure default to a vanilla Array, but allow the user to specify a different target and/or add methods. A lot of this is getting the right dispatch patterns in place to allow extending through interrupting the dispatch at different points.

[chriselrod]

The StrideArrays being created are statically sized, 8 x 8.
Julia stack-allocates statically sized mutable arrays if they don't escape (so this goes for MArrays too).

The tricky part is not escaping. Getting passed to a function that isn't inlined counts, because who knows what that function may do with it (e.g., push! into into a vector).

So the @gc_preserve macro GC.@preserves all arguments to a function, and tries to replace them with a PtrArray. The PtrArray contains information such as size and strides of the array, but it only holds a pointer to the original memory.
This way, the mutable arrays don't escape, yet are still protected from the GC while we use their memory in another function.
This of course means the functions you pass it to can't do things like store the array somewhere. E.g., if it's a sampled parameter vector, you could copy the data into a preallocated chain, but you couldn't store the PtrArray itself. Or you could, but then the memory would turn to garbage and be overwritten as soon as you leave the GC.@preserved region and the stack gets used for something else.

Maybe at some point I should see what it takes to get a dynamic alloca in Julia -- i.e., stack allocation of dynamically sized arrays. As I said earlier, it exists in LLVM, it's just not usable from Julia/with llvmcall because of how llvmcall works (creating a function that gets inlined; the stack memory is bound to the scope of the inlined function).

A hacky workaround would be to just pick some upper bound for the maximum size you want to use, and allocate that. Then you can make a dynamically sized PtrArray of whatever size smaller than that upper bound. The two disadvantages are (a) there's an upper bound, and (b) you waste stack memory when smaller than that. Do it too much, and it's not hard to get stack overflows.
I'm not sure how much stack memory gets allocated exactly, but 32-bit Julia does seem to have less, so you'd probably have to test on them to ensure support. Octavian.jl takes this approach for 64-bit builds only; it ran out of memory on 32-bit builds, and hence for them just uses a big global constant array.

[cscherrer]

Thanks @chriselrod . I'm thinking some more about this, and it's really interesting. I had been assuming we'd need to allocate once and then reuse the memory in order to avoid reallocation. But I guess stack allocation really changes the whole paradigm, so I need to rethink things.

I'd love to hear your thoughts on any limitations you see for the stack-based approach as you've implemented in StrideArrays, especially for probabilistic applications. In your experience, is it just better all around to use them instead of Arrays, or are there situations where you'd avoid them?

[matbesancon]

Yes rng usually goes first because it's the mutated argument for all rand(!) methods, which is not super practical otherwise many implementations could be written as rand(d, params, rng=DefaultRNG)

[cscherrer]

@matbesancon What do you think of this approach?

• rand always dispatches to Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractMeasure), which allocates if needed, then calls rand!
• rand! always dispatches to Base.rand!(rng::AbstractRNG, d::AbstractMeasure, x)

So for any new distribution, you only define the latter

[matbesancon]

which allocates if needed, then calls rand!

sounds good if you are always assigning to a mutable container (defaulting to Array). I haven't found a good design yet to let people specify the output container they want (i.e., you want small arrays of random values returned in StatisArrays for example)

[cscherrer]

I haven't found a good design yet to let people specify the output container they want

I was thinking this is determined by the second argument of rand or the third argument of rand!. In the 3-argument methods I describe above, I mean. Would that not do it?

[matbesancon]

oh ok I see, I thought T was an eltype that you commit to return from rand

[cscherrer]

eltype doesn't compose well, for example

julia> using Distributions

julia> eltype(Normal())
Float64

julia> eltype(MvNormal(ones(3)))
Float64

So we're using sampletype in its place. Each measure has a default sampletype, and the ::Type{T} argument here lets you choose a different one if you like. So for example, you may want your MvNormal to return SArray{Float32}, this is how you'd do that.