add cpu!, gpu!, and cu! macros

Author: marius311

src/CMBLensing.jl

@@ -68,7 +68,7 @@ import Statistics: std
 
 
 export
-    @⌛, @show⌛, @ismain, @namedtuple, @repeated, @unpack, animate,
+    @⌛, @show⌛, @ismain, @namedtuple, @repeated, @unpack, @cpu!, animate,
     argmaxf_lnP, BandPassOp, BaseDataSet, batch, batch_index, batch_length, beamCℓs, cache,
     CachedLenseFlow, camb, cov_to_Cℓ, cpu, Cℓ_2D, Cℓ_to_Cov, DataSet, DerivBasis,
     diag, Diagonal, DiagOp, dot, EBFourier, EBMap, expnorm, Field, FieldArray, fieldinfo,

src/gpu.jl

@@ -4,7 +4,7 @@ using CUDA: cufunc, curand_rng
 using CUDA.CUSPARSE: CuSparseMatrix, CuSparseMatrixCSR, CuSparseMatrixCOO
 using CUDA.CUSOLVER: CuQR
 
-export cuda_gc, gpu
+export cuda_gc, gpu, @gpu!, @cu!
 
 const CuBaseField{B,M,T,A<:CuArray} = BaseField{B,M,T,A}
 
@@ -20,10 +20,40 @@ end
 is_gpu_backed(::BaseField{B,M,T,A}) where {B,M,T,A<:CuArray} = true
 global_rng_for(::Type{<:CuArray}) = curand_rng()
 
+# handy conversion functions and macros
+@doc doc"""
 
+    @gpu! x y
+
+Equivalent to `x = gpu(x)`, `y = gpu(y)`, etc... for any number of
+listed variables. See [`gpu`](@ref).
+"""
+macro gpu!(vars...)
+    :(begin; $((:($(esc(var)) = gpu($(esc(var)))) for var in vars)...); nothing; end)
+end
+@doc doc"""
+
+    gpu(x)
+
+Recursively moves x to GPU, but unlike `CUDA.cu`, doesn't also convert
+to Float32. Equivalent to `adapt_structure(CuArray, x)`. Returns nothing.
+"""
 gpu(x) = adapt_structure(CuArray, x)
 
 
+@doc doc"""
+
+    @cu! x y
+    
+Equivalent to `x = cu(x)`, `y = cu(y)`, etc... for any number of
+listed variables. See `CUDA.cu`. Returns nothing.
+"""
+macro cu!(vars...)
+    :(begin; $((:($(esc(var)) = cu($(esc(var)))) for var in vars)...); nothing; end)
+end
+
+
+
 adapt_structure(::CUDA.Float32Adaptor, proj::ProjLambert) = adapt_structure(CuArray{Float32}, proj)
 
 

src/util.jl

@@ -225,6 +225,18 @@ Recursively move an object to CPU memory. See also [`gpu`](@ref).
 """
 cpu(x) = adapt_structure(Array, x)
 
+@doc doc"""
+
+    @cpu! x y
+
+Equivalent to `x = cpu(x)`, `y = cpu(y)`, etc... for any number of
+listed variables. See [`cpu`](@ref).
+"""
+macro cpu!(vars...)
+    :(begin; $((:($(esc(var)) = cpu($(esc(var)))) for var in vars)...); nothing; end)
+end
+
+
 @doc doc"""
 
     gpu(x)