Update diffusion benchmarks

scripts_future_API/Project.toml

@@ -4,3 +4,4 @@ KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

scripts_future_API/benchmark_dbc.jl

@@ -13,43 +13,44 @@ using MPI
     if !isnothing(offset)
         I += offset
     end
-    f[I] = val + I[1]
+    f[I] = val
 end
 
 MPI.Init()
 
 arch = Architecture(CPU(), (0, 0))
-grid = CartesianGrid(; origin=(0.0, 0.0), extent=(1.0, 1.0), size=(10, 10))
-field = Field(backend(arch), grid, (Vertex(), Vertex()); halo=1)
-
+grid = CartesianGrid(; origin=(0.0, 0.0), extent=(1.0, 1.0), size=(5, 5))
+field = Field(backend(arch), grid, (Center(), Center()); halo=1)
 
 me = global_rank(details(arch))
 
 fill!(parent(field), Inf)
 
-bc = FieldBoundaryConditions((field,), (DirichletBC{FullCell}(-me-10),))
+bc = BoundaryConditionsBatch((field,), (DirichletBC{FullCell}(-me-10),))
 
-boundary_conditions = ((bc, bc),
-                       (bc, bc))
-
-boundary_conditions = ntuple(Val(length(boundary_conditions))) do D
-    ntuple(Val(2)) do S
-        if neighbor(details(arch), D, S) != MPI.PROC_NULL
-            DistributedBoundaryConditions(Val(S), Val(D), (field, ))
-        else
-            boundary_conditions[D][S]
-        end
-    end
-end
+boundary_conditions = override_boundary_conditions(arch, ((bc, bc), (bc, bc)); exchange=true)
 
 hide_boundaries = HideBoundaries{2}(arch)
 
-outer_width = (4, 4)
+outer_width = (2, 2)
 
 launch!(arch, grid, fill_field! => (field, me); location=location(field), hide_boundaries, boundary_conditions, outer_width)
 
 sleep(me)
 @show coordinates(details(arch))
 display(parent(field))
 
+field_g = if global_rank(details(arch)) == 0
+    KernelAbstractions.allocate(Architectures.backend(arch), eltype(field), dimensions(details(arch)) .* size(field))
+else
+    nothing
+end
+
+gather!(arch, field_g, field)
+
+if global_rank(details(arch)) == 0
+    println("global matrix:")
+    display(field_g)
+end
+
 MPI.Finalize()

scripts_future_API/benchmark_diffusion_2D.jl

@@ -0,0 +1,95 @@
+using FastIce.Grids
+using FastIce.GridOperators
+using FastIce.Fields
+using FastIce.Architectures
+using FastIce.BoundaryConditions
+using FastIce.Distributed
+using FastIce.KernelLaunch
+
+using KernelAbstractions
+using MPI
+
+using Plots
+
+@kernel function update_C!(C, qC, dt, Δ, offset=nothing)
+    I = @index(Global, Cartesian)
+    isnothing(offset) || (I += offset)
+    @inbounds if checkbounds(Bool, C, I)
+        C[I] -= dt * (∂ᶜx(qC.x, I) / Δ.x +
+                      ∂ᶜy(qC.y, I) / Δ.y)
+    end
+end
+
+@kernel function update_qC!(qC, C, dc, Δ, offset=nothing)
+    I = @index(Global, Cartesian)
+    isnothing(offset) || (I += offset)
+    @inbounds if checkbounds(Bool, qC.x, I)
+        qC.x[I] = -dc * ∂ᵛx(C, I) / Δ.x
+    end
+    @inbounds if checkbounds(Bool, qC.y, I)
+        qC.y[I] = -dc * ∂ᵛy(C, I) / Δ.y
+    end
+end
+
+function diffusion_2D(ka_backend=CPU())
+    # setup arch
+    arch = Architecture(ka_backend, (0, 0))
+    topology = details(arch)
+    # physics
+    lx, ly = 10.0, 10.0
+    dc = 1
+    # numerics
+    nx, ny = 32, 32
+    nt = 100
+    # preprocessing
+    grid = CartesianGrid(; origin=(-0.5lx, -0.5ly),
+                         extent=(lx, ly),
+                         size=(nx, ny))
+    Δ = NamedTuple{(:x, :y)}(spacing(grid))
+    dt = minimum(Δ)^2 / dc / ndims(grid) / 2.1
+    hide_boundaries = HideBoundaries{ndims(grid)}(arch)
+    outer_width = (4, 4)
+    # fields
+    C = Field(arch, grid, Center(); halo=1)
+    qC = (x = Field(arch, grid, (Vertex(), Center())),
+          y = Field(arch, grid, (Center(), Vertex())))
+    C_g = if global_rank(topology) == 0
+        KernelAbstractions.allocate(Architectures.backend(arch), eltype(C), dimensions(topology) .* size(C))
+    else
+        nothing
+    end
+    # initial condition
+    foreach(comp -> set!(comp, 0.0), qC)
+    set!(C, grid, (x, y) -> exp(-x^2 - y^2))
+    # boundary conditions
+    zero_flux_bc = DirichletBC{FullCell}(0.0)
+    bc_q = NamedTuple(comp => BoundaryConditionsBatch((qC[comp],), (zero_flux_bc,)) for comp in eachindex(qC))
+    # zero flux at physical boundaries and nothing at MPI boundaries
+    bc_q = override_boundary_conditions(arch, ((bc_q.x, bc_q.x), (bc_q.y, bc_q.y)))
+    # nothing at physical boundaries and communication at MPI boundaries
+    bc_c = BoundaryConditionsBatch((C,), nothing)
+    bc_c = override_boundary_conditions(arch, ((bc_c, bc_c), (bc_c, bc_c)); exchange=true)
+    ntuple(Val(ndims(grid))) do D
+        apply_boundary_conditions!(Val(1), Val(D), arch, grid, bc_c[D][1])
+        apply_boundary_conditions!(Val(2), Val(D), arch, grid, bc_c[D][2])
+    end
+    # time loop
+    for it in 1:nt
+        launch!(arch, grid, update_qC! => (qC, C, dc, Δ); location=Vertex(), hide_boundaries, boundary_conditions=bc_q, outer_width)
+        launch!(arch, grid, update_C! => (C, qC, dt, Δ); location=Center(), hide_boundaries, boundary_conditions=bc_c, outer_width)
+    end
+    Architectures.synchronize(arch)
+
+    gather!(arch, C_g, C)
+
+    if global_rank(topology) == 0
+        p1 = heatmap(C_g; aspect_ratio=1, size=(600,600))
+        png(p1, "C.png")
+    end
+
+    return
+end
+
+MPI.Init()
+diffusion_2D()
+MPI.Finalize()

