diff --git a/moment_kinetics/Project.toml b/moment_kinetics/Project.toml index fe15a78560..e8a430cc7b 100644 --- a/moment_kinetics/Project.toml +++ b/moment_kinetics/Project.toml @@ -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" @@ -50,6 +51,7 @@ file_io_netcdf = "NCDatasets" manufactured_solns_ext = ["Symbolics", "IfElse"] [compat] +MPIQR = "0.1.0" julia = "1.9.0" [extras] @@ -57,4 +59,4 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Random", "Test", "NCDatasets"] +test = ["Random", "Test"] diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl index 556490a6d3..d33ad29e6a 100644 --- a/moment_kinetics/src/electron_kinetic_equation.jl +++ b/moment_kinetics/src/electron_kinetic_equation.jl @@ -2,6 +2,7 @@ module electron_kinetic_equation using LinearAlgebra using MPI +using MPIQR using SparseArrays export get_electron_critical_velocities @@ -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 @@ -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[] @@ -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) @@ -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 @@ -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[] @@ -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[] @@ -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 @@ -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) @@ -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 diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl index 3615479705..568dba47ff 100644 --- a/moment_kinetics/src/nonlinear_solvers.jl +++ b/moment_kinetics/src/nonlinear_solvers.jl @@ -39,6 +39,7 @@ using ..type_definitions: mk_float, mk_int using LinearAlgebra using MPI +using MPIQR using SparseArrays using StatsBase: mean @@ -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),