From 3eb4cad7d604fbe620c68e4b425edb4496f485e7 Mon Sep 17 00:00:00 2001
From: Ivan Utkin <iutkin@ethz.ch>
Date: Mon, 6 Nov 2023 15:36:32 +0100
Subject: [PATCH] Fix bug in boundary condition application

---
 scripts_future_API/benchmark_diffusion_2D.jl  | 26 +++++--------------
 .../field_boundary_conditions.jl              | 20 +++++++++-----
 2 files changed, 21 insertions(+), 25 deletions(-)

diff --git a/scripts_future_API/benchmark_diffusion_2D.jl b/scripts_future_API/benchmark_diffusion_2D.jl
index b7039b75..ba7e3fb7 100644
--- a/scripts_future_API/benchmark_diffusion_2D.jl
+++ b/scripts_future_API/benchmark_diffusion_2D.jl
@@ -33,14 +33,14 @@ end
 
 function diffusion_2D(ka_backend=CPU())
     # setup arch
-    arch = Architecture(ka_backend, (0, 0))
+    arch     = Architecture(ka_backend, (0, 0))
     topology = details(arch)
     # physics
-    lx, ly = 10.0, 10.0
+    lx, ly = 1.0, 1.0
     dc = 1
     # numerics
     nx, ny = 32, 32
-    nt = 500
+    nt = 1000
     # preprocessing
     grid = CartesianGrid(; origin=(-0.5lx, -0.5ly),
                          extent=(lx, ly),
@@ -60,8 +60,9 @@ function diffusion_2D(ka_backend=CPU())
     end
     # initial condition
     foreach(comp -> fill!(parent(comp), 0.0), qC)
-    fill!(parent(C), 0.0)
-    set!(C, grid, (x, y) -> exp(-(x - 0.1lx)^2 - (y-0.05ly)^2))
+    # fill!(parent(C), me)
+    set!(C, grid, (x, y) -> exp(-x^2 - y^2))
+    # set!(C, me)
     # boundary conditions
     zero_flux_bc = DirichletBC{FullCell}(0.0)
     bc_q = (x = BoundaryConditionsBatch((qC.x, qC.y), (zero_flux_bc, nothing)),
@@ -85,21 +86,8 @@ function diffusion_2D(ka_backend=CPU())
         if global_rank(topology) == 0
             println("it = $it")
         end
-        MPI.Barrier(cartesian_communicator(topology))
-        update_qC!(ka_backend, 256, (nx+1, ny+1))(qC, C, dc, Δ)
-        for D in ndims(grid):-1:1
-            apply_boundary_conditions!(Val(1), Val(D), arch, grid, bc_q[D][1])
-            apply_boundary_conditions!(Val(2), Val(D), arch, grid, bc_q[D][2])
-        end
+        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(), expand=1)
-        Architectures.synchronize(arch)
-        # sleep(0.5global_rank(topology))
-        # @show coordinates(topology)
-        # display(parent(C))
-        # update_C!(ka_backend, 256, (nx+2, ny+2))(C, qC, dt, Δ, -CartesianIndex(1,1))
-        # launch!(arch, grid, update_qC! => (qC, C, dc, Δ); location=Vertex(), hide_boundaries, boundary_conditions=bc_q, outer_width)
-        MPI.Barrier(cartesian_communicator(topology))
-
         if it % 5 == 0
             gather!(arch, C_g, C)
             if global_rank(topology) == 0
diff --git a/src/BoundaryConditions/field_boundary_conditions.jl b/src/BoundaryConditions/field_boundary_conditions.jl
index 0d9a48c4..7803b718 100644
--- a/src/BoundaryConditions/field_boundary_conditions.jl
+++ b/src/BoundaryConditions/field_boundary_conditions.jl
@@ -6,22 +6,30 @@ function apply_boundary_conditions!(::Val{S}, ::Val{D},
                                     fields::NTuple{N,Field},
                                     conditions::NTuple{N,FieldOrNothing}; async=true) where {S,D,N}
     _validate_fields(fields, D, S)
-    sizes = ntuple(ifield -> remove_dim(Val(D), size(fields[ifield])), Val(length(fields)))
-    worksize = remove_dim(Val(D), size(grid, Vertex()))
-    _apply_boundary_conditions!(backend(arch), 256, worksize)(Val(S), Val(D), grid, sizes, fields, conditions)
+    sizes = ntuple(ifield -> remove_dim(Val(D), size(parent(fields[ifield]))), Val(N))
+    halos = ntuple(ifield -> remove_dim(Val(D), halo(fields[ifield])), Val(N))
+    offsets = ntuple(Val(N)) do ifield
+        NF = ndims(fields[ifield]) - 1
+        ntuple(Val(NF)) do RD
+            -halos[ifield][RD][1]
+        end |> CartesianIndex
+    end
+    worksize = maximum(sizes)
+    _apply_boundary_conditions!(backend(arch), 256, worksize)(Val(S), Val(D), grid, sizes, offsets, fields, conditions)
     async || KernelAbstractions.synchronize(backend(arch))
     return
 end
 
-@kernel function _apply_boundary_conditions!(side, dim, grid, sizes, fields, conditions)
+@kernel function _apply_boundary_conditions!(side, dim, grid, sizes, offsets, fields, conditions)
     I = @index(Global, Cartesian)
     # ntuple here unrolls the loop over fields
     ntuple(Val(length(fields))) do ifield
         Base.@_inline_meta
         if _checkindices(sizes[ifield], I)
-            field = fields[ifield]
+            Ibc       = I + offsets[ifield]
+            field     = fields[ifield]
             condition = conditions[ifield]
-            _apply_field_boundary_condition!(side, dim, grid, field, location(field), I, condition)
+            _apply_field_boundary_condition!(side, dim, grid, field, location(field), Ibc, condition)
         end
     end
 end