From c01d9ad4331e7caba751f782729babbbfaecaa4e Mon Sep 17 00:00:00 2001
From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com>
Date: Wed, 24 Jul 2024 14:57:00 +0200
Subject: [PATCH] Make `SolutionSavingCallback` work on GPUs (#579)

* Make `SolutionSavingCallback` work on GPUs

* Avoid overhead on CPUs

---------

Co-authored-by: Sven Berger <berger.sven@gmail.com>
Co-authored-by: Niklas Neher <73897120+LasNikas@users.noreply.github.com>
---
 src/visualization/write2vtk.jl | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/visualization/write2vtk.jl b/src/visualization/write2vtk.jl
index 8e36f136a..dabd463dc 100644
--- a/src/visualization/write2vtk.jl
+++ b/src/visualization/write2vtk.jl
@@ -75,12 +75,15 @@ function trixi2vtk(vu_ode, semi, t; iter=nothing, output_directory="out", prefix
 end
 
 # Convert data for a single TrixiParticle system to VTK format
-function trixi2vtk(v, u, t, system, periodic_box; output_directory="out", prefix="",
-                   iter=nothing, system_name=vtkname(system), write_meta_data=true,
+function trixi2vtk(v_, u_, t, system_, periodic_box; output_directory="out", prefix="",
+                   iter=nothing, system_name=vtkname(system_), write_meta_data=true,
                    max_coordinates=Inf, git_hash=compute_git_hash(),
                    custom_quantities...)
     mkpath(output_directory)
 
+    # Transfer to CPU if data is on the GPU. Do nothing if already on CPU.
+    v, u, system = transfer2cpu(v_, u_, system_)
+
     # handle "_" on optional pre/postfix strings
     add_opt_str_pre(str) = (str === "" ? "" : "$(str)_")
     add_opt_str_post(str) = (str === nothing ? "" : "_$(str)")
@@ -135,6 +138,18 @@ function trixi2vtk(v, u, t, system, periodic_box; output_directory="out", prefix
     vtk_save(pvd)
 end
 
+function transfer2cpu(v_, u_, system_::GPUSystem)
+    v = Adapt.adapt(Array, v_)
+    u = Adapt.adapt(Array, u_)
+    system = Adapt.adapt(Array, system_)
+
+    return v, u, system
+end
+
+function transfer2cpu(v_, u_, system_)
+    return v_, u_, system_
+end
+
 function custom_quantity(quantity::AbstractArray, v, u, t, system)
     return quantity
 end
@@ -357,3 +372,7 @@ function write2vtk!(vtk, v, u, t, model::BoundaryModelDummyParticles, system;
         vtk["wall_velocity"] = view(model.cache.wall_velocity, 1:ndims(system), :)
     end
 end
+
+function write2vtk!(vtk, v, u, t, system::BoundaryDEMSystem; write_meta_data=true)
+    return vtk
+end