Use Adapt.jl to change storage and element type #2212

vchuravy · 2024-12-17T16:43:17Z

In order to eventually support GPU computation we need
to use Adapt.jl to allow GPU backend packages to swap
out host-array types like CuArray with device-side types
like CuDeviceArray.

Additionally this will allow us to change the element type
of a simulation by using adapt(Array{Float32}.

Co-authored-by: Lars Christmann [email protected]
Co-authored-by: Benedict Geihe [email protected]

github-actions · 2024-12-17T16:43:30Z

Review checklist

This checklist is meant to assist creators of PRs (to let them know what reviewers will typically look for) and reviewers (to guide them in a structured review process). Items do not need to be checked explicitly for a PR to be eligible for merging.

Purpose and scope

The PR has a single goal that is clear from the PR title and/or description.
All code changes represent a single set of modifications that logically belong together.
No more than 500 lines of code are changed or there is no obvious way to split the PR into multiple PRs.

Code quality

The code can be understood easily.
Newly introduced names for variables etc. are self-descriptive and consistent with existing naming conventions.
There are no redundancies that can be removed by simple modularization/refactoring.
There are no leftover debug statements or commented code sections.
The code adheres to our conventions and style guide, and to the Julia guidelines.

Documentation

New functions and types are documented with a docstring or top-level comment.
Relevant publications are referenced in docstrings (see example for formatting).
Inline comments are used to document longer or unusual code sections.
Comments describe intent ("why?") and not just functionality ("what?").
If the PR introduces a significant change or new feature, it is documented in NEWS.md with its PR number.

Testing

The PR passes all tests.
New or modified lines of code are covered by tests.
New or modified tests run in less then 10 seconds.

Performance

There are no type instabilities or memory allocations in performance-critical parts.
If the PR intent is to improve performance, before/after time measurements are posted in the PR.

Verification

The correctness of the code was verified using appropriate tests.
If new equations/methods are added, a convergence test has been run and the results
are posted in the PR.

Created with ❤️ by the Trixi.jl community.

src/semidiscretization/semidiscretization_hyperbolic.jl

src/Trixi.jl

src/solvers/dgsem_p4est/containers_parallel.jl

src/solvers/dgsem_p4est/containers.jl

ranocha

Should be close to final when tests are added

In order to eventually support GPU computation we need to use Adapt.jl to allow GPU backend packages to swap out host-array types like `CuArray` with device-side types like `CuDeviceArray`. Additionally this will allow us to change the element type of a simulation by using `adapt(Array{Float32}`. Co-authored-by: Lars Christmann <[email protected]> Co-authored-by: Benedict Geihe <[email protected]>

codecov · 2025-02-03T16:26:02Z

Codecov Report

Attention: Patch coverage is 37.90323% with 77 lines in your changes missing coverage. Please review.

Project coverage is 95.86%. Comparing base (24107f4) to head (4bf0864).

Files with missing lines	Patch %	Lines
src/solvers/dgsem_p4est/containers.jl	32.86%	47 Missing ⚠️
src/solvers/dgsem_p4est/containers_parallel.jl	20.83%	19 Missing ⚠️
src/auxiliary/containers.jl	14.29%	6 Missing ⚠️
src/auxiliary/vector_of_arrays.jl	50.00%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2212      +/-   ##
==========================================
- Coverage   96.42%   95.86%   -0.56%     
==========================================
  Files         490      491       +1     
  Lines       39426    39518      +92     
==========================================
- Hits        38015    37881     -134     
- Misses       1411     1637     +226

Flag	Coverage Δ
unittests	`95.86% <37.90%> (-0.56%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

ranocha

Thanks! Could you please adapt the t8code stuff as well? It causes the remaining failing CI tests like

  LoadError: MethodError: Cannot `convert` an object of type Vector{Vector{Int64}} to an object of type Trixi.VecOfArrays{Vector{Int64}}
  
  Closest candidates are:
    convert(::Type{T}, ::T) where T
     @ Base Base.jl:84
    (::Type{Trixi.VecOfArrays{T}} where T<:AbstractArray)(::Any)
     @ Trixi ~/work/Trixi.jl/Trixi.jl/src/auxiliary/vector_of_arrays.jl:12
  
  Stacktrace:
    [1] setproperty!(x::Trixi.P4estMPICache{Vector{Float64}, Vector{Int64}}, f::Symbol, v::Vector{Vector{Int64}})
      @ Base ./Base.jl:40
    [2] init_mpi_cache!(mpi_cache::Trixi.P4estMPICache{Vector{Float64}, Vector{Int64}}, mesh::Trixi.T8codeMesh{2, Float64, Static.True, 4, 4}, mpi_mesh_info::@NamedTuple{mpi_mortars::Trixi.P4estMPIMortarContainer{2, Float64, Float64, 3, 4, 5, Array{Float64, 5}, Vector{Float64}, Array}, mpi_interfaces::Trixi.P4estMPIInterfaceContainer{2, Float64, 4, Array{Float64, 4}, Vector{Int64}, Vector{Tuple{Symbol, Symbol}}, Vector{Float64}, Array}, global_mortar_ids::Vector{UInt128}, global_interface_ids::Vector{UInt128}, neighbor_ranks_mortar::Vector{Vector{Int64}}, neighbor_ranks_interface::Vector{Int64}}, nvars::Int64, nnodes::Int64, uEltype::Type)
      @ Trixi ~/work/Trixi.jl/Trixi.jl/src/solvers/dgsem_t8code/dg_parallel.jl:93

ranocha · 2025-02-04T04:54:06Z

src/solvers/dgsem_unstructured/sort_boundary_conditions.jl

-mutable struct UnstructuredSortedBoundaryTypes{N, BCs <: NTuple{N, Any}}
+mutable struct UnstructuredSortedBoundaryTypes{N, BCs <: NTuple{N, Any},
+                                               Vec <: AbstractVector{<:Integer}}
    boundary_condition_types::BCs # specific boundary condition type(s), e.g. BoundaryConditionDirichlet
-    boundary_indices::NTuple{N, Vector{Int}} # integer vectors containing global boundary indices
+    boundary_indices::NTuple{N, Vec} # integer vectors containing global boundary indices
    boundary_dictionary::Dict{Symbol, Any} # boundary conditions as set by the user in the elixir file
    boundary_symbol_indices::Dict{Symbol, Vector{Int}} # integer vectors containing global boundary indices per boundary identifier
 end


Aqua.jl complains about unbound type parameters here, see https://github.com/trixi-framework/Trixi.jl/actions/runs/13117124962/job/36594056579?pr=2212#step:7:6391

Ja, das verstehe ich nicht ganz.

This looks like JuliaTesting/Aqua.jl#86.

See also JuliaTesting/Aqua.jl#265 (comment), which is basically exactly the case here.

Hm I tried a couple of ways, but I am not sure how to fix this false positive.

If there is no way, you can set unbound_args = false for Aqua.jl tests so that we can proceed.

sloede

Just a few minor comments/suggestions, then this should be ready to merge (once the other convo by Hendrik is resolved as well)

sloede · 2025-02-06T09:32:53Z

src/auxiliary/containers.jl

+# The first type parameter must be `Array` by default and if 
+# `Adapt.adapt_structure(to, container)` is called then this must be
+# `typeof(to)`. This is used in downstream code, e.g. for calling
+# `wrap_array` with an appropriate type after `resize!`ing.


I don't understand this comment, in particular

The first type parameter must be Array by default and if Adapt.adapt_structure(to, container) is called then this must be typeof(to).

What is "this" referring to? To me, the sentence roughly reads like "the first type parameter must by Array and Array must always be the result of typeof(to)", which is very likely not the actual intent 😅

Could you please clarify a little?

sloede · 2025-02-06T09:34:26Z

src/auxiliary/containers.jl

+function unsafe_wrap_or_alloc(to, vec, size)
+    if length(vec) == 0
+        return similar(vec, size)
+    else
+        return unsafe_wrap(to, pointer(vec), size)
+    end
+end


Suggested change

function unsafe_wrap_or_alloc(to, vec, size)

if length(vec) == 0

return similar(vec, size)

else

return unsafe_wrap(to, pointer(vec), size)

end

end

function unsafe_wrap_or_alloc(to, vector, size)

if length(vector) == 0

return similar(vector, size)

else

return unsafe_wrap(to, pointer(vector), size)

end

end

To not shadow Base.vec?

I fixed the suggestion above. It's not necessary, from my point of view, to avoid shadowing here, but it's fine to be explicit and verbose.

sloede · 2025-02-06T09:36:13Z

src/auxiliary/vector_of_arrays.jl

+struct VecOfArrays{T <: AbstractArray}
+    arrays::Vector{T}
+end
+Base.getindex(v::VecOfArrays, i::Int) = Base.getindex(v.arrays, i)
+Base.IndexStyle(v::VecOfArrays) = Base.IndexStyle(v.arrays)
+Base.size(v::VecOfArrays) = Base.size(v.arrays)
+Base.length(v::VecOfArrays) = Base.length(v.arrays)
+Base.eltype(v::VecOfArrays{T}) where {T} = T
+function Adapt.adapt_structure(to, v::VecOfArrays)
+    return VecOfArrays([Adapt.adapt(to, arr) for arr in v.arrays])
+end
+function Base.convert(::Type{<:VecOfArrays}, v::Vector{<:AbstractArray})
+    VecOfArrays(v)
+end


Suggested change

struct VecOfArrays{T <: AbstractArray}

arrays::Vector{T}

end

Base.getindex(v::VecOfArrays, i::Int) = Base.getindex(v.arrays, i)

Base.IndexStyle(v::VecOfArrays) = Base.IndexStyle(v.arrays)

Base.size(v::VecOfArrays) = Base.size(v.arrays)

Base.length(v::VecOfArrays) = Base.length(v.arrays)

Base.eltype(v::VecOfArrays{T}) where {T} = T

function Adapt.adapt_structure(to, v::VecOfArrays)

return VecOfArrays([Adapt.adapt(to, arr) for arr in v.arrays])

end

function Base.convert(::Type{<:VecOfArrays}, v::Vector{<:AbstractArray})

VecOfArrays(v)

end

struct VectorOfArrays{T <: AbstractArray}

arrays::Vector{T}

end

Base.getindex(v::VectorOfArrays, i::Int) = Base.getindex(v.arrays, i)

Base.IndexStyle(v::VectorOfArrays) = Base.IndexStyle(v.arrays)

Base.size(v::VectorOfArrays) = Base.size(v.arrays)

Base.length(v::VectorOfArrays) = Base.length(v.arrays)

Base.eltype(v::VectorOfArrays{T}) where {T} = T

function Adapt.adapt_structure(to, v::VectorOfArrays)

return VectorOfArrays([Adapt.adapt(to, arr) for arr in v.arrays])

end

function Base.convert(::Type{<:VectorOfArrays}, v::Vector{<:AbstractArray})

VectorOfArrays(v)

end

We try to avoid abbreviations unless really necessary for sanity 😬 Also, it would be consistent with the file name.

Note, however, that I do not have a super strong opinion here, thus if you deem VecOfArrays to be better, we can stay like this.

I just saw that RecursiveArrayTools.jl seems to already have a VectorOfArrays type (https://github.com/trixi-framework/Trixi.jl/pull/2150/files#diff-6058d487915aa254cfe1a5f8da6315b8c1fef4da04df38b0089a599803e28e5aR63). Should we use that one instead?

If not, then of course we should use a different name for our own data structure, and keeping VecOfArrays as the name seems reasonable to me. In that case, it might be good to leave a short note in the comment why we do not want to use the SciML package's type for future reference.

using RecursiveArrayTools: VectorOfArray is added in #2150 for DGMulti support

Yeah, I saw that. I was just wondering whether their data type could (or should) be used instead of rolling our own.

sloede · 2025-02-06T09:38:53Z

src/solvers/dgsem_p4est/containers.jl

+    elements.node_coordinates = unsafe_wrap(ArrayType, pointer(_node_coordinates),
                                            (n_dims, ntuple(_ -> n_nodes, n_dims)...,
                                             capacity))


Why don't we need unsafe_wrap_or_alloc here?

sloede · 2025-02-06T09:40:22Z

src/solvers/dgsem_p4est/containers.jl

+                                       uArray <: DenseArray{uEltype, NDIMSP2},
+                                       IdsMatrix <: DenseMatrix{Int},


Why the restriction to dense types? (just curious)

sloede · 2025-02-06T09:41:04Z

src/solvers/dgsem_p4est/containers.jl

 end

 @inline function ninterfaces(interfaces::P4estInterfaceContainer)
    size(interfaces.neighbor_ids, 2)
 end
 @inline Base.ndims(::P4estInterfaceContainer{NDIMS}) where {NDIMS} = NDIMS
+@inline function Base.eltype(::P4estInterfaceContainer{NDIMS, uEltype}) where {NDIMS,
+                                                                               uEltype}
+    uEltype


Suggested change

uEltype

return uEltype

We like to be explicit for the sake of less experienced users...

sloede · 2025-02-06T09:44:23Z

src/solvers/dgsem_p4est/containers_parallel.jl

+    # TODO: Vector of Vector type data structure does not work on GPUs,
+    # must be redesigned. This skeleton implementation here just exists just
+    # for compatibility with the rest of the KA.jl solver code


Wasn't this solved by the VecOfArrays/VectorOfArrays type? That is, is the comment still relevant? If yes, maybe add one sentence why VecOfArrays does not solve the issue?

sloede · 2025-02-06T09:45:32Z

src/solvers/dgsem_p4est/dg_parallel.jl

+    mpi_neighbor_interfaces = VecOfArrays(_mpi_neighbor_interfaces)
+    mpi_neighbor_mortars = VecOfArrays(_mpi_neighbor_mortars)
+    mpi_send_buffers = VecOfArrays(_mpi_send_buffers)
+    mpi_recv_buffers = VecOfArrays(_mpi_recv_buffers)


Instead of this indirection here, would it make sense to already return the proper VecOfArrays type from the init_mpi_... functions?

github-actions bot reviewed Dec 17, 2024

View reviewed changes

src/semidiscretization/semidiscretization_hyperbolic.jl Outdated Show resolved Hide resolved

src/semidiscretization/semidiscretization_hyperbolic.jl Outdated Show resolved Hide resolved

src/semidiscretization/semidiscretization_hyperbolic.jl Outdated Show resolved Hide resolved

vchuravy force-pushed the vc/adapt branch from 3fff1fa to f6866b0 Compare January 19, 2025 18:44

ranocha reviewed Feb 3, 2025

View reviewed changes

src/Trixi.jl Outdated Show resolved Hide resolved

src/solvers/dgsem_p4est/containers_parallel.jl Outdated Show resolved Hide resolved

ranocha reviewed Feb 3, 2025

View reviewed changes

src/solvers/dgsem_p4est/containers.jl Outdated Show resolved Hide resolved

ranocha requested changes Feb 3, 2025

View reviewed changes

vchuravy and others added 4 commits February 3, 2025 14:01

use adapt_structure macro where-ever possible

89020ee

fixup! use adapt_structure macro where-ever possible

cedd86c

fixup! Use Adapt.jl to change storage and element type

72cb0f2

vchuravy force-pushed the vc/adapt branch from a920040 to 72cb0f2 Compare February 3, 2025 13:05

vchuravy added 8 commits February 3, 2025 14:29

remove AbstractHeterogenousContainer for now

cb1f54e

actually include vector_of_arrays

66ddd65

fix p4est

d5abb1a

fix downgrade and typos

1dc7c9e

Add HeterogenousContainer

410ee91

fixup! Add HeterogenousContainer

fa0239e

fixup! fixup! Add HeterogenousContainer

ad02c33

fixup! fixup! fixup! Add HeterogenousContainer

ca99ff6

vchuravy marked this pull request as ready for review February 3, 2025 15:38

vchuravy requested review from ranocha and benegee February 3, 2025 15:38

ranocha requested changes Feb 4, 2025

View reviewed changes

define conversion to VecOfArrays

4bf0864

sloede reviewed Feb 6, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use Adapt.jl to change storage and element type #2212

Use Adapt.jl to change storage and element type #2212

vchuravy commented Dec 17, 2024

github-actions bot commented Dec 17, 2024

ranocha left a comment

codecov bot commented Feb 3, 2025 •

edited

Loading

ranocha left a comment

ranocha Feb 4, 2025

vchuravy Feb 4, 2025

JoshuaLampert Feb 4, 2025

JoshuaLampert Feb 4, 2025

vchuravy Feb 5, 2025

ranocha Feb 6, 2025

sloede left a comment

sloede Feb 6, 2025

sloede Feb 6, 2025 •

edited by ranocha

Loading

ranocha Feb 6, 2025

sloede Feb 6, 2025

sloede Feb 6, 2025

ranocha Feb 7, 2025

sloede Feb 7, 2025

sloede Feb 6, 2025

sloede Feb 6, 2025

sloede Feb 6, 2025

sloede Feb 6, 2025

sloede Feb 6, 2025

		uArray <: DenseArray{uEltype, NDIMSP2},
		IdsMatrix <: DenseMatrix{Int},

Use Adapt.jl to change storage and element type #2212

Are you sure you want to change the base?

Use Adapt.jl to change storage and element type #2212

Conversation

vchuravy commented Dec 17, 2024

github-actions bot commented Dec 17, 2024

Review checklist

Purpose and scope

Code quality

Documentation

Testing

Performance

Verification

ranocha left a comment

Choose a reason for hiding this comment

codecov bot commented Feb 3, 2025 • edited Loading

Codecov Report

ranocha left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sloede left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sloede Feb 6, 2025 • edited by ranocha Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Feb 3, 2025 •

edited

Loading

sloede Feb 6, 2025 •

edited by ranocha

Loading