Skip to content

[WIP] Swap serial lu for MPI parallel qr using MPIQR.jl #307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jwscook wants to merge 16 commits into
base: master
Choose a base branch
from
Open

[WIP] Swap serial lu for MPI parallel qr using MPIQR.jl #307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
87809e4
this is running with my dev'd MPIQR with 3arg ldiv!
Jan 17, 2025
57588a5
block_rank[] >= 0 plus others
jwscook Jan 17, 2025
a839dc8
this works with latest MPIQR.jl
jwscook Jan 20, 2025
7b80960
put println statements back in
Jan 21, 2025
bc7f5f4
delete dead line
Jan 21, 2025
f7468a6
timeit_debug_enabled set back to false
Jan 21, 2025
43f2414
delete dead line
Jan 21, 2025
676a9be
time_advance was forced to use MPIQR rather than ADI - revert this to…
Jan 21, 2025
35ba4b1
tidy up
Jan 21, 2025
933d749
increase MPIQRBLOCKSIZE to 8
Jan 21, 2025
52f939e
delete obsolete comment
Jan 21, 2025
2ed2b17
make using statements alphabetical
Jan 21, 2025
2f932ee
make using statements alphabetical
Jan 21, 2025
733e919
make all caps const MPIQRBLOCKSIZE mpiqr_blocksize as inkeeping with …
Jan 21, 2025
207ba22
remove unnecessary defensive _block_synchonize() calls
Jan 21, 2025
792de1d
use in-place 2-arg ldiv!
Jan 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion moment_kinetics/Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
LsqFit = "2fda8390-95c7-5789-9bda-21331edee243"
MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
MPIQR = "4acd9526-7714-4958-a7d5-1632fb49fc8a"
Measures = "442fdcdd-2543-5da2-b0f3-8c86c306513e"
OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
Expand Down Expand Up @@ -50,11 +51,12 @@ file_io_netcdf = "NCDatasets"
manufactured_solns_ext = ["Symbolics", "IfElse"]

[compat]
MPIQR = "0.1.0"
julia = "1.9.0"

[extras]
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[targets]
test = ["Random", "Test", "NCDatasets"]
test = ["Random", "Test"]
73 changes: 46 additions & 27 deletions moment_kinetics/src/electron_kinetic_equation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ module electron_kinetic_equation

using LinearAlgebra
using MPI
using MPIQR
using SparseArrays

export get_electron_critical_velocities
Expand Down Expand Up @@ -71,6 +72,8 @@ import ..runge_kutta

const integral_correction_sharpness = 4.0

const mpiqr_blocksize = 8 # 1 will be slow, 4 is good for smaller matrices, 8 is good for larger matrices, 16 - who knows?!

