Multithreaded simulations

Overview

Multithreaded simulations are possible via use of standard Julia threads. The approach is relatively straightforward:

  1. The domain is decomposed into n_chunks "cell chunks", which are lists of cells. Within each chunk the cell indices should be continuous. For example, a 4-cell domain may be decomposed into [[1,2],[3,4]] (2 chunks), [[1], [2,3], [4] (3 chunks), [[1], [2], [3], [4]] (4 chunks), etc. A LoadBalancerCellQ instance can be used to perform both the initial decomposition and on-the-fly decomposition with load balancing; it stores the chunked indices as ranges of cell indices, i.e. [1:2, 3:4] for the first example above.
  2. Several variables are instantiated for each chunk:
    • n_chunks arrays of ParticleVectors (each hold n_species ParticleVectors)
    • n_chunks RNGs
    • n_chunks ParticleIndexerArray instances
    • n_chunks GridSortInPlace instances
    • n_chunks+1 SurfProps instances (one additional instance is required for the reduce operation)
  3. Collisions, convection, sorting are all performed per-chunk, completely independently, so multithreading can be used
  4. After particles have been sorted, some might need to be moved between chunks, if they ended up in cells assigned to another chunk. This is done in two steps: first, one calls exchange_particles! to move the particle data between chunks. This is can be done either in serial mode or in threaded mode, see below for details.
  5. Then, sort_particles_after_exchange! is called to reset indexing without having to completely resort all the new particles. This can be done using multithreading. It also clears the entries of the ChunkExchanger as it reads them, so the exchanger is ready for the next timestep and does not need to be reset! explicitly.
  6. Physical grid properties are computed using multithreading, as they can easily be computed only for cells assigned to the chunk, thus avoiding any race conditions. In case surface properties were computed during particle movement, a reduce operation is needed to sum the per-chunk values. This is done by a serial call to reduce_surf_props!.

The trade-off of the approach is that even in case no new particles are created in the simulation, particle data will be moved around, as when after a movement step a particle ends up in a cell corresponding to a different chunk, it needs to be moved to the ParticleVector for that chunk. In addition, some duplication of data structures (i.e. multiple GridSortInPlace instances) is required. A special struct ChunkExchanger is used to facilitate exchange of particles between chunks.

However, the code logic is more straightforward, as most of the data is independent and no race conditions can occur. Some of this functionality will be re-used to make MPI simulations possible.

Serial vs parallel particle exchange

Particle exchange needs to be done carefully to avoid race conditions. For the serial version, the call looks like this:

exchange_particles!(chunk_exchanger, particles_chunks, pia_chunks, cell_chunks, 1)

The function exchange_particles! contains two loops over the chunks: one outer loop for i in 1:n_chunks-1 and one inner loop for j in i+1:n_chunks.

The equivalent threaded version looks like this:

for exchange_partial in exchange_list    @threads for exchange_pair in exchange_partial        exchange_particles!(chunk_exchanger, particles_chunks, pia_chunks, cell_chunks,                             1, exchange_pair[1], exchange_pair[2])    endend

Here exchange_list is a fixed list of chunk pairs that allows for thread-safe particle exchange; it is pre-computed by calling generate_1_factorization before the start of the collision loop: exchange_list = generate_1_factorization(n_chunks). The version of the exchange_particles! called here exchanges particles only between two specific chunks: exchange_pair[1] and exchange_pair[2]. In fact, the serial version of the function simply calls this exchange-between-two-chunks function for every (i,j) pair inside the double loop described above. It should be noted that threaded particle exchange is not necessarily faster than the serial exchange, due to the overhead of threading. Example multi-threaded simulations in the simulations directory have an optional parameter parallel_exchange that can be used to toggle serial and parallel particle exchange to see which is faster for a given problem.

Example: multithreaded Couette flow simulation with load balancing

Below is an example of a fixed-weight DSMC multithreaded Couette flow simulation. To run in multithreaded mode, one should start julia specifying the number of threads: julia --threads NTHREADS. The file can also be found under simulations/1D/couette_multithreaded.jl. A variable-weight example can be found under simulations/1D/couette_multithreaded_varweight_octree.jl.

The ChunkSplitters.jl library is used internally to perform domain decomposition, by splitting the range of cell indices 1:nx into independent chunks. By default, the number of chunks is set equal to the number of threads; it can be set to a multiple of the number of threads by setting a value of the chunk_count_multiplier parameter. preallocation_margin_multiplier allocates additional unused particles in the per-chunk ParticleVector instances, since otherwise the transfer of particles might lead to frequent calls to resize! at the start of the simulation as the solution approaches steady state and the average number of particles in a chunk changes significantly. A value of 1.0 means no additional particle storage is allocated. We perform the domain decomposition into chunks indirectly by instantiating a LoadBalancerCellQ object, which is also used to perform load balancing (see below).

The surface properties are collected into surf_props_reduced via a call to reduce_surf_props!. Before the start of the time loop, the sampling procedure is multithreaded via the @threads macro. The physical properties are also computed in multithreaded mode.

Inside the time loop, collisions, convection, and sorting are performed inside a @threads block. Once the block finishes, the particles are moved between chunks. To improve the speed of this procedure, first the update_occupancy_bounds! is called, which takes in the grid sorting structure from the sorting step and reads the upper and lower bounds of the cells in which the chunk holds particles, including those outside of the range of cells owned by the chunk, to which particles have moved during the convection step. The bounds are recorded by the sorting step itself, so this is a constant-time operation. Therefore, one can avoid scanning across all cells in the grid in the exchange step, but one can also set the lower bound to 1 and upper bound to n_cells if the occupancy bound computation step is not performed. Next, a serial call to exchange_particles! is used. Note that the occupancy bounds computed in the sorting step are not species-specific, i.e. first sorting multiple species and only then updating occupancy bounds will lead to errors, as the occupancy data will be overwritten by the data corresponding to the last species.

After the particles have been exchanged, they are sorted via a threaded call to sort_particles_after_exchange!. Finally, the indexing is reset, and physical grid properties are computed, again inside a @threads block. The reduction operation for the surface properties, as well as averaging of grid and surface properties, is performed serially at the end of the timestep.

The LoadBalancerCellQ performs load-balancing by tracking a per-cell quantity (i.e. number of collisions, number of particles, density, etc.) and rebalancing the assignment of cells to chunks so that each chunk contains approximately the same total quantity. The update_lb_cellq! function can be used to update the per-cell quantity in the load balancer, and the rebalance_lb! function can be used to rebalance the assignment of cells to chunks based on that quantity. reset_lb! then resets the tracking of the per-cell quantity to 0, if required to do so. In the example below, the number of collisions performed in a cell is used to perform load-balancing, and the load balancer is updated every 5000 timesteps.

using Merzbildusing Randomusing TimerOutputsusing Base.Threadsusing ChunkSplittersfunction run(seed, T_wall, v_wall, L, ndens, nx, ppc, Δt, n_timesteps, avg_start; chunk_count_multiplier=1,             preallocation_margin_multiplier=1.0)    reset_timer!()    n_threads = Threads.nthreads()    n_chunks = n_threads * chunk_count_multiplier    println("Running on $n_threads threads, will split cells into $n_chunks chunks")    rng_chunks = [Xoshiro(seed + i) for i in 0:n_chunks-1]    # load particle and interaction data    particles_data_path = joinpath(MERZBILD_DATA_PATH, "particles.toml")    species_data = load_species_data(particles_data_path, "Ar")    interaction_data_path = joinpath(MERZBILD_DATA_PATH, "vhs.toml")    interaction_data::Array{Interaction, 2} = load_interaction_data(interaction_data_path, species_data)    # create our grid and BCs    grid = Grid1DUniform(L, nx)    bc_list = (FullyDiffuseBC1D(1, species_data, T_wall, [0.0, -v_wall, 0.0]),               FullyDiffuseBC1D(1, species_data, T_wall, [0.0, v_wall, 0.0]))    # split cell indices into chunks    lbq = LoadBalancerCellQ(nx, n_chunks)    cell_chunks = lbq.chunked_indices    # init per-chunk particle vectors, particle indexers, grid particle sorters    n_particles_chunks = [floor(Int64, ppc * length(cell_chunk) * preallocation_margin_multiplier) for cell_chunk in cell_chunks]    particles_chunks = [[ParticleVector(n_particles)] for n_particles in n_particles_chunks]    pia_chunks = [ParticleIndexerArray(grid.n_cells, 1) for cell_chunk in cell_chunks]    gridsorter_chunks = [GridSortInPlace(grid, n_particles) for n_particles in n_particles_chunks]    # this is used for moving particles between chunks after they have been sorted into grid cells    chunk_exchanger = ChunkExchanger(cell_chunks, nx)    # sample particles    # Fnum * ppc = Np in cell = ndens * V_cell    Fnum = grid.cells[1].V * ndens / ppc    # sample particles per-chunk    @timeit "sampling" @threads for chunk_id in 1:n_chunks        @inbounds sample_particles_equal_weight!(rng_chunks[chunk_id], grid, particles_chunks[chunk_id][1],                                                    pia_chunks[chunk_id],                                                    1, species_data, ndens, T_wall, Fnum, cell_chunks[chunk_id])    end         # create collision structs    collision_data = [CollisionData() for cell_chunk in cell_chunks]        # create struct for computation of physical properties, sizes of pia are the same    phys_props = PhysProps(pia_chunks[1])    # create second struct for averaging of physical properties, sizes of pia are the same    phys_props_avg = PhysProps(pia_chunks[1])    # create struct for computation of surface properties, need a SurfProps instance per chunk    surf_props_chunks = [SurfProps(pia_chunks[1], grid) for cell_chunk in cell_chunks]    # we sum up all the surf props here    surf_props_reduced = SurfProps(pia_chunks[1], grid)    # create second struct for averaging of physical properties    surf_props_avg = SurfProps(pia_chunks[1], grid)    # create struct for netCDF for time-averaged physical properties I/O    ds_avg = NCDataHolder("scratch/data/avg_mt_couette_$(L)_$(nx)_$(v_wall)_$(T_wall)_$(ppc)_after$(avg_start).nc",                          species_data, phys_props)    # create struct for netCDF for time-averaged surface properties I/O    ds_surf_avg = NCDataHolderSurf("scratch/data/avg_mt_couette_$(L)_$(nx)_$(v_wall)_$(T_wall)_$(ppc)_surf_after$(avg_start).nc",                                   species_data, surf_props_avg)    # create and estimate collision factors    collision_factors = [create_collision_factors_array(pia, interaction_data, species_data, T_wall, Fnum)                         for pia in pia_chunks]    # compute data at t=0    @timeit "props compute" @threads for chunk_id in 1:n_chunks        compute_props_sorted!(particles_chunks[chunk_id], pia_chunks[chunk_id], species_data, phys_props, cell_chunks[chunk_id])    end    n_avg = n_timesteps - avg_start + 1    for t in 1:n_timesteps        if t % 500 == 0            println(t)        end                # collide, convect, sort particles        @timeit "collide+convect+sort" @threads for chunk_id in 1:n_chunks            for cell in cell_chunks[chunk_id]                @inbounds ntc!(rng_chunks[chunk_id], collision_factors[chunk_id][1, 1, cell],                               collision_data[chunk_id], interaction_data, particles_chunks[chunk_id][1],                               pia_chunks[chunk_id], cell, 1, Δt, grid.cells[cell].V)                update_lb_cellq!(lbq, chunk_id, cell, collision_factors[chunk_id][1, 1, cell].n_coll_performed; averaging_window=1.0)            end            if (t >= avg_start)                @inbounds convect_particles!(rng_chunks[chunk_id], grid, bc_list,                                    particles_chunks[chunk_id][1], pia_chunks[chunk_id],                                    1, species_data, surf_props_chunks[chunk_id], Δt)            else                # we don't need to compute surface properties before we start averaging                @inbounds convect_particles!(rng_chunks[chunk_id], grid, bc_list,                                    particles_chunks[chunk_id][1], pia_chunks[chunk_id],                                    1, species_data, Δt)            end            # sort particles            @inbounds sort_particles!(gridsorter_chunks[chunk_id], grid, particles_chunks[chunk_id][1], pia_chunks[chunk_id], 1)            # tell the exchanger which cells this chunk holds particles in, so that pairs of            # chunks with nothing to exchange are rejected without scanning any cells            @inbounds update_occupancy_bounds!(chunk_exchanger, gridsorter_chunks[chunk_id], pia_chunks[chunk_id], chunk_id, 1)        end        if t%5000 == 0            @timeit "rebalance" rebalance_lb!(lbq)            @timeit "rebalance" reset_lb!(lbq)        end        # move particles between chunks        @timeit "exchange" exchange_particles!(chunk_exchanger, particles_chunks, pia_chunks, cell_chunks, 1)        # reset indexing, compute physical properties if needed        @timeit "re-sort + compute props" @threads for chunk_id in 1:n_chunks            sort_particles_after_exchange!(chunk_exchanger, gridsorter_chunks[chunk_id],                                           particles_chunks[chunk_id][1], pia_chunks[chunk_id],                                           cell_chunks[chunk_id], 1)            if (t >= avg_start)                @inbounds compute_props_sorted!(particles_chunks[chunk_id], pia_chunks[chunk_id],                                                species_data, phys_props, cell_chunks[chunk_id])            end        end        # reduce surface properties, average grid and surface properties        if (t >= avg_start)            @timeit "avg physprops" avg_props!(phys_props_avg, phys_props, n_avg)            @timeit "reduce surf props" reduce_surf_props!(surf_props_reduced, surf_props_chunks)            @timeit "avg surfprops" avg_props!(surf_props_avg, surf_props_reduced, n_avg)        end    end    @timeit "I/O" write_netcdf(ds_avg, phys_props_avg, n_timesteps)    @timeit "I/O" write_netcdf(ds_surf_avg, surf_props_avg, n_timesteps)    close_netcdf(ds_avg)    close_netcdf(ds_surf_avg)    # print out how many cells in each chunk    for i in 1:n_chunks        println("chunk $i has $(cell_chunks[i][end] - cell_chunks[i][1] + 1) cells ($(cell_chunks[i][1]):$(cell_chunks[i][end]))")    end    print_timer()endrun(1234, 300.0, 500.0, 5e-4, 5e22, 100, 100, 2.59e-9, 50000, 14000; chunk_count_multiplier=1, preallocation_margin_multiplier=1.5)