Skip to content

Commit

Permalink
Reformat src\solvers
Browse files Browse the repository at this point in the history
  • Loading branch information
@huiyuxie
huiyuxie committed Aug 20, 2024
1 parent 2ebe44c commit dc29ce6
Show file tree
Hide file tree
Showing 6 changed files with 313 additions and 673 deletions.
2 changes: 1 addition & 1 deletion .JuliaFormatter.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,6 @@
style = "sciml"

# Additional Options
margin = 150
margin = 100
yas_style_nesting = true
align_struct_field = true
32 changes: 11 additions & 21 deletions src/TrixiGPU.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,27 +3,17 @@ module TrixiGPU
# Include other packages that are used in TrixiGPU.jl (# FIXME: Remember to reorder)
# using Reexport: @reexport

using CUDA:
@cuda, CuArray, HostKernel, threadIdx, blockIdx, blockDim, similar, launch_configuration

using Trixi:
AbstractEquations,
TreeMesh,
DGSEM,
BoundaryConditionPeriodic,
SemidiscretizationHyperbolic,
VolumeIntegralWeakForm,
VolumeIntegralFluxDifferencing,
flux,
ntuple,
nvariables,
True,
False,
wrap_array,
compute_coefficients,
have_nonconservative_terms,
set_log_type,
set_sqrt_type
using CUDA: @cuda, CuArray, HostKernel,
threadIdx, blockIdx, blockDim, similar, launch_configuration

using Trixi: AbstractEquations, TreeMesh, DGSEM,
BoundaryConditionPeriodic, SemidiscretizationHyperbolic,
VolumeIntegralWeakForm, VolumeIntegralFluxDifferencing,
flux, ntuple, nvariables,
True, False,
wrap_array, compute_coefficients,
have_nonconservative_terms,
set_log_type, set_sqrt_type

import Trixi: get_node_vars, get_node_coords, get_surface_node_vars

Expand Down
211 changes: 63 additions & 148 deletions src/solvers/dg_1d.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,8 @@ function weak_form_kernel!(du, derivative_dhat, flux_arr)
end

# Kernel for calculating volume fluxes
function volume_flux_kernel!(
volume_flux_arr, u, equations::AbstractEquations{1}, volume_flux::Function
)
function volume_flux_kernel!(volume_flux_arr, u, equations::AbstractEquations{1},
volume_flux::Function)
j = (blockIdx().x - 1) * blockDim().x + threadIdx().x
k = (blockIdx().y - 1) * blockDim().y + threadIdx().y

Expand Down Expand Up @@ -85,9 +84,7 @@ function volume_integral_kernel!(du, derivative_split, volume_flux_arr)
end

# Kernel for prolonging two interfaces
function prolong_interfaces_kernel!(
interfaces_u, u, neighbor_ids, equations::AbstractEquations{1}
)
function prolong_interfaces_kernel!(interfaces_u, u, neighbor_ids, equations::AbstractEquations{1})
j = (blockIdx().x - 1) * blockDim().x + threadIdx().x
k = (blockIdx().y - 1) * blockDim().y + threadIdx().y

