fix batched in gpu transforms

Author: fgerick

ext/SHTnsCUDAExt/sht.jl

@@ -1,24 +1,24 @@
-function cu_spat_to_SH(cfg, Vr::CuMatrix{Float64}, Qlm::CuVector{Complex{Float64}}, lmax)
+function cu_spat_to_SH(cfg, Vr::CuArray{Float64}, Qlm::CuVector{Complex{Float64}}, lmax)
     ccall((:cu_spat_to_SH, libshtns[]), Nothing, (shtns_cfg, CuPtr{Float64}, CuPtr{Complex{Float64}}, Clong), cfg, Vr, Qlm, lmax)
 end
 
-function cu_SH_to_spat(cfg, Qlm::CuVector{Complex{Float64}}, Vr::CuMatrix{Float64}, lmax)
+function cu_SH_to_spat(cfg, Qlm::CuVector{Complex{Float64}}, Vr::CuArray{Float64}, lmax)
     ccall((:cu_SH_to_spat, libshtns[]), Nothing, (shtns_cfg, CuPtr{Complex{Float64}}, CuPtr{Float64}, Clong), cfg, Qlm, Vr, lmax)
 end
 
-function cu_spat_to_SHsphtor(cfg, Vt::CuMatrix{Float64}, Vp::CuMatrix{Float64}, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, lmax)
+function cu_spat_to_SHsphtor(cfg, Vt::CuArray{Float64}, Vp::CuArray{Float64}, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, lmax)
     ccall((:cu_spat_to_SHsphtor, libshtns[]), Nothing, (shtns_cfg,CuPtr{Float64},CuPtr{Float64},CuPtr{ComplexF64},CuPtr{ComplexF64}, Clong), cfg, Vt, Vp, Slm, Tlm, lmax)
 end
 
-function cu_SHsphtor_to_spat(cfg, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, Vt::CuMatrix{Float64}, Vp::CuMatrix{Float64}, lmax)
+function cu_SHsphtor_to_spat(cfg, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, Vt::CuArray{Float64}, Vp::CuArray{Float64}, lmax)
     ccall((:cu_SHsphtor_to_spat, libshtns[]), Nothing, (shtns_cfg,CuPtr{ComplexF64},CuPtr{ComplexF64},CuPtr{Float64},CuPtr{Float64}, Clong), cfg, Slm, Tlm, Vt, Vp, lmax)
 end
 
-function spat_to_SHqst(cfg, Vr::CuMatrix{Float64}, Vt::CuMatrix{Float64}, Vp::CuMatrix{Float64}, Qlm::CuVector{Complex{Float64}}, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, lmax)
+function spat_to_SHqst(cfg, Vr::CuArray{Float64}, Vt::CuArray{Float64}, Vp::CuArray{Float64}, Qlm::CuVector{Complex{Float64}}, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, lmax)
     ccall((:cu_spat_to_SHqst, libshtns[]), Nothing, (shtns_cfg,CuPtr{Float64},CuPtr{Float64},CuPtr{Float64},CuPtr{ComplexF64},CuPtr{ComplexF64},CuPtr{ComplexF64}, Clong), cfg, Vr, Vt, Vp, Qlm, Slm, Tlm, lmax)
 end
 
-function SHqst_to_spat(cfg, Qlm::CuVector{Complex{Float64}}, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, Vr::CuMatrix{Float64}, Vt::CuMatrix{Float64}, Vp::CuMatrix{Float64}, lmax)
+function SHqst_to_spat(cfg, Qlm::CuVector{Complex{Float64}}, Slm::CuVector{Complex{Float64}}, Tlm::CuVector{Complex{Float64}}, Vr::CuArray{Float64}, Vt::CuArray{Float64}, Vp::CuArray{Float64}, lmax)
     ccall((:cu_SHqst_to_spat, libshtns[]), Nothing, (shtns_cfg,CuPtr{ComplexF64},CuPtr{ComplexF64},CuPtr{ComplexF64},CuPtr{Float64},CuPtr{Float64},CuPtr{Float64}, Clong), cfg, Qlm, Slm, Tlm, Vr, Vt, Vp, lmax)
 end
 

ext/SHTnsCUDAExt/synth.jl

@@ -7,7 +7,7 @@ function synth(cfg::SHTnsCfg{TR,T,N}, qlm::CuVector{ComplexF64}) where {TR,T,N}
     nx = cfg.shtype.contiguous_phi ? cfg.nphi : cfg.nlat_padded
     ny = cfg.shtype.contiguous_phi ? cfg.nlat_padded : cfg.nphi
 