scripts_future_API/benchmark_diffusion_3D.jl

@@ -0,0 +1,100 @@
+using FastIce.Grids
+using FastIce.GridOperators
+using FastIce.Fields
+using FastIce.Architectures
+using FastIce.BoundaryConditions
+using FastIce.Distributed
+using FastIce.KernelLaunch
+
+using KernelAbstractions
+using MPI
+
+using Plots
+
+@kernel function update_C!(C, qC, dt, Δ, offset=nothing)
+    I = @index(Global, Cartesian)
+    isnothing(offset) || (I += offset)
+    @inbounds if checkbounds(Bool, C, I)
+        C[I] -= dt * (∂ᶜx(qC.x, I) / Δ.x +
+                      ∂ᶜy(qC.y, I) / Δ.y +
+                      ∂ᶜz(qC.z, I) / Δ.z)
+    end
+end
+
+@kernel function update_qC!(qC, C, dc, Δ, offset=nothing)
+    I = @index(Global, Cartesian)
+    isnothing(offset) || (I += offset)
+    @inbounds if checkbounds(Bool, qC.x, I)
+        qC.x[I] = -dc * ∂ᵛx(C, I) / Δ.x
+    end
+    @inbounds if checkbounds(Bool, qC.y, I)
+        qC.y[I] = -dc * ∂ᵛy(C, I) / Δ.y
+    end
+    @inbounds if checkbounds(Bool, qC.z, I)
+        qC.z[I] = -dc * ∂ᵛz(C, I) / Δ.z
+    end
+end
+
+function diffusion_3D(ka_backend=CPU())
+    # setup arch
+    arch = Architecture(ka_backend, (0, 0, 0))
+    topology = details(arch)
+    # physics
+    lx, ly, lz = 10.0, 10.0, 10.0
+    dc = 1
+    # numerics
+    nx, ny, nz = 32, 32, 32
+    nt = 100
+    # preprocessing
+    grid = CartesianGrid(; origin=(-0.5lx, -0.5ly, -0.5lz),
+                         extent=(lx, ly, lz),
+                         size=(nx, ny, nz))
+    Δ = NamedTuple{(:x, :y, :z)}(spacing(grid))
+    dt = minimum(Δ)^2 / dc / ndims(grid) / 2.1
+    hide_boundaries = HideBoundaries{ndims(grid)}(arch)
+    outer_width = (4, 4, 4)
+    # fields
+    C = Field(arch, grid, Center(); halo=1)
+    qC = (x = Field(arch, grid, (Vertex(), Center(), Center())),
+          y = Field(arch, grid, (Center(), Vertex(), Center())),
+          z = Field(arch, grid, (Center(), Center(), Vertex())))
+    C_g = if global_rank(topology) == 0
+        KernelAbstractions.allocate(Architectures.backend(arch), eltype(C), dimensions(topology) .* size(C))
+    else
+        nothing
+    end
+    # initial condition
+    foreach(comp -> set!(comp, 0.0), qC)
+    set!(C, grid, (x, y, z) -> exp(-x^2 - y^2 - z^2))
+    # boundary conditions
+    zero_flux_bc = DirichletBC{FullCell}(0.0)
+    bc_q = NamedTuple(comp => BoundaryConditionsBatch((qC[comp],), (zero_flux_bc,)) for comp in eachindex(qC))
+    # zero flux at physical boundaries and nothing at MPI boundaries
+    boundary_conditions_q = ((bc_q.x, bc_q.x), (bc_q.y, bc_q.y), (bc_q.z, bc_q.z))
+    boundary_conditions_q = override_boundary_conditions(arch, boundary_conditions_q)
+    # nothing at physical boundaries and communication at MPI boundaries
+    bc_c = BoundaryConditionsBatch((C,), nothing)
+    boundary_conditions_c = ((bc_c, bc_c), (bc_c, bc_c), (bc_c, bc_c))
+    boundary_conditions_c = override_boundary_conditions(arch, boundary_conditions_c; exchange=true)
+    # time loop
+    for it in 1:nt
+        launch!(arch, grid, update_qC! => (qC, C, dc, Δ);
+                location=Vertex(), hide_boundaries, boundary_conditions=boundary_conditions_q, outer_width)
+        launch!(arch, grid, update_C! => (C, qC, dt, Δ);
+                location=Center(), hide_boundaries, boundary_conditions=boundary_conditions_c, outer_width)
+        Architectures.synchronize(arch)
+    end
+
+    gather!(arch, C_g, C)
+
+    if global_rank(topology) == 0
+        p1 = heatmap(interior(C)[:, :, nz÷2]; aspect_ratio=1)
+        png(p1, "C.png")
+    end
+
+    return
+end
+
+MPI.Init()
+diffusion_3D()
+MPI.Finalize()