Expand All @@ -105,9 +102,8 @@ function prolong_interfaces_kernel!(
end

# Kernel for calculating surface fluxes
function surface_flux_kernel!(
surface_flux_arr, interfaces_u, equations::AbstractEquations{1}, surface_flux::Any
)
function surface_flux_kernel!(surface_flux_arr, interfaces_u, equations::AbstractEquations{1},
surface_flux::Any)
k = (blockIdx().x - 1) * blockDim().x + threadIdx().x

if (k <= size(surface_flux_arr, 3))
Expand All @@ -126,9 +122,8 @@ function surface_flux_kernel!(
end

# Kernel for setting interface fluxes
function interface_flux_kernel!(
surface_flux_values, surface_flux_arr, neighbor_ids, equations::AbstractEquations{1}
)
function interface_flux_kernel!(surface_flux_values, surface_flux_arr, neighbor_ids,
equations::AbstractEquations{1})
i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
k = (blockIdx().y - 1) * blockDim().y + threadIdx().y

Expand All @@ -146,18 +141,16 @@ function interface_flux_kernel!(
end

# Kernel for calculating surface integrals
function surface_integral_kernel!(
du, factor_arr, surface_flux_values, equations::AbstractEquations{1}
)
function surface_integral_kernel!(du, factor_arr, surface_flux_values,
equations::AbstractEquations{1})
i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
j = (blockIdx().y - 1) * blockDim().y + threadIdx().y
k = (blockIdx().z - 1) * blockDim().z + threadIdx().z

if (i <= size(du, 1) && j <= size(du, 2) && k <= size(du, 3))
@inbounds begin
du[i, j, k] -= surface_flux_values[i, 1, k] * isequal(j, 1) * factor_arr[1]
du[i, j, k] +=
surface_flux_values[i, 2, k] * isequal(j, size(du, 2)) * factor_arr[2]
du[i, j, k] += surface_flux_values[i, 2, k] * isequal(j, size(du, 2)) * factor_arr[2]
end
end

Expand All @@ -182,69 +175,41 @@ end
# and run them on the device (i.e., GPU).

# Pack kernels for calculating volume integrals
function cuda_volume_integral!(
du,
u,
mesh::TreeMesh{1},
nonconservative_terms,
equations,
volume_integral::VolumeIntegralWeakForm,
dg::DGSEM,
)
function cuda_volume_integral!(du, u, mesh::TreeMesh{1}, nonconservative_terms,
equations, volume_integral::VolumeIntegralWeakForm, dg::DGSEM)
derivative_dhat = CuArray{Float64}(dg.basis.derivative_dhat)
flux_arr = similar(u)

size_arr = CuArray{Float64}(undef, size(u, 2), size(u, 3))

flux_kernel = @cuda launch = false flux_kernel!(flux_arr, u, equations, flux)
flux_kernel = @cuda launch=false flux_kernel!(flux_arr, u, equations, flux)
flux_kernel(flux_arr, u, equations, flux; configurator_2d(flux_kernel, size_arr)...)

weak_form_kernel = @cuda launch = false weak_form_kernel!(du, derivative_dhat, flux_arr)
weak_form_kernel(
du, derivative_dhat, flux_arr; configurator_3d(weak_form_kernel, du)...
)
weak_form_kernel = @cuda launch=false weak_form_kernel!(du, derivative_dhat, flux_arr)
weak_form_kernel(du, derivative_dhat, flux_arr; configurator_3d(weak_form_kernel, du)...)

return nothing
end

function cuda_volume_integral!(
du,
u,
mesh::TreeMesh{1},
nonconservative_terms::False,
equations,
volume_integral::VolumeIntegralFluxDifferencing,
dg::DGSEM,
)
function cuda_volume_integral!(du, u, mesh::TreeMesh{1}, nonconservative_terms::False,
equations, volume_integral::VolumeIntegralFluxDifferencing,
dg::DGSEM)
volume_flux = volume_integral.volume_flux

derivative_split = CuArray{Float64}(dg.basis.derivative_split)
volume_flux_arr = CuArray{Float64}(
undef, size(u, 1), size(u, 2), size(u, 2), size(u, 3)
)
volume_flux_arr = CuArray{Float64}(undef, size(u, 1), size(u, 2), size(u, 2), size(u, 3))

size_arr = CuArray{Float64}(undef, size(u, 2)^2, size(u, 3))

volume_flux_kernel = @cuda launch = false volume_flux_kernel!(
volume_flux_arr, u, equations, volume_flux
)
volume_flux_kernel(
volume_flux_arr,
u,
equations,
volume_flux;
configurator_2d(volume_flux_kernel, size_arr)...,
)

volume_integral_kernel = @cuda launch = false volume_integral_kernel!(
du, derivative_split, volume_flux_arr
)
volume_integral_kernel(
du,
derivative_split,
volume_flux_arr;
configurator_3d(volume_integral_kernel, du)...,
)
volume_flux_kernel = @cuda launch=false volume_flux_kernel!(volume_flux_arr, u, equations,
volume_flux)
volume_flux_kernel(volume_flux_arr, u, equations, volume_flux;
configurator_2d(volume_flux_kernel, size_arr)...,)

volume_integral_kernel = @cuda launch=false volume_integral_kernel!(du, derivative_split,
volume_flux_arr)
volume_integral_kernel(du, derivative_split, volume_flux_arr;
configurator_3d(volume_integral_kernel, du)...,)

return nothing
end
Expand All @@ -256,26 +221,20 @@ function cuda_prolong2interfaces!(u, mesh::TreeMesh{1}, equations, cache)

size_arr = CuArray{Float64}(undef, size(interfaces_u, 2), size(interfaces_u, 3))

prolong_interfaces_kernel = @cuda launch = false prolong_interfaces_kernel!(
interfaces_u, u, neighbor_ids, equations
)
prolong_interfaces_kernel(
interfaces_u,
u,
neighbor_ids,
equations;
configurator_2d(prolong_interfaces_kernel, size_arr)...,
)
prolong_interfaces_kernel = @cuda launch=false prolong_interfaces_kernel!(interfaces_u, u,
neighbor_ids,
equations)
prolong_interfaces_kernel(interfaces_u, u, neighbor_ids, equations;
configurator_2d(prolong_interfaces_kernel, size_arr)...,)

cache.interfaces.u = interfaces_u # copy back to host automatically

return nothing
end

# Pack kernels for calculating interface fluxes
function cuda_interface_flux!(
mesh::TreeMesh{1}, nonconservative_terms::False, equations, dg::DGSEM, cache
)
function cuda_interface_flux!(mesh::TreeMesh{1}, nonconservative_terms::False, equations, dg::DGSEM,
cache)
surface_flux = dg.surface_integral.surface_flux

neighbor_ids = CuArray{Int64}(cache.interfaces.neighbor_ids)
Expand All @@ -285,75 +244,52 @@ function cuda_interface_flux!(

size_arr = CuArray{Float64}(undef, size(interfaces_u, 3))

surface_flux_kernel = @cuda launch = false surface_flux_kernel!(
surface_flux_arr, interfaces_u, equations, surface_flux
)
surface_flux_kernel(
surface_flux_arr,
interfaces_u,
equations,
surface_flux;
configurator_1d(surface_flux_kernel, size_arr)...,
)
surface_flux_kernel = @cuda launch=false surface_flux_kernel!(surface_flux_arr, interfaces_u,
equations, surface_flux)
surface_flux_kernel(surface_flux_arr, interfaces_u, equations, surface_flux;
configurator_1d(surface_flux_kernel, size_arr)...,)

size_arr = CuArray{Float64}(undef, size(surface_flux_values, 1), size(interfaces_u, 3))

interface_flux_kernel = @cuda launch = false interface_flux_kernel!(
surface_flux_values, surface_flux_arr, neighbor_ids, equations
)
interface_flux_kernel(
surface_flux_values,
surface_flux_arr,
neighbor_ids,
equations;
configurator_2d(interface_flux_kernel, size_arr)...,
)
interface_flux_kernel = @cuda launch=false interface_flux_kernel!(surface_flux_values,
surface_flux_arr,
neighbor_ids, equations)
interface_flux_kernel(surface_flux_values, surface_flux_arr, neighbor_ids, equations;
configurator_2d(interface_flux_kernel, size_arr)...,)

cache.elements.surface_flux_values = surface_flux_values # copy back to host automatically

return nothing
end

# Dummy function for asserting boundaries
function cuda_prolong2boundaries!(
u, mesh::TreeMesh{1}, boundary_condition::BoundaryConditionPeriodic, cache
)
function cuda_prolong2boundaries!(u, mesh::TreeMesh{1},
boundary_condition::BoundaryConditionPeriodic, cache)
@assert iszero(length(cache.boundaries.orientations))
end

# Dummy function for asserting boundary fluxes
function cuda_boundary_flux!(
t,
mesh::TreeMesh{1},
boundary_condition::BoundaryConditionPeriodic,
equations,
dg::DGSEM,
cache,
)
function cuda_boundary_flux!(t, mesh::TreeMesh{1}, boundary_condition::BoundaryConditionPeriodic,
equations, dg::DGSEM, cache)
@assert iszero(length(cache.boundaries.orientations))
end

# Pack kernels for calculating surface integrals
function cuda_surface_integral!(du, mesh::TreeMesh{1}, equations, dg::DGSEM, cache)
# FIXME: Check `surface_integral`
factor_arr = CuArray{Float64}([
dg.basis.boundary_interpolation[1, 1],
dg.basis.boundary_interpolation[size(du, 2), 2],
])
dg.basis.boundary_interpolation[1, 1],
dg.basis.boundary_interpolation[size(du, 2), 2]
])
surface_flux_values = CuArray{Float64}(cache.elements.surface_flux_values)

size_arr = CuArray{Float64}(undef, size(du, 1), size(du, 2), size(du, 3))

surface_integral_kernel = @cuda launch = false surface_integral_kernel!(
du, factor_arr, surface_flux_values, equations
)
surface_integral_kernel(
du,
factor_arr,
surface_flux_values,
equations;
configurator_3d(surface_integral_kernel, size_arr)...,
)
surface_integral_kernel = @cuda launch=false surface_integral_kernel!(du, factor_arr,
surface_flux_values,
equations)
surface_integral_kernel(du, factor_arr, surface_flux_values, equations;
configurator_3d(surface_integral_kernel, size_arr)...,)

return nothing
end
Expand All @@ -362,47 +298,26 @@ end
function cuda_jacobian!(du, mesh::TreeMesh{1}, equations, cache)
inverse_jacobian = CuArray{Float64}(cache.elements.inverse_jacobian)

jacobian_kernel = @cuda launch = false jacobian_kernel!(du, inverse_jacobian, equations)
jacobian_kernel(
du, inverse_jacobian, equations; configurator_3d(jacobian_kernel, du)...
)
jacobian_kernel = @cuda launch=false jacobian_kernel!(du, inverse_jacobian, equations)
jacobian_kernel(du, inverse_jacobian, equations; configurator_3d(jacobian_kernel, du)...)

return nothing
end

# Dummy function returning nothing
function cuda_sources!(
du, u, t, source_terms::Nothing, equations::AbstractEquations{1}, cache
)
function cuda_sources!(du, u, t, source_terms::Nothing, equations::AbstractEquations{1}, cache)
return nothing
end

# Put everything together into a single function
# Ref: `rhs!` function in Trixi.jl

function rhs_gpu!(
du_cpu,
u_cpu,
t,
mesh::TreeMesh{1},
equations,
initial_condition,
boundary_conditions,
source_terms::Source,
dg::DGSEM,
cache,
) where {Source}
function rhs_gpu!(du_cpu, u_cpu, t, mesh::TreeMesh{1}, equations, initial_condition,
boundary_conditions, source_terms::Source, dg::DGSEM, cache) where {Source}
du, u = copy_to_device!(du_cpu, u_cpu)

cuda_volume_integral!(
du,
u,
mesh,
have_nonconservative_terms(equations),
equations,
dg.volume_integral,
dg,
)
cuda_volume_integral!(du, u, mesh, have_nonconservative_terms(equations), equations,
dg.volume_integral, dg)

cuda_prolong2interfaces!(u, mesh, equations, cache)

Expand Down
Loading

0 comments on commit dc29ce6

Please sign in to comment.