"""
update_electron_pdf is a function that uses the electron kinetic equation
to solve for the updated electron pdf
Expand Down Expand Up @@ -1217,7 +1220,7 @@ pressure \$p_{e∥}\$.
end

if nl_solver_params.solves_since_precon_update[] ≥ nl_solver_params.preconditioner_update_interval
global_rank[] == 0 && println("recalculating precon")
global_rank[] == 0 && println("recalculating precon :electron_lu 1")
nl_solver_params.solves_since_precon_update[] = 0
nl_solver_params.precon_dt[] = t_params.dt[]

Expand All @@ -1231,27 +1234,37 @@ global_rank[] == 0 && println("recalculating precon")
external_source_settings, num_diss_params, t_params, ion_dt,
ir, evolve_ppar)

begin_serial_region()
if block_rank[] == 0
if size(orig_lu) == (1, 1)
if block_rank[] >= 0
localcols = MPIQR.localcolumns(block_rank[], size(precon_matrix, 2),
mpiqr_blocksize, block_size[])
if size(orig_lu, 1) == 1
# Have not properly created the LU decomposition before, so
# cannot reuse it.
@timeit_debug global_timer "lu" nl_solver_params.preconditioners[ir] =
(lu(sparse(precon_matrix)), precon_matrix, input_buffer,
output_buffer)
@timeit_debug global_timer "lu" begin
mpiqrmatrix = MPIQR.MPIQRMatrix(precon_matrix[:, localcols],
size(precon_matrix); blocksize=mpiqr_blocksize, comm=comm_block[])
mpiqrstruct = qr!(mpiqrmatrix)
nl_solver_params.preconditioners[ir] =
(mpiqrstruct, precon_matrix, input_buffer, output_buffer)
end
else
# LU decomposition was previously created. The Jacobian always
# has the same sparsity pattern, so by using `lu!()` we can
# reuse some setup.
try
@timeit_debug global_timer "lu!" lu!(orig_lu, sparse(precon_matrix); check=false)
@timeit_debug global_timer "lu!" begin
orig_lu[:, localcols] .= precon_matrix[:, localcols]
qr!(orig_lu)
end
catch e
if !isa(e, ArgumentError)
rethrow(e)
end
println("Sparsity pattern of matrix changed, rebuilding "
* " LU from scratch")
@timeit_debug global_timer "lu" orig_lu = lu(sparse(precon_matrix))
mpiqrmatrix = MPIQR.MPIQRMatrix(precon_matrix[:, localcols],
size(precon_matrix); blocksize=mpiqr_blocksize, comm=comm_block[])
@timeit_debug global_timer "lu" orig_lu = qr!(mpiqrmatrix)
end
nl_solver_params.preconditioners[ir] =
(orig_lu, precon_matrix, input_buffer, output_buffer)
Expand Down Expand Up @@ -1280,10 +1293,8 @@ global_rank[] == 0 && println("recalculating precon")
counter += 1
end

begin_serial_region()
@serial_region begin
@timeit_debug global_timer "ldiv!" ldiv!(this_output_buffer, precon_lu, this_input_buffer)
end
this_output_buffer .= this_input_buffer
@timeit_debug global_timer "ldiv!" MPIQR.ldiv!(precon_lu, this_output_buffer)

begin_serial_region()
counter = 1
Expand Down Expand Up @@ -1338,6 +1349,7 @@ global_rank[] == 0 && println("recalculating precon")
left_preconditioner = identity
right_preconditioner = lu_precon!
elseif nl_solver_params.preconditioner_type === Val(:electron_adi)
global_rank[] == 0 && println("recalculating precon :electron_adi")

if t_params.dt[] > 1.5 * nl_solver_params.precon_dt[] ||
t_params.dt[] < 2.0/3.0 * nl_solver_params.precon_dt[]
Expand All @@ -1347,7 +1359,6 @@ global_rank[] == 0 && println("recalculating precon")
end

if nl_solver_params.solves_since_precon_update[] ≥ nl_solver_params.preconditioner_update_interval
global_rank[] == 0 && println("recalculating precon")
nl_solver_params.solves_since_precon_update[] = 0
nl_solver_params.precon_dt[] = t_params.dt[]

Expand Down Expand Up @@ -1931,7 +1942,7 @@ to allow the outer r-loop to be parallelised.

function recalculate_preconditioner!()
if nl_solver_params.preconditioner_type === Val(:electron_lu)
global_rank[] == 0 && println("recalculating precon")
global_rank[] == 0 && println("recalculating precon :electron_lu 2")
nl_solver_params.solves_since_precon_update[] = 0
nl_solver_params.precon_dt[] = ion_dt

Expand All @@ -1944,27 +1955,37 @@ global_rank[] == 0 && println("recalculating precon")
z_advect, vpa_advect, scratch_dummy, external_source_settings,
num_diss_params, t_params, ion_dt, ir, true, :all, true, false)

begin_serial_region()
if block_rank[] == 0
if size(orig_lu) == (1, 1)
if block_rank[] >= 0
localcols = MPIQR.localcols(block_rank[], size(precon_matrix, 2),
mpiqr_blocksize, block_size[])
if size(orig_lu, 1) == 1
# Have not properly created the LU decomposition before, so
# cannot reuse it.
@timeit_debug global_timer "lu" nl_solver_params.preconditioners[ir] =
(lu(sparse(precon_matrix)), precon_matrix, input_buffer,
output_buffer)
@timeit_debug global_timer "lu" begin
mpiqrmatrix = MPIQR.MPIQRMatrix((precon_matrix[:, localcols]), size(precon_matrix);
blocksize=mpiqr_blocksize, comm=comm_block[])
mpiqrstruct = qr!(mpiqrmatrix)
nl_solver_params.preconditioners[ir] =
(mpiqrstruct, precon_matrix, input_buffer, output_buffer)
end
else
# LU decomposition was previously created. The Jacobian always
# has the same sparsity pattern, so by using `lu!()` we can
# reuse some setup.
try
@timeit_debug global_timer "lu!" lu!(orig_lu, sparse(precon_matrix); check=false)
@timeit_debug global_timer "lu!" begin
orig_lu[:, localcols] .= precon_matrix[:, localcols]
qr!(orig_lu)
end
catch e
if !isa(e, ArgumentError)
rethrow(e)
end
println("Sparsity pattern of matrix changed, rebuilding "
* " LU from scratch")
@timeit_debug global_timer "lu" orig_lu = lu(sparse(precon_matrix))
mpiqrmatrix = MPIQR.MPIQRMatrix(precon_matrix[:, localcols], size(precon_matrix);
blocksize=mpiqr_blocksize, comm=comm_block[])
@timeit_debug global_timer "lu" orig_lu = qr!(mpiqrmatrix)
end
nl_solver_params.preconditioners[ir] =
(orig_lu, precon_matrix, input_buffer, output_buffer)
Expand Down Expand Up @@ -2007,10 +2028,8 @@ global_rank[] == 0 && println("recalculating precon")
counter += 1
end

begin_serial_region()
@serial_region begin
@timeit_debug global_timer "ldiv!" ldiv!(this_output_buffer, precon_lu, this_input_buffer)
end
this_output_buffer .= this_input_buffer
@timeit_debug global_timer "ldiv!" MPIQR.ldiv!(precon_lu, this_output_buffer)

begin_serial_region()
counter = 1
Expand Down
15 changes: 9 additions & 6 deletions moment_kinetics/src/nonlinear_solvers.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ using ..type_definitions: mk_float, mk_int

using LinearAlgebra
using MPI
using MPIQR
using SparseArrays
using StatsBase: mean

Expand Down Expand Up @@ -166,23 +167,25 @@ function setup_nonlinear_solve(active, input_dict, coords, outer_coords=(); defa
V .= 0.0
end
end

mpiqr = MPIQR.MPIQRMatrix(zeros(Float64, 1, 1), (1, block_size[]);
blocksize=1, comm=comm_block[])
mpiqrstruct = MPIQR.MPIQRStruct(mpiqr)
if preconditioner_type === Val(:lu)
# Create dummy LU solver objects so we can create an array for preconditioners.
# These will be calculated properly within the time loop.
preconditioners = fill(lu(sparse(1.0*I, total_size_coords, total_size_coords)),
preconditioners = fill(mpiqrstruct,
reverse(outer_coord_sizes))
elseif preconditioner_type === Val(:electron_split_lu)
preconditioners = (z=fill(lu(sparse(1.0*I, coords.z.n, coords.z.n)),
preconditioners = (z=fill(mpiqrstruct,
tuple(coords.vpa.n, reverse(outer_coord_sizes)...)),
vpa=fill(lu(sparse(1.0*I, coords.vpa.n, coords.vpa.n)),
vpa=fill(mpqirstruct,
tuple(coords.z.n, reverse(outer_coord_sizes)...)),
ppar=fill(lu(sparse(1.0*I, coords.z.n, coords.z.n)),
ppar=fill(mpiqrstruct,
reverse(outer_coord_sizes)),
)
elseif preconditioner_type === Val(:electron_lu)
pdf_plus_ppar_size = total_size_coords + coords.z.n
preconditioners = fill((lu(sparse(1.0*I, 1, 1)),
preconditioners = fill((mpiqrstruct,
allocate_shared_float(pdf_plus_ppar_size, pdf_plus_ppar_size),
allocate_shared_float(pdf_plus_ppar_size),
allocate_shared_float(pdf_plus_ppar_size),
Expand Down