Skip to content

Commit 4d33b11

Browse files
vchuravyvchuravy
committed
port integrate_via_indices to the GPU
1 parent 40264c1 commit 4d33b11

2 files changed

Lines changed: 111 additions & 42 deletions

File tree

‎src/callbacks_step/analysis_dg2d.jl‎

Lines changed: 58 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -338,24 +338,25 @@ function integrate_via_indices(func::Func, u,
338338
return integral
339339
end
340340

341-
function integrate_via_indices(func::Func, _u,
341+
function integrate_via_indices(func::Func, u,
342342
mesh::Union{StructuredMesh{2}, StructuredMeshView{2},
343343
UnstructuredMesh2D,
344344
P4estMesh{2}, P4estMeshView{2},
345345
T8codeMesh{2}},
346346
equations, dg::DGSEM, cache,
347347
args...; normalize = true) where {Func}
348-
# TODO GPU AnalysiCallback currently lives on CPU
349-
backend = trixi_backend(_u)
350-
if backend isa Nothing # TODO GPU KA CPU backend
351-
@unpack weights = dg.basis
352-
@unpack inverse_jacobian = cache.elements
353-
u = _u
354-
else
355-
weights = Array(dg.basis.weights)
356-
inverse_jacobian = Array(cache.elements.inverse_jacobian)
357-
u = Array(_u)
358-
end
348+
return integrate_via_indices(func, trixi_backend(u), u, mesh, equations, dg, cache, args...; normalize = normalize)
349+
end
350+
351+
function integrate_via_indices(func::Func, ::Nothing, u,
352+
mesh::Union{StructuredMesh{2}, StructuredMeshView{2},
353+
UnstructuredMesh2D,
354+
P4estMesh{2}, P4estMeshView{2},
355+
T8codeMesh{2}},
356+
equations, dg::DGSEM, cache,
357+
args...; normalize = true) where {Func}
358+
@unpack weights = dg.basis
359+
@unpack inverse_jacobian = cache.elements
359360

360361
# Initialize integral with zeros of the right shape
361362
integral = zero(func(u, 1, 1, 1, equations, dg, args...))
@@ -380,6 +381,50 @@ function integrate_via_indices(func::Func, _u,
380381
return integral
381382
end
382383

384+
function integrate_via_indices(func::Func, backend::Backend, u,
385+
mesh::Union{StructuredMesh{2}, StructuredMeshView{2},
386+
UnstructuredMesh2D,
387+
P4estMesh{2}, P4estMeshView{2},
388+
T8codeMesh{2}},
389+
equations, dg::DGSEM, cache,
390+
args...; normalize = true) where {Func}
391+
@unpack weights = dg.basis
392+
@unpack inverse_jacobian = cache.elements
393+
394+
function local_plus((integral, total_volume), (integral_element, volume_element))
395+
return (integral + integral_element, total_volume + volume_element)
396+
end
397+
398+
# `func(u, 1,1,1, equations, dg, args...)` might access GPU memory
399+
# so we have to rely on the compiler to correctly infer the type of the integral here.
400+
# TODO: Technically we need device_promote_op here that "infers" the function within the context of the GPU.
401+
integral₀ = zero(Base.promote_op(func, typeof(u), Int, Int, Int, typeof(equations), typeof(dg), map(typeof, args)...))
402+
init = neutral = (integral₀, zero(real(mesh)))
403+
404+
# Use quadrature to numerically integrate over entire domain
405+
num_elements = nelements(dg, cache)
406+
_integral, _total_volume = AcceleratedKernels.mapreduce(local_plus, 1:num_elements, backend; init, neutral) do element
407+
# Initialize integral with zeros of the right shape
408+
integral, total_volume = neutral
409+
410+
for j in eachnode(dg), i in eachnode(dg)
411+
volume_jacobian = abs(inv(inverse_jacobian[i, j, element]))
412+
integral += volume_jacobian * weights[i] * weights[j] *
413+
func(u, i, j, element, equations, dg, args...)
414+
total_volume += volume_jacobian * weights[i] * weights[j]
415+
end
416+
417+
return (integral, total_volume)
418+
end
419+
420+
# Normalize with total volume
421+
if normalize
422+
_integral = _integral / _total_volume
423+
end
424+
425+
return _integral
426+
end
427+
383428
function integrate(func::Func, u,
384429
mesh::Union{TreeMesh{2}, StructuredMesh{2}, StructuredMeshView{2},
385430
UnstructuredMesh2D, P4estMesh{2}, P4estMeshView{2},
@@ -410,18 +455,10 @@ function integrate(func::Func, u,
410455
end
411456
end
412457

413-
function analyze(::typeof(entropy_timederivative), _du, u, t,
458+
function analyze(::typeof(entropy_timederivative), du, u, t,
414459
mesh::Union{TreeMesh{2}, StructuredMesh{2}, StructuredMeshView{2},
415460
UnstructuredMesh2D, P4estMesh{2}, T8codeMesh{2}},
416461
equations, dg::Union{DGSEM, FDSBP}, cache)
417-
# TODO GPU AnalysiCallback currently lives on CPU
418-
backend = trixi_backend(u)
419-
if backend isa Nothing # TODO GPU KA CPU backend
420-
du = _du
421-
else
422-
du = Array(_du)
423-
end
424-
425462
# Calculate
426463
# Calculate ∫(∂S/∂u ⋅ ∂u/∂t)dΩ
427464
integrate_via_indices(u, mesh, equations, dg, cache,

‎src/callbacks_step/analysis_dg3d.jl‎

Lines changed: 53 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -389,22 +389,21 @@ function integrate_via_indices(func::Func, u,
389389
return integral
390390
end
391391

392-
function integrate_via_indices(func::Func, _u,
392+
function integrate_via_indices(func::Func, u,
393393
mesh::Union{StructuredMesh{3}, P4estMesh{3},
394394
T8codeMesh{3}},
395395
equations, dg::DGSEM, cache,
396396
args...; normalize = true) where {Func}
397-
# TODO GPU AnalysiCallback currently lives on CPU
398-
backend = trixi_backend(_u)
399-
if backend isa Nothing # TODO GPU KA CPU backend
400-
@unpack weights = dg.basis
401-
@unpack inverse_jacobian = cache.elements
402-
u = _u
403-
else
404-
weights = Array(dg.basis.weights)
405-
inverse_jacobian = Array(cache.elements.inverse_jacobian)
406-
u = Array(_u)
407-
end
397+
return integrate_via_indices(func, trixi_backend(u), u, mesh, equations, dg, cache, args...; normalize = normalize)
398+
end
399+
400+
function integrate_via_indices(func::Func, ::Nothing, u,
401+
mesh::Union{StructuredMesh{3}, P4estMesh{3},
402+
T8codeMesh{3}},
403+
equations, dg::DGSEM, cache,
404+
args...; normalize = true) where {Func}
405+
@unpack weights = dg.basis
406+
@unpack inverse_jacobian = cache.elements
408407

409408
# Initialize integral with zeros of the right shape
410409
integral = zero(func(u, 1, 1, 1, 1, equations, dg, args...))
@@ -429,6 +428,47 @@ function integrate_via_indices(func::Func, _u,
429428
return integral
430429
end
431430

431+
function integrate_via_indices(func::Func, backend::Backend, u,
432+
mesh::Union{StructuredMesh{3}, P4estMesh{3},
433+
T8codeMesh{3}},
434+
equations, dg::DGSEM, cache,
435+
args...; normalize = true) where {Func}
436+
@unpack weights = dg.basis
437+
@unpack inverse_jacobian = cache.elements
438+
439+
function local_plus((integral, total_volume), (integral_element, volume_element))
440+
return (integral + integral_element, total_volume + volume_element)
441+
end
442+
443+
# `func(u, 1,1,1,1, equations, dg, args...)` might access GPU memory
444+
# so we have to rely on the compiler to correctly infer the type of the integral here.
445+
# TODO: Technically we need device_promote_op here that "infers" the function within the context of the GPU.
446+
integral₀ = zero(Base.promote_op(func, typeof(u), Int, Int, Int, Int, typeof(equations), typeof(dg), map(typeof, args)...))
447+
init = neutral = (integral₀, zero(real(mesh)))
448+
449+
# Use quadrature to numerically integrate over entire domain
450+
num_elements = nelements(dg, cache)
451+
_integral, _total_volume = AcceleratedKernels.mapreduce(local_plus, 1:num_elements, backend; init, neutral) do element
452+
# Initialize integral with zeros of the right shape
453+
integral, total_volume = neutral
454+
455+
for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
456+
volume_jacobian = abs(inv(inverse_jacobian[i, j, k, element]))
457+
integral += volume_jacobian * weights[i] * weights[j] * weights[k] *
458+
func(u, i, j, k, element, equations, dg, args...)
459+
total_volume += volume_jacobian * weights[i] * weights[j] * weights[k]
460+
end
461+
return integral, total_volume
462+
end
463+
464+
# Normalize with total volume
465+
if normalize
466+
_integral = _integral / _total_volume
467+
end
468+
469+
return _integral
470+
end
471+
432472
function integrate(func::Func, u,
433473
mesh::Union{TreeMesh{3}, StructuredMesh{3}, P4estMesh{3},
434474
T8codeMesh{3}},
@@ -460,18 +500,10 @@ function integrate(func::Func, u,
460500
end
461501
end
462502

463-
function analyze(::typeof(entropy_timederivative), _du, u, t,
503+
function analyze(::typeof(entropy_timederivative), du, u, t,
464504
mesh::Union{TreeMesh{3}, StructuredMesh{3}, P4estMesh{3},
465505
T8codeMesh{3}},
466506
equations, dg::Union{DGSEM, FDSBP}, cache)
467-
# TODO GPU AnalysiCallback currently lives on CPU
468-
backend = trixi_backend(u)
469-
if backend isa Nothing # TODO GPU KA CPU backend
470-
du = _du
471-
else
472-
du = Array(_du)
473-
end
474-
475507
# Calculate ∫(∂S/∂u ⋅ ∂u/∂t)dΩ
476508
integrate_via_indices(u, mesh, equations, dg, cache,
477509
du) do u, i, j, k, element, equations, dg, du

0 commit comments

Comments
 (0)