GPU Metric Part II

Jobs/JobNHBaroWaveDrySphere

@@ -30,8 +30,8 @@ mpirun -n 1 julia --project Examples/testNHSphere.jl \
   --Stretch=true \
   --StretchType="Exp" \
   --GridType="CubedSphere" \
-  --nz=10 \
-  --nPanel=10 \
+  --nz=64 \
+  --nPanel=30 \
   --H=30000.0 \
   --OrdPoly=3 \
   --HyperVisc=true \

README.md

@@ -1 +1,7 @@
-# CGDycore.jl
+# CGDycore.jl
+
+GDycore is an expermental Julia code [Julia](https://julialang.org/) to study numerical methods and implementational details of numerical dycores for weather prediction. The code is parallelized by MPI and the time integration part of the code is already ported to accelerator by using the Julia package KernelAbstractions.jl. within this programming paradigm the code runs on CPUs and different GPU backends (CUDA, Metal). At the moment a spectral continuous Galerkin method following the HOMME philosphy is implemented as first numerical dycore.
+
+## Code structure and usage
+* The code has a general data structure for unstructured grids on the sphere. There exist grid generators for cubed sphere grids, icosahedral triangular and hexagonal grids. In addition there are options to read in other types of unstructured quad grids n the sphere.
+* Three dimensional grids (extension to the vertical) are constructed by extruding the two-dimensional horizontal grids.

src/DyCore/DiscretizationCG.jl

@@ -13,7 +13,26 @@ function DiscretizationCG(backend,FT,Jacobi,CG,Exchange,Global,zs)
   FS = zeros(nQuad,NF)
   J = zeros(nQuad,OPZ,nz,NF)
   X = zeros(nQuad,OPZ,3,nz,NF)
+  F = zeros(4,3,NF)
+  FGPU = KernelAbstractions.zeros(backend,FT,4,3,NF)
+  for iF = 1 : NF
+    F[1,1,iF] = Grid.Faces[iF].P[1].x  
+    F[1,2,iF] = Grid.Faces[iF].P[1].y  
+    F[1,3,iF] = Grid.Faces[iF].P[1].z  
+    F[2,1,iF] = Grid.Faces[iF].P[2].x  
+    F[2,2,iF] = Grid.Faces[iF].P[2].y  
+    F[2,3,iF] = Grid.Faces[iF].P[2].z  
+    F[3,1,iF] = Grid.Faces[iF].P[3].x  
+    F[3,2,iF] = Grid.Faces[iF].P[3].y  
+    F[3,3,iF] = Grid.Faces[iF].P[3].z  
+    F[4,1,iF] = Grid.Faces[iF].P[4].x  
+    F[4,2,iF] = Grid.Faces[iF].P[4].y  
+    F[4,3,iF] = Grid.Faces[iF].P[4].z  
+  end  
+  copyto!(FGPU,F)
+  Grids.JacobiSphere3GPU!(Metric.X,Metric.dXdxI,Metric.J,CG,FGPU,Grid.z,zs,Grid.Rad)
 
+#=
   for iF = 1 : NF
     for iz = 1 : nz
       zI = [Grid.z[iz],Grid.z[iz+1]]
@@ -36,8 +55,9 @@ function DiscretizationCG(backend,FT,Jacobi,CG,Exchange,Global,zs)
   copyto!(Metric.dXdxI,dXdxI)
   copyto!(Metric.nS,nS)
   copyto!(Metric.FS,FS)
+=# 
 
-  MassCGGPU!(backend,FT,CG,Metric.J,CG.Glob,Exchange,Global)
+  MassCGGPU!(CG,Metric.J,CG.Glob,Exchange,Global)
 
   Metric.dz = KernelAbstractions.zeros(backend,FT,nz,CG.NumG)
   Metric.zP = KernelAbstractions.zeros(backend,FT,nz,CG.NumG)
@@ -62,7 +82,7 @@ end
 
 function DiscretizationCG(backend,FT,Jacobi,CG,Exchange,Global)
   DiscretizationCG(backend,FT,Jacobi,CG,Exchange,Global,
-  zeros(CG.OrdPoly+1,CG.OrdPoly+1,Global.Grid.NumFaces))
+  KernelAbstractions.zeros(backend,FT,CG.OrdPoly+1,CG.OrdPoly+1,Global.Grid.NumFaces))
 end  
 
 @kernel function GridSizeSphereKernel!(lat,zP,dz,@Const(X),@Const(Glob),Rad)

src/DyCore/FiniteElement.jl

