Open BrendonChau opened 1 year ago
Thanks for the report. How many threads were used?
I wrote matrix-vector multiplication three years ago, and there could be room for improvement. I will look into the GX
part. Something that sounds counterintuitive to me is that if multiple threads work on the same tile, there could be some cache conflict, but that doesn't seem to be the case here. Let's check these parts in other environments.
When I read the code today, I see things I could have done differently. For example, recursive tiling strategies would have made code maintenance easier than the current fixed-size block approach (at the cost of slightly worse performance). The machinery based on Tullio.jl it is in the OpenADMIXTURE package's loop.jl
file.
Sorry for interjecting 2 questions here.
@BrendonChau, just want to confirm, are you saying that G_mul_x_tiled!
is 5x faster than SnpLinAlg
because SnpLinAlg
has an extra parallel layer? I thought it was okay to have multiple layers of parallelization as long as it is not nested @threads
.
Also, can you explain a bit more on this:
Even better, by not parallelizing over tiles, we avoid all heap allocations.
where did the heap allocations come from? I'm asking because previously I also observed lots of allocations but I couldn't figure out where they came from.
@kose-y I ran it with 8 threads on my local machine.
julia> versioninfo()
Julia Version 1.9.2
Commit e4ee485e90 (2023-07-05 09:39 UTC)
Platform Info:
OS: macOS (arm64-apple-darwin22.4.0)
CPU: 10 × Apple M1 Max
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-14.0.6 (ORCJIT, apple-m1)
Threads: 8 on 8 virtual cores
Environment:
JULIA_EDITOR = code
JULIA_NUM_THREADS = 8
@biona001 I believe the allocations come from Julia's own resource management using Threads.@threads
and from assignment into array views. Since we must ensure thread-safety, my understanding is that GC.@preserve
happens to allocate for array views. Using @tturbo
or @turbo thread=true
is non-allocating for a single tile operation.
I did some more investigating and it seems that my implementation for $\mathbf{Gx}$ is only faster for 'wide' SnpArray
s. For 'square' or 'tall' SnpArrays
I actually see a small performance regression relative to SnpLinAlg
which is still mysterious to me.
@BrendonChau While I haven't had a chance to look into it yet, that certainly makes sense to me. The previous example wasn't just "tall" enough--perhaps we can change behavior in that particular case. Allocations from multithreading didn't really seem major there--how is it with large arrays?
I modified my functions so that remainder handling is less verbose. I also wrote a version of $\mathbf{Gx}$ that uses Threads.@threads
to parallelize over the tiles in a column, this has more allocations than computing over each tile in parallel, but is faster for large SnpArrays
, being consistently faster than SnpLinAlg
. Also, for the chunk calculations, I use Base.Cartesian.@nexprs
to directly unroll the innermost loop, this is about 20-30% faster than just letting @turbo
handle it.
using StatsBase, Statistics, Plots, LinearAlgebra, BenchmarkTools, Random
using SnpArrays, DelimitedFiles, Glob
using Plots, LoopVectorization, Polyester
function G_mul_x_chunk!(
y::AbstractVector{T},
G::AbstractMatrix{UInt8},
x::AbstractVector{Tx},
μ::AbstractVector{T},
σinv::AbstractVector{T}
) where {T<:AbstractFloat, Tx<:Union{Signed, AbstractFloat}}
@turbo thread=true for k in axes(G, 2)
for j in axes(G, 1)
# for i in 1:4
# # index into the correct value in compressed SNP array
# Gijk = (G[j, k] >> ((i - 1) << 1)) & 0x03
# y[((j - 1) << 2) + i] += x[k] * (T((Gijk >= 0x02) * (Gijk - 0x01)) - !isone(Gijk) * μ[k]) * σinv[k]
# end
# Macro directly unrolls the i ∈ 1:4 loop, ~20-30% faster
Base.Cartesian.@nexprs 4 i -> begin
G_i_jk = (G[j, k] >> ((i - 1) * 2)) & 0x03
y[((j - 1) << 2) + i] += x[k] * (T((G_i_jk >= 0x02) * (G_i_jk - 0x01)) - !isone(G_i_jk) * μ[k]) * σinv[k]
end
end
end
end
function G_mul_x_tiled!(
y::AbstractVector{T},
G::SnpArray,
x::AbstractVector{Tx},
μ::AbstractVector{T},
σinv::AbstractVector{T};
chunkwidth::Int = 1024
) where {T<:AbstractFloat, Tx<:Union{Signed, AbstractFloat}}
m, d = size(G)
mpacked = size(G.data, 1)
packedchunkwidth = chunkwidth >> 2
# need K and rem to handle remainder rows
K, rem = divrem(m, 4)
# Slice up row iteration space
nrowchunks, rowchunkrem = divrem(m, chunkwidth)
# if remainder is nonzero, number of rowslices is nrowchunks + 1
rowchunkiters = rowchunkrem == 0 ? nrowchunks : nrowchunks + 1
# Slice up col iteration space
ncolchunks, colchunkrem = divrem(d, chunkwidth)
# if remainder is nonzero, number of colslices is ncolchunks + 1
colchunkiters = colchunkrem == 0 ? ncolchunks : ncolchunks + 1
# Initialize output vector to zero,
# * could be modified to agree with 5-argument mul!
fill!(y, zero(T))
# multithreading only on inner tile loop avoids pointer allocations using @turbo thread=true
# a more sophisticated way might use buffers and multi-level parallelism
@inbounds for col in 1:colchunkiters
colstart = (col - 1) * chunkwidth
colstop = col <= ncolchunks ? col * chunkwidth : d
_x = view(x, (colstart + 1):colstop)
_μ = view(μ, (colstart + 1):colstop)
_σinv = view(σinv, (colstart + 1):colstop)
for l in 1:rowchunkiters
rowstart = (l - 1) * chunkwidth
rowstop = l <= nrowchunks ? l * chunkwidth : min(m, K << 2)
packedrowstart = (l - 1) * packedchunkwidth
packedrowstop = l <= nrowchunks ? l * packedchunkwidth : min(mpacked, K)
_y = view(y, (rowstart + 1):rowstop)
_G_data = view(G.data, (packedrowstart + 1):packedrowstop, (colstart + 1):colstop)
G_mul_x_chunk!(_y, _G_data, _x, _μ, _σinv)
end
if rem > 0
width = col <= ncolchunks ? chunkwidth : colchunkrem
@turbo thread=true for k in 1:width
for i in 1:rem
# index into the correct value in compressed SNP array
Gijk = (G.data[K + 1, colstart + k] >> ((i - 1) << 1)) & 0x03
y[K << 2 + i] += _x[k] * (T((Gijk >= 0x02) * (Gijk - 0x01)) - !isone(Gijk) * _μ[k]) * _σinv[k]
end
end
end
end
return y
end
function G_mul_x_chunk_2!(
y::AbstractVector{T},
G::AbstractMatrix{UInt8},
x::AbstractVector{Tx},
μ::AbstractVector{T},
σinv::AbstractVector{T}
) where {T<:AbstractFloat, Tx<:Union{Signed, AbstractFloat}}
@turbo thread=false for k in axes(G, 2)
for j in axes(G, 1)
# for i in 1:4
# # index into the correct value in compressed SNP array
# Gijk = (G[j, k] >> ((i - 1) << 1)) & 0x03
# y[((j - 1) << 2) + i] += x[k] * (T((Gijk >= 0x02) * (Gijk - 0x01)) - !isone(Gijk) * μ[k]) * σinv[k]
# end
# Macro directly unrolls the i ∈ 1:4 loop, ~20-30% faster
Base.Cartesian.@nexprs 4 i -> begin
G_i_jk = (G[j, k] >> ((i - 1) * 2)) & 0x03
y[((j - 1) << 2) + i] += x[k] * (T((G_i_jk >= 0x02) * (G_i_jk - 0x01)) - !isone(G_i_jk) * μ[k]) * σinv[k]
end
end
end
end
function G_mul_x_tiled_2!(
y::AbstractVector{T},
G::SnpArray,
x::AbstractVector{Tx},
μ::AbstractVector{T},
σinv::AbstractVector{T};
chunkwidth::Int = 1024
) where {T<:AbstractFloat, Tx<:Union{Signed, AbstractFloat}}
m, d = size(G)
mpacked = size(G.data, 1)
packedchunkwidth = chunkwidth >> 2
# need K and rem to handle remainder rows
K, rem = divrem(m, 4)
# Slice up row iteration space
nrowchunks, rowchunkrem = divrem(m, chunkwidth)
# if remainder is nonzero, number of rowslices in nrowchunks + 1
rowchunkiters = rowchunkrem == 0 ? nrowchunks : nrowchunks + 1
# Slice up col iteration space
ncolchunks, colchunkrem = divrem(d, chunkwidth)
# if remainder is nonzero, number of colslices in ncolchunks + 1
colchunkiters = colchunkrem == 0 ? ncolchunks : ncolchunks + 1
# Initialize output vector to zero,
# * could be modified to agree with 5-argument mul!
fill!(y, zero(T))
@inbounds for col in 1:colchunkiters
colstart = (col - 1) * chunkwidth
colstop = col <= ncolchunks ? col * chunkwidth : d
_x = view(x, (colstart + 1):colstop)
_μ = view(μ, (colstart + 1):colstop)
_σinv = view(σinv, (colstart + 1):colstop)
# multi-threading on inside loop should be safe since rowchunks
# assign into disjoint slices of y, dont need to bother with
# @sync and Threads.@spawn
Threads.@threads for l in 1:rowchunkiters
rowstart = (l - 1) * chunkwidth
rowstop = l <= nrowchunks ? l * chunkwidth : min(m, K << 2)
packedrowstart = (l - 1) * packedchunkwidth
packedrowstop = l <= nrowchunks ? l * packedchunkwidth : min(mpacked, K)
_y = view(y, (rowstart + 1):rowstop)
_G_data = view(G.data, (packedrowstart + 1):packedrowstop, (colstart + 1):colstop)
G_mul_x_chunk_2!(_y, _G_data, _x, _μ, _σinv)
if l > nrowchunks && rem > 0
width = col <= ncolchunks ? chunkwidth : colchunkrem
@turbo thread=false for k in 1:width
for i in 1:rem
# index into the correct value in compressed SNP array
Gijk = (G.data[K + 1, colstart + k] >> ((i - 1) << 1)) & 0x03
y[K << 2 + i] += _x[k] * (T((Gijk >= 0x02) * (Gijk - 0x01)) - !isone(Gijk) * _μ[k]) * _σinv[k]
end
end
end
end
end
return y
end
First the smaller simulation, G
is 65_667
by 131_072
rng = MersenneTwister(1234)
# SnpArrays.simulate! is from my fork
m = 1_024 * 64 + 128 + 3
q = 1_024 * 128
minfreq = 0.01
MAFs = minfreq .+ (0.5 - minfreq) .* rand(rng, q)
G = SnpArrays.simulate!(rng, m, q, MAFs)
# Second row of `columncounts` indicates number of missing
G_counts = counts(G, dims = 1)
m_nonmissing = m .- G_counts[2, :]
# Allele Frequencies
ρ = (G_counts[3, :] + 2 * G_counts[4, :]) ./ (2 * m_nonmissing)
# Expected Values
μ = 2 * ρ
# Standard Deviations
σ = sqrt.(2 .* (1 .- ρ) .* ρ)
σinv = inv.(σ)
# Testing Gx
xvecq = rand(rng, (-1.0, 1.0), q);
resm = Vector{Float64}(undef, m);
resm2 = similar(resm);
@benchmark G_mul_x_tiled!($resm2, $G, $xvecq, $μ, $σinv; chunkwidth=1024) evals=1
BenchmarkTools.Trial: 7 samples with 1 evaluation.
Range (min … max): 713.627 ms … 819.101 ms ┊ GC (min … max): 0.00% … 0.00%
Time (median): 740.464 ms ┊ GC (median): 0.00%
Time (mean ± σ): 750.719 ms ± 36.066 ms ┊ GC (mean ± σ): 0.00% ± 0.00%
█ █ █ █ █ █ █
█▁▁▁▁▁▁█▁█▁▁▁▁▁█▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
714 ms Histogram: frequency by time 819 ms <
Memory estimate: 0 bytes, allocs estimate: 0.
resm3 = similar(resm);
@benchmark G_mul_x_tiled_2!($resm3, $G, $xvecq, $μ, $σinv; chunkwidth=1024) evals=1
BenchmarkTools.Trial: 7 samples with 1 evaluation.
Range (min … max): 707.115 ms … 759.956 ms ┊ GC (min … max): 0.00% … 0.00%
Time (median): 716.626 ms ┊ GC (median): 0.00%
Time (mean ± σ): 721.859 ms ± 17.641 ms ┊ GC (mean ± σ): 0.00% ± 0.00%
█ █ ██ █ █ █
█▁▁▁█▁▁▁▁██▁▁█▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
707 ms Histogram: frequency by time 760 ms <
Memory estimate: 828.09 KiB, allocs estimate: 6403.
# Compare to SnpLinAlg
Gsla_ = SnpLinAlg{Float64}(G; model=ADDITIVE_MODEL, center=true, scale=true, impute=true);
resm_sla = similar(resm);
@benchmark mul!($resm_sla, $Gsla_, $xvecq) evals=1
BenchmarkTools.Trial: 7 samples with 1 evaluation.
Range (min … max): 747.357 ms … 765.554 ms ┊ GC (min … max): 0.00% … 0.00%
Time (median): 750.175 ms ┊ GC (median): 0.00%
Time (mean ± σ): 752.828 ms ± 6.519 ms ┊ GC (mean ± σ): 0.00% ± 0.00%
█ █ █ █ █ █ █
█▁█▁▁▁█▁▁█▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
747 ms Histogram: frequency by time 766 ms <
Memory estimate: 1.53 MiB, allocs estimate: 24069.
Then, the larger simulation, G
is 131_203
by `262_144
rng = MersenneTwister(1234)
# SnpArrays.simulate! is from my fork
m = 1_024 * 128 + 128 + 3
q = 1_024 * 256
minfreq = 0.01
MAFs = minfreq .+ (0.5 - minfreq) .* rand(rng, q)
G = SnpArrays.simulate!(rng, m, q, MAFs)
# Second row of `columncounts` indicates number of missing
G_counts = counts(G, dims = 1)
m_nonmissing = m .- G_counts[2, :]
# Allele Frequencies
ρ = (G_counts[3, :] + 2 * G_counts[4, :]) ./ (2 * m_nonmissing)
# Expected Values
μ = 2 * ρ
# Standard Deviations
σ = sqrt.(2 .* (1 .- ρ) .* ρ)
σinv = inv.(σ)
# Testing Gx
xvecq = rand(rng, (-1.0, 1.0), q);
resm = Vector{Float64}(undef, m);
resm2 = similar(resm);
@benchmark G_mul_x_tiled!($resm2, $G, $xvecq, $μ, $σinv; chunkwidth=1024) evals=1
BenchmarkTools.Trial: 2 samples with 1 evaluation.
Range (min … max): 3.212 s … 3.321 s ┊ GC (min … max): 0.00% … 0.00%
Time (median): 3.267 s ┊ GC (median): 0.00%
Time (mean ± σ): 3.267 s ± 76.880 ms ┊ GC (mean ± σ): 0.00% ± 0.00%
█ █
█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
3.21 s Histogram: frequency by time 3.32 s <
Memory estimate: 0 bytes, allocs estimate: 0.
resm3 = similar(resm);
@benchmark G_mul_x_tiled_2!($resm3, $G, $xvecq, $μ, $σinv; chunkwidth=1024) evals=1
BenchmarkTools.Trial: 2 samples with 1 evaluation.
Range (min … max): 2.708 s … 2.753 s ┊ GC (min … max): 0.00% … 0.00%
Time (median): 2.730 s ┊ GC (median): 0.00%
Time (mean ± σ): 2.730 s ± 31.661 ms ┊ GC (mean ± σ): 0.00% ± 0.00%
█ █
█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
2.71 s Histogram: frequency by time 2.75 s <
Memory estimate: 1.62 MiB, allocs estimate: 12816.
# Compare to SnpLinAlg
Gsla_ = SnpLinAlg{Float64}(G; model=ADDITIVE_MODEL, center=true, scale=true, impute=true);
resm_sla = similar(resm);
@benchmark mul!($resm_sla, $Gsla_, $xvecq) evals=1
BenchmarkTools.Trial: 2 samples with 1 evaluation.
Range (min … max): 2.970 s … 2.979 s ┊ GC (min … max): 0.00% … 0.00%
Time (median): 2.974 s ┊ GC (median): 0.00%
Time (mean ± σ): 2.974 s ± 6.346 ms ┊ GC (mean ± σ): 0.00% ± 0.00%
█ █
█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
2.97 s Histogram: frequency by time 2.98 s <
Memory estimate: 5.83 MiB, allocs estimate: 93187.
In the larger simulation G_mul_x_tiled!
is no longer performant, but G_mul_x_tiled_2!
manages to keep up. Allocations are still relatively small for both G_mul_x_tiled_2!
and SnpLinAlg
. The performance gap is now closer to a 10% improvement, so not particularly impressive. For larger SnpArray
s, the bottleneck ends up being memory latency, as we would expect. I could pick a better chunkwidth
like 128 or 192 to make sure all working memory can fit in the L1 cache for each core, but that would only be marginally beneficial for large arrays.
Which one do you think more affects the performance, parallelizing over the tiles in a column or loop unrolling via Base.Cartesian.@nexprs
?
I'm fairly confident that parallelization has a much larger impact than the loop unrolling, at least for large SnpArray
s.
I need some clarification, as there already is parallelization happening with the current package. So do you mean limiting the parallelization to a single column has more impact than 20~30% acceleration from the loop unrolling?
For stochastic trace estimation, we require being able to evaluate terms like $\mathbf{G}^\top \mathbf{x}$ and $\mathbf{G}\mathbf{x}$ efficiently, where $\mathbf{G}$ is a SnpArray. I was looking into
SnpLinAlg
and noticed that I was having alot of memory allocations, presumably due to use of JuliaBase.@threads
and pointer allocation from the tiling procedure. However, it turns out that only multithreading within a tile evaluation rather than multithreading over several tiles yields a substantial performance improvement.First the section defining the functions
Gt_mul_x!
,Gt_mul_x_tiled!
,G_mul_x!
, andG_mul_x_tiled!
And setting up data for simulation:
First, we can look at $\mathbf{G}^\top \mathbf{x}$,
then $\mathbf{Gx}$,
And finally, looking at the same operations using
SnpLinAlg
For $\mathbf{G}^\top \mathbf{x}$, tiling provides no improvement over just using
@turbo thread=true
, when using SnpLinAlg, there is substantial overhead from using Base threads and tiled views. Whereas for $\mathbf{Gx}$ there is an over 5x performance improvement from using sequential tiling with limiting multithreading to a single tile computation. Even better, by not parallelizing over tiles, we avoid all heap allocations.I suspect we might see a similar improvement with $\mathbf{G}^\top \mathbf{X}_1$ and $\mathbf{GX}_2$ where $\mathbf{X}_1$ and $\mathbf{X}_2$ are matrices, but implementation of that might be very tedious dealing with tile remainder chunks.
I did these simulations on my M1 Max Macbook, please let me know if the improvement is comparable for other hardware. It is possible that for something like a 32-core Xeon processor on a server we would see different behavior.