-    v = CuMatrix{Tv}(undef, nx, ny)
+    v = cfg.howmany > 1 ? CuArray{Tv}(undef, nx, ny, cfg.howmany) : CuMatrix{Tv}(undef, nx, ny)
     synth!(cfg, qlm, v)
     return v
 end
@@ -21,8 +21,8 @@ function synth(cfg::SHTnsCfg{TR,T,N}, slm::CuVector{ComplexF64}, tlm::CuVector{C
     nx = cfg.shtype.contiguous_phi ? cfg.nphi : cfg.nlat_padded
     ny = cfg.shtype.contiguous_phi ? cfg.nlat_padded : cfg.nphi
 
-    utheta = CuMatrix{Tv}(undef, nx, ny)
-    uphi = CuMatrix{Tv}(undef, nx, ny)
+    utheta = cfg.howmany > 1 ? CuArray{Tv}(undef, nx, ny, cfg.howmany) : CuMatrix{Tv}(undef, nx, ny)
+    uphi = cfg.howmany > 1 ? CuArray{Tv}(undef, nx, ny, cfg.howmany) : CuMatrix{Tv}(undef, nx, ny)
     synth!(cfg, slm, tlm, utheta, uphi)
     return utheta, uphi
 end
@@ -35,27 +35,27 @@ function synth(cfg::SHTnsCfg{TR,T,N}, qlm::CuVector{ComplexF64}, slm::CuVector{C
     nx = cfg.shtype.contiguous_phi ? cfg.nphi : cfg.nlat_padded
     ny = cfg.shtype.contiguous_phi ? cfg.nlat_padded : cfg.nphi
 
-    ur = CuMatrix{Tv}(undef, nx, ny)
-    utheta = CuMatrix{Tv}(undef, nx, ny)
-    uphi = CuMatrix{Tv}(undef, nx, ny)
+    ur = cfg.howmany > 1 ? CuArray{Tv}(undef, nx, ny, cfg.howmany) : CuMatrix{Tv}(undef, nx, ny)
+    utheta = cfg.howmany > 1 ? CuArray{Tv}(undef, nx, ny, cfg.howmany) : CuMatrix{Tv}(undef, nx, ny)
+    uphi = cfg.howmany > 1 ? CuArray{Tv}(undef, nx, ny, cfg.howmany) : CuMatrix{Tv}(undef, nx, ny)
     synth!(cfg, qlm, slm, tlm, ur, utheta, uphi)
     return ur, utheta, uphi
 end
 
 
-function synth!(cfg::SHTnsCfg{Real,T,N}, qlm::CuVector{ComplexF64}, v::CuMatrix{Float64}) where {T,N}
+function synth!(cfg::SHTnsCfg{Real,T,N}, qlm::CuVector{ComplexF64}, v::CuArray{Float64}) where {T,N}
     @assert cfg.shtype.gpu
     cu_SH_to_spat(cfg.cfg, qlm, v, cfg.lmax)
     return v
 end
 
-function synth!(cfg::SHTnsCfg{Real,T,N}, slm::CuVector{ComplexF64}, tlm::CuVector{ComplexF64}, utheta::Tv, uphi::Tv) where {T,N,Tv<:CuMatrix{Float64}}
+function synth!(cfg::SHTnsCfg{Real,T,N}, slm::CuVector{ComplexF64}, tlm::CuVector{ComplexF64}, utheta::Tv, uphi::Tv) where {T,N,Tv<:CuArray{Float64}}
     @assert cfg.shtype.gpu
     cu_SHsphtor_to_spat(cfg.cfg, slm, tlm, utheta, uphi, cfg.lmax)
     return utheta, uphi
 end
 
-function synth!(cfg::SHTnsCfg{Real,T,N}, qlm::CuVector{ComplexF64}, slm::CuVector{ComplexF64}, tlm::CuVector{ComplexF64}, ur::Tv, utheta::Tv, uphi::Tv) where {T,N,Tv<:CuMatrix{Float64}}
+function synth!(cfg::SHTnsCfg{Real,T,N}, qlm::CuVector{ComplexF64}, slm::CuVector{ComplexF64}, tlm::CuVector{ComplexF64}, ur::Tv, utheta::Tv, uphi::Tv) where {T,N,Tv<:CuArray{Float64}}
     @assert cfg.shtype.gpu
     cu_SHqst_to_spat(cfg.cfg, qlm, slm, tlm, ur, utheta, uphi, cfg.lmax)
     return ur, utheta, uphi