@@ -25,10 +25,10 @@ mutable struct CGStruct{FT<:AbstractFloat,
     NumG::Int
     NumI::Int
     w::AT1
-    xw::Array{FT, 1}
+    xw::AT1
     xe::Array{FT, 1}
     IntXE2F::Array{FT, 2}
-    xwZ::Array{FT, 1}
+    xwZ::AT1
     IntZE2F::Array{FT, 2}
     DW::AT2
     DWT::Array{FT, 2}
@@ -112,11 +112,15 @@ function CGStruct{FT}(backend,OrdPoly,OrdPolyZ,Grid) where FT<:AbstractFloat
   OrdPolyZ=OrdPolyZ
   DoF = OP * OP
 
-  (wCPU,xw)=DG.GaussLobattoQuad(OrdPoly)
+  (wCPU,xwCPU)=DG.GaussLobattoQuad(OrdPoly)
   w = KernelAbstractions.zeros(backend,FT,size(wCPU))
+  xw = KernelAbstractions.zeros(backend,FT,size(xwCPU))
   copyto!(w,wCPU)
+  copyto!(xw,xwCPU)
 
-  (wZ,xwZ)=DG.GaussLobattoQuad(OrdPolyZ)
+  (wZ,xwZCPU)=DG.GaussLobattoQuad(OrdPolyZ)
+  xwZ = KernelAbstractions.zeros(backend,FT,size(xwZCPU))
+  copyto!(xwZ,xwZCPU)
   xe = zeros(OrdPoly+1)
   xe[1] = -1.0
   for i = 2 : OrdPoly

src/DyCore/InitDriver.jl

@@ -8,7 +8,7 @@ function InitSphere(backend,FT,OrdPoly,OrdPolyZ,nz,nPanel,H,GridType,Topography,
   ParallelCom.ProcNumber  = ProcNumber
 
   TimeStepper = TimeStepperStruct{FT}(backend)
-  Grid = Grids.GridStruct(nz,Topography)
+  Grid = Grids.GridStruct{FT}(backend,nz,Topography)
 
   if GridType == "HealPix"
   # Grid=CGDycore.InputGridH("Grid/mesh_H12_no_pp.nc",
@@ -102,7 +102,7 @@ function InitCart(backend,FT,OrdPoly,OrdPolyZ,nx,ny,Lx,Ly,x0,y0,nz,H,Boundary,Gr
   ParallelCom.ProcNumber  = ProcNumber
 
   TimeStepper = TimeStepperStruct{FT}(backend)
-  Grid = Grids.GridStruct(nz,Topography)
+  Grid = Grids.GridStruct{FT}(backend,nz,Topography)
 
   Grid = Grids.CartGrid(nx,ny,Lx,Ly,x0,y0,Grids.OrientFaceCart,Boundary,Grid)
 

src/DyCore/MassCG.jl

@@ -29,7 +29,9 @@ function MassCG(CG,J,Glob,Exchange)
   return (M,MW,MMass)
 end
 
-function MassCGGPU!(backend,FT,CG,J,Glob,Exchange,Global)
+function MassCGGPU!(CG,J,Glob,Exchange,Global)
+  backend = get_backend(J)
+  FT = eltype(J)
   N = CG.OrdPoly + 1
   DoF = CG.DoF
   w = CG.w

src/Grids/AddVerticalGrid.jl

@@ -1,17 +1,18 @@
 function AddVerticalGrid!(Grid::GridStruct,nz::Int,H::Float64)
-Grid.zP=zeros(nz);
-Grid.z=zeros(nz+1);
-Grid.dzeta = zeros(nz)
-Grid.H = H
-@. Grid.dzeta = H/nz;
-Grid.H=H
-for i=2:nz+1
-  Grid.z[i] = Grid.z[i-1] + Grid.dzeta[i-1]
-end
-for i=1:nz
-  Grid.dzeta[i] = Grid.z[i+1] - Grid.z[i]  
-  Grid.zP[i] = 0.5 * (Grid.z[i] + Grid.z[i+1])
-end
+  Grid.zP=zeros(nz);
+  z=zeros(nz+1);
+  Grid.dzeta = zeros(nz)
+  Grid.H = H
+  @. Grid.dzeta = H/nz;
+  Grid.H=H
+  for i=2:nz+1
+    z[i] = z[i-1] + Grid.dzeta[i-1]
+  end
+  for i=1:nz
+    Grid.dzeta[i] = z[i+1] - z[i]  
+    Grid.zP[i] = 0.5 * (z[i] + z[i+1])
+  end
+  copyto!(Grid.z,z)
 end
 
 function AddStretchICONVerticalGrid!(Grid::GridStruct,nz::Int,H::Float64,sigma::Float64,lambda::Float64)
@@ -55,7 +56,7 @@ function AddExpStretchVerticalGrid!(Grid::GridStruct,nz::Int,H::Float64,
   dz_bottom::Float64,dz_top::Float64)
 
   Grid.zP=zeros(nz);
-  Grid.z=zeros(nz+1);
+  z=zeros(nz+1);
   Grid.dzeta = zeros(nz)
   Grid.H = H
 
@@ -118,10 +119,11 @@ function AddExpStretchVerticalGrid!(Grid::GridStruct,nz::Int,H::Float64,
   # add the bottom level
   faces = [z_bottom; faces...]
   for iz = 1 : nz + 1
-    Grid.z[iz] = faces[iz]
+    z[iz] = faces[iz]
   end
   for i = 1 : nz
-    Grid.dzeta[i] = Grid.z[i+1] - Grid.z[i]
-    Grid.zP[i] = 0.5 * (Grid.z[i] + Grid.z[i+1])
+    Grid.dzeta[i] = z[i+1] - z[i]
+    Grid.zP[i] = 0.5 * (z[i] + z[i+1])
   end
+  copyto!(Grid.z,z)
 end

src/Grids/GridStruct.jl

@@ -14,36 +14,37 @@ function Boundary()
   )
 end  
 
-mutable struct GridStruct
-    nz::Int
-    zP::Array{Float64, 1}
-    z::Array{Float64, 1}
-    dzeta::Array{Float64, 1}
-    H::Float64
-    NumFaces::Int
-    NumGhostFaces::Int
-    Faces::Array{Face, 1}
-    NumEdges::Int
-    Edges::Array{Edge, 1}
-    NumNodes::Int
-    Nodes::Array{Node, 1}
-    Form::String
-    Type ::String
-    Dim::Int
-    Rad::Float64
-    NumEdgesI::Int
-    NumEdgesB::Int
-    nBar3::Array{Float64, 2}
-    nBar::Array{Float64, 2}
-    Topography::NamedTuple
-    colors::Array{Array{Int, 1}, 1}
-    Spline_2d::Dierckx.Spline2D
-    BoundaryFaces::Array{Int,1}
-    InteriorFaces::Array{Int,1}
+mutable struct GridStruct{FT<:AbstractFloat,
+                           AT1<:AbstractArray}
+  nz::Int
+  zP::Array{FT, 1}
+  z::AT1
+  dzeta::Array{FT, 1}
+  H::FT
+  NumFaces::Int
+  NumGhostFaces::Int
+  Faces::Array{Face, 1}
+  NumEdges::Int
+  Edges::Array{Edge, 1}
+  NumNodes::Int
+  Nodes::Array{Node, 1}
+  Form::String
+  Type ::String
+  Dim::Int
+  Rad::FT
+  NumEdgesI::Int
+  NumEdgesB::Int
+  nBar3::Array{FT, 2}
+  nBar::Array{FT, 2}
+  Topography::NamedTuple
+  colors::Array{Array{Int, 1}, 1}
+  Spline_2d::Dierckx.Spline2D
+  BoundaryFaces::Array{Int,1}
+  InteriorFaces::Array{Int,1}
 end
-function GridStruct(nz,Topography)
+function GridStruct{FT}(backend,nz,Topography) where FT <: AbstractFloat
   zP=zeros(nz)
-  z=zeros(nz+1)
+  z=KernelAbstractions.zeros(backend,FT,nz+1)
   dzeta=zeros(nz)
   H=0.0
   NumFaces=0
@@ -65,7 +66,8 @@ function GridStruct(nz,Topography)
   Spline_2d = Spline2D(zeros(0),zeros(0),zeros(0),0,0,0.0)
   BoundaryFaces = zeros(Int,0)
   InteriorFaces = zeros(Int,0)
-   return GridStruct(
+   return GridStruct{FT,
+                     typeof(z)}(
     nz,
     zP,
     z,

src/Grids/Grids.jl

@@ -8,6 +8,9 @@ using RootSolvers
 using Dierckx
 using NCDatasets
 using PairedLinkedLists
+using KernelAbstractions
+using KernelAbstractions: @atomic, @atomicswap, @atomicreplace
+
 
 include("Geometry.jl")
 include("Node.jl")
@@ -20,6 +23,7 @@ include("CubedGrid.jl")
 include("FacesInNodes.jl")
 include("JacobiDG3.jl")
 include("JacobiSphere3.jl")
+include("JacobiSphere3GPU.jl")
 include("OrientFaceCart.jl")
 include("OrientFaceSphere.jl")
 include("Orientation.jl")