Fix style

PTsolvers · Mar 13, 2024 · f6fe4f4 · f6fe4f4
1 parent e6bfd9b
commit f6fe4f4
Show file tree

Hide file tree

Showing 10 changed files with 56 additions and 54 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Chmy"
 uuid = "33a72cf0-4690-46d7-b987-06506c2248b9"
 authors = ["Ivan Utkin <[email protected]>, Ludovic Raess <[email protected]>, and contributors"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/examples/diffusion_2d.jl b/examples/diffusion_2d.jl
@@ -1,8 +1,10 @@
 using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch
 using KernelAbstractions
-# using AMDGPU
+using Printf
 using CairoMakie
 
+# using AMDGPU
+
 @kernel inbounds = true function compute_q!(q, C, χ, g::StructuredGrid, O)
     I = @index(Global, NTuple)
     I = I + O
@@ -16,22 +18,21 @@ end
     C[I...] -= Δt * divg(q, g, I...)
 end
 
-@views function main(backend=CPU())
+@views function main(backend=CPU(); nxy=(126, 126))
     arch = Arch(backend)
     # geometry
-    grid = UniformGrid(arch; origin=(-1, -1), extent=(2, 2), dims=(1022, 1022))
+    grid   = UniformGrid(arch; origin=(-1, -1), extent=(2, 2), dims=nxy)
+    launch = Launcher(arch, grid; outer_width=(16, 8))
     # physics
     χ = 1.0
     # numerics
-    Δt = minimum(spacing(grid, Center(), 1, 1))^2 / χ / ndims(grid) / 2.1
+    Δt = minimum(spacing(grid))^2 / χ / ndims(grid) / 2.1
     # allocate fields
     C = Field(backend, grid, Center())
     q = VectorField(backend, grid)
     # initial conditions
     set!(C, grid, (_, _) -> rand())
     bc!(arch, grid, C => Neumann(); exchange=C)
-    launch = Launcher(arch, grid)
-    @show typeof(launch)
     # visualisation
     fig = Figure(; size=(400, 320))
     ax  = Axis(fig[1, 1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="it = 0")
@@ -41,8 +42,9 @@ end
     # action
     nt = 100
     for it in 1:nt
-        @time launch(arch, grid, compute_q! => (q, C, χ, grid))
-        @time launch(arch, grid, update_C! => (C, q, Δt, grid); bc=batch(grid, C => Neumann(); exchange=C))
+        @printf("it = %d/%d \n", it, nt)
+        launch(arch, grid, compute_q! => (q, C, χ, grid))
+        launch(arch, grid, update_C! => (C, q, Δt, grid); bc=batch(grid, C => Neumann(); exchange=C))
     end
     KernelAbstractions.synchronize(backend)
     plt[3] = interior(C) |> Array
@@ -51,5 +53,6 @@ end
     return
 end
 
-# main(ROCBackend())
-main()
+n = 128
+# main(ROCBackend(); nxy=(n, n) .- 2)
+main(; nxy=(n, n) .- 2)
diff --git a/examples/diffusion_2d_mpi.jl b/examples/diffusion_2d_mpi.jl
@@ -1,12 +1,12 @@
 using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch
 using KernelAbstractions
-using CairoMakie
 using Printf
+# using CairoMakie
 
 # using AMDGPU
 # AMDGPU.allowscalar(false)
-using CUDA
-CUDA.allowscalar(false)
+# using CUDA
+# CUDA.allowscalar(false)
 
 using Chmy.Distributed
 using MPI
@@ -28,24 +28,24 @@ end
 @views function main(backend=CPU(); nxy_l=(126, 126))
     arch = Arch(backend, MPI.COMM_WORLD, (0, 0))
     topo = topology(arch)
-    me = global_rank(topo)
+    me   = global_rank(topo)
     # geometry
     dims_l = nxy_l
     dims_g = dims_l .* dims(topo)
-    grid = UniformGrid(arch; origin=(-2, -2), extent=(4, 4), dims=dims_g)
+    grid   = UniformGrid(arch; origin=(-2, -2), extent=(4, 4), dims=dims_g)
+    launch = Launcher(arch, grid, outer_width=(16, 8))
     nx, ny = dims_g
     # physics
     χ = 1.0
     # numerics
-    Δt = minimum(spacing(grid, Center(), 1, 1))^2 / χ / ndims(grid) / 2.1
+    Δt = minimum(spacing(grid))^2 / χ / ndims(grid) / 2.1
     # allocate fields
     C = Field(backend, grid, Center())
     q = VectorField(backend, grid)
     C_v = (me==0) ? KernelAbstractions.zeros(CPU(), Float64, size(interior(C)) .* dims(topo)) : nothing
     # initial conditions
     set!(C, grid, (x, y) -> exp(-x^2 - y^2))
     bc!(arch, grid, C => Neumann(); exchange=C)
-    launch = Launcher(arch, grid, outer_width=(16, 8)) # ; outer_width=(16, 8)
     # visualisation
     fig = Figure(; size=(400, 320))
     ax  = Axis(fig[1, 1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="it = 0")
@@ -60,8 +60,8 @@ end
     end
     KernelAbstractions.synchronize(backend)
     # local postprocess
-    plt[3] = interior(C) |> Array
-    ax.title = "it = $nt"
+    # plt[3] = interior(C) |> Array
+    # ax.title = "it = $nt"
     # display(fig)
     # save("out$me.png", fig)
     # global postprocess
@@ -76,8 +76,10 @@ end
     return
 end
 
-# main(ROCBackend(); n=1022)
-main(CUDABackend(); nxy_l=(254, 254))
-# main(; n=1022)
+n = 128
+
+# main(ROCBackend(); nxy_l=(n, n) .- 2)
+# main(CUDABackend(); nxy_l=(n, n) .- 2)
+main(; nxy_l=(n, n) .- 2)
 
 MPI.Finalize()
diff --git a/examples/diffusion_2d_mpi_perf.jl b/examples/diffusion_2d_mpi_perf.jl
@@ -1,7 +1,7 @@
 using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch
 using KernelAbstractions
-# using CairoMakie
 using Printf
+# using CairoMakie
 
 using AMDGPU
 AMDGPU.allowscalar(false)
@@ -42,21 +42,20 @@ end
     # geometry
     dims_l = nxy_l
     dims_g = dims_l .* dims(topo)
-    grid = UniformGrid(arch; origin=(-2, -2), extent=(4, 4), dims=dims_g)
+    grid   = UniformGrid(arch; origin=(-2, -2), extent=(4, 4), dims=dims_g)
+    launch = Launcher(arch, grid, outer_width=(128, 8))
     nx, ny = dims_g
     # physics
     χ = 1.0
     # numerics
-    Δt = minimum(spacing(grid, Center(), 1, 1))^2 / χ / ndims(grid) / 2.1
+    Δt = minimum(spacing(grid))^2 / χ / ndims(grid) / 2.1
     # allocate fields
     C = Field(backend, grid, Center())
     q = VectorField(backend, grid)
     # C_v = (me==0) ? KernelAbstractions.zeros(CPU(), Float64, size(interior(C)) .* dims(topo)) : nothing
     # initial conditions
     set!(C, grid, (x, y) -> exp(-x^2 - y^2))
     bc!(arch, grid, C => Neumann(); exchange=C)
-    # launch = Launcher(arch, grid)
-    launch = Launcher(arch, grid, outer_width=(128, 8))
     # visualisation
     # fig = Figure(; size=(400, 320))
     # ax  = Axis(fig[1, 1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="it = 0")
@@ -99,6 +98,6 @@ end
 res = 1024 * 16
 main(ROCBackend(); nxy_l=(res, res) .- 2)
 # main(CUDABackend(); nxy_l=(res, res) .- 2)
-# main(; nxy_l=(256, 256))
+# main(; nxy_l=(res, res) .- 2)
 
 MPI.Finalize()
diff --git a/examples/diffusion_2d_perf.jl b/examples/diffusion_2d_perf.jl
@@ -31,10 +31,10 @@ function compute(arch, grid, launch, q, C, χ, Δt, nt)
     return wtime
 end
 
-@views function main(backend=CPU(); nxy_l=(126, 126))
+@views function main(backend=CPU(); nxy=(126, 126))
     arch = Arch(backend; device_id=1)
     # geometry
-    grid = UniformGrid(arch; origin=(-1, -1), extent=(2, 2), dims=nxy_l)
+    grid = UniformGrid(arch; origin=(-1, -1), extent=(2, 2), dims=nxy)
     # physics
     χ = 1.0
     # numerics
@@ -66,7 +66,7 @@ end
             println("Experiment = $ex")
             wtime = compute(arch, grid, lau, q, C, χ, Δt, (iters - warmup))
             # report
-            A_eff = 6 / 1e9 * prod(nxy_l) * sizeof(Float64)
+            A_eff = 6 / 1e9 * prod(nxy) * sizeof(Float64)
             wtime_it = wtime ./ (iters - warmup)
             T_eff = A_eff ./ wtime_it
             @printf("  Executed %d steps in = %1.3e sec (@ T_eff = %1.2f GB/s - device %s) \n", (iters - warmup), wtime,
@@ -80,8 +80,8 @@ end
 end
 
 res = 1024 * 1
-# main(ROCBackend(); nxy_l=(1022, 1022))
-# main(CUDABackend(); nxy_l=(16382, 16382))
-main(ROCBackend(); nxy_l=(res, res) .- 2)
-# main(CUDABackend(); nxy_l=(res, res) .- 2)
-# main(; nxy_l=(256, 256))
+# main(ROCBackend(); nxy=(1022, 1022))
+# main(CUDABackend(); nxy=(16382, 16382))
+main(ROCBackend(); nxy=(res, res) .- 2)
+# main(CUDABackend(); nxy=(res, res) .- 2)
+# main(; nxy=(256, 256))
diff --git a/examples/stokes_2d_inc_ve_T.jl b/examples/stokes_2d_inc_ve_T.jl
@@ -1,7 +1,7 @@
 using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch
 using KernelAbstractions
-using CairoMakie
 using Printf
+using CairoMakie
 
 # using AMDGPU
 # AMDGPU.allowscalar(false)
@@ -77,7 +77,8 @@ end
     # numerics
     nx = ny = nxy
     grid   = UniformGrid(arch; origin=(-lx/2, -ly/2), extent=(lx, ly), dims=(nx, ny))
-    dx, dy = spacing(grid, Center(), 1, 1)
+    launch = Launcher(arch, grid)
+    dx, dy = spacing(grid)
     nt     = 4
     niter  = 50nx
     ncheck = 2nx
@@ -106,7 +107,6 @@ end
     ρgy = FunctionField(init_incl, grid, (Center(), Vertex()); parameters=(x0=0.0, y0=0.0, r=0.1lx, in=ρg.y, out=0.0))
     set!(T, grid, init_incl; parameters=(x0=0.0, y0=0.0, r=0.1lx, in=T0, out=Ta))
     η_ve = 0.0
-    launch = Launcher(arch, grid)
     # boundary conditions
     bc_V = (V.x => (x=Dirichlet(), y=Neumann()),
             V.y => (x=Neumann(), y=Dirichlet()))

diff --git a/examples/stokes_2d_inc_ve_T_mpi.jl b/examples/stokes_2d_inc_ve_T_mpi.jl
@@ -1,7 +1,7 @@
 using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch
 using KernelAbstractions
-using CairoMakie
 using Printf
+using CairoMakie
 
 using AMDGPU
 AMDGPU.allowscalar(false)
@@ -90,8 +90,9 @@ end
     dims_l = nxy_l
     dims_g = dims_l .* dims(topo)
     grid   = UniformGrid(arch; origin=(-lx/2, -ly/2), extent=(lx, ly), dims=dims_g)
+    launch = Launcher(arch, grid; outer_width=(16, 8))
     nx, ny = dims_g
-    dx, dy = spacing(grid, Center(), 1, 1)
+    dx, dy = spacing(grid)
     nt     = 4
     niter  = 50nx
     ncheck = 2nx
@@ -122,7 +123,6 @@ end
     ρgy = FunctionField(init_incl, grid, (Center(), Vertex()); parameters=(x0=0.0, y0=0.0, r=0.1lx, in=ρg.y, out=0.0))
     set!(T, grid, init_incl; parameters=(x0=0.0, y0=0.0, r=0.1lx, in=T0, out=Ta))
     η_ve = 0.0
-    launch = Launcher(arch, grid; outer_width=(16, 8))
     # boundary conditions
     bc_V = (V.x => (x=Dirichlet(), y=Neumann()),
             V.y => (x=Neumann(), y=Dirichlet()))

diff --git a/examples/stokes_3d_inc_ve_T.jl b/examples/stokes_3d_inc_ve_T.jl
@@ -1,7 +1,7 @@
 using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch
 using KernelAbstractions
-using CairoMakie
 using Printf
+using CairoMakie
 
 # using AMDGPU
 # AMDGPU.allowscalar(false)
@@ -94,7 +94,8 @@ end
     # numerics
     nx = ny = nz = nxyz
     grid   = UniformGrid(arch; origin=(-lx/2, -ly/2, -lz/2), extent=(lx, ly, lz), dims=(nx, ny, nz))
-    dx, dy, dz = spacing(grid, Center(), 1, 1, 1)
+    launch = Launcher(arch, grid)
+    dx, dy, dz = spacing(grid)
     nt     = 2
     niter  = 50nx
     ncheck = 2nx
@@ -125,7 +126,6 @@ end
     set!(ρgz, grid, init_incl; parameters=(x0=0.0, y0=0.0, z0=0.0, r=0.1lx, in=ρg.z, out=0.0))
     set!(T, grid, init_incl; parameters=(x0=0.0, y0=0.0, z0=0.0, r=0.1lx, in=T0, out=Ta))
     η_ve = 0.0
-    launch = Launcher(arch, grid)
     # boundary conditions
     bc_V = (V.x => (x=Dirichlet(), y=Neumann(), z=Neumann()),
             V.y => (x=Neumann(), y=Dirichlet(), z=Neumann()),

diff --git a/examples/stokes_3d_inc_ve_T_mpi.jl b/examples/stokes_3d_inc_ve_T_mpi.jl
@@ -1,8 +1,8 @@
 using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch
 using KernelAbstractions
-# using CairoMakie
 using Printf
 using JSON
+# using CairoMakie
 
 using AMDGPU
 AMDGPU.allowscalar(false)
@@ -109,7 +109,7 @@ end
     grid   = UniformGrid(arch; origin=(-lx/2, -ly/2, -lz/2), extent=(lx, ly, lz), dims=dims_g)
     launch = Launcher(arch, grid; outer_width=(128, 8, 4))
     nx, ny, nz = dims_g
-    dx, dy, dz = spacing(grid, Center(), 1, 1, 1)
+    dx, dy, dz = spacing(grid)
     nt     = params.nt #1e3
     niter  = 50nx
     ncheck = 0.5nx

diff --git a/examples/stokes_3d_inc_ve_T_mpi_perf.jl b/examples/stokes_3d_inc_ve_T_mpi_perf.jl
@@ -106,16 +106,17 @@ end
     dims_l = nxyz_l
     dims_g = dims_l .* dims(topo)
     grid   = UniformGrid(arch; origin=(-lx/2, -ly/2, -lz/2), extent=(lx, ly, lz), dims=dims_g)
+    launch = Launcher(arch, grid; outer_width=(128, 8, 4))
     nx, ny, nz = dims_g
     dx, dy, dz = spacing(grid, Center(), 1, 1, 1)
-    nt     = 1
+    nt     = 8
     # niter  = 50nx
     niter, warmup = 110, 10
     ncheck = 2nx
     ϵ_it   = 1e-6
     ysl    = ceil(Int, length(ycenters(grid)) / 2) # for visu
     # PT params
-    re_m    = 2.3π
+    re_m    = 2.5π
     r       = 0.5
     lτ_re_m = min(lx, ly, lz) / re_m
     vdτ     = min(dx, dy, dz) / sqrt(ndims(grid) * 1.1)
@@ -139,8 +140,6 @@ end
     ρgz = FunctionField(init_incl, grid, (Center(), Center(), Vertex()); parameters=(x0=0.0, y0=0.0, z0=0.0, r=0.1lx, in=ρg.z, out=0.0))
     set!(T, grid, init_incl; parameters=(x0=0.0, y0=0.0, z0=0.0, r=0.1lx, in=T0, out=Ta))
     η_ve = 0.0
-    # launch = Launcher(arch, grid)
-    launch = Launcher(arch, grid; outer_width=(128, 8, 4))
     # boundary conditions
     bc_V = (V.x => (x=Dirichlet(), y=Neumann(), z=Neumann()),
             V.y => (x=Neumann(), y=Dirichlet(), z=Neumann()),
@@ -178,9 +177,8 @@ end
         η_ve = 1.0 / (1.0 / η + 1.0 / (G * dt))
         (it > 2) && (ncheck = ceil(Int, 0.5nx))
 
-        MPI.Barrier(cart_comm(topo))
         for iter in 1:niter
-            (iter == warmup) && (tic = time_ns())
+            (iter == warmup) && (MPI.Barrier(cart_comm(topo)); tic = time_ns())
             launch(arch, grid, update_stress! => (τ, Pr, ∇V, V, τ_old, η, η_ve, G, dt, dτ_Pr, dτ_r, grid))
             launch(arch, grid, update_velocity! => (V, r_V, Pr, τ, ρgz, η_ve, νdτ, grid); bc=batch(grid, bc_V...; exchange=(V.x, V.y, V.z)))
             # if it > 1