From 61bf925a51809f916965c0ac22bfebc0c08c3ee3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 31 Mar 2024 18:40:41 +0100
Subject: [PATCH 1/2] Special version of `view()` for `MPIDebugSharedArray`

We want `view(::MPIDebugSharedArray, I...)` to return an
`MPIDebugSharedArray` whose fields are views, rather than a view of an
`MPIDebugSharedArray`, treated as an `AbstractArray`. This allows other
special methods (e.g. `MPI.Buffer(::MPIDebugSharedArray)`) to work
properly.
---
 moment_kinetics/src/communication.jl | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/communication.jl b/moment_kinetics/src/communication.jl
index 9bbd07e9f..84b41202f 100644
--- a/moment_kinetics/src/communication.jl
+++ b/moment_kinetics/src/communication.jl
@@ -390,23 +390,23 @@ end
     Special type for debugging race conditions in accesses to shared-memory arrays.
     Only used if debugging._debug_level is high enough.
     """
-    struct DebugMPISharedArray{T, N} <: AbstractArray{T, N}
-        data::Array{T,N}
+    struct DebugMPISharedArray{T, N, TArray <: AbstractArray{T,N}, TIntArray <: AbstractArray{mk_int,N}, TBoolArray <: AbstractArray{Bool,N}} <: AbstractArray{T, N}
+        data::TArray
         accessed::Ref{Bool}
-        is_initialized::Array{mk_int,N}
-        is_read::Array{Bool,N}
-        is_written::Array{Bool, N}
+        is_initialized::TIntArray
+        is_read::TBoolArray
+        is_written::TBoolArray
         creation_stack_trace::String
         @debug_detect_redundant_block_synchronize begin
-            previous_is_read::Array{Bool,N}
-            previous_is_written::Array{Bool, N}
+            previous_is_read::TBoolArray
+            previous_is_written::TBoolArray
         end
     end
 
     export DebugMPISharedArray
 
     # Constructors
-    function DebugMPISharedArray(array::Array, comm)
+    function DebugMPISharedArray(array::AbstractArray, comm)
         dims = size(array)
         is_initialized = allocate_shared(mk_int, dims; comm=comm, maybe_debug=false)
         if block_rank[] == 0
@@ -484,6 +484,16 @@ end
               * "silently disable the debug checks")
     end
 
+    # Explicit overload for view() so the result is a DebugMPISharedArray
+    import Base: view
+    function view(A::DebugMPISharedArray, inds...)
+        return DebugMPISharedArray(
+            (isa(getfield(A, name), AbstractArray) ?
+             view(getfield(A, name), inds...) :
+             getfield(A, name)
+             for name ∈ fieldnames(typeof(A)))...)
+    end
+
     # Explicit overload for vec() so the result is a DebugMPISharedArray
     import Base: vec
     function vec(A::DebugMPISharedArray)

From 0cd6007be2e1f4284856f0292c1308eaacc0f4c6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 31 Mar 2024 18:43:36 +0100
Subject: [PATCH 2/2] Set read/write tracking flags when passing debug array to
 MPI.Buffer()

If a shared-memory array (or part of one) is passed to `MPI.Buffer()` it
will be communicated. This probably implies it is read on at least one
of the distributed-MPI processes, and read on others. To check for
errors as conservatively as possible (I think, but not sure if this is
really true), set the `is_read`, `is_written` and `is_initialized` and
`accessed` flags of an `MPIDebugSharedArray` when it is passed to
`MPI.Buffer()`, so that a shared-memory synchronization is required
before the array is used. This should be a sensible thing to do because:
distributed-MPI communication is done by the root process of each
shared-memory block, so any array that is touched by distributed-memory
communication requires a synchronization before processes in the
shared-memory block (apart from the root process) are allowed to use it.
---
 moment_kinetics/src/communication.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/moment_kinetics/src/communication.jl b/moment_kinetics/src/communication.jl
index 84b41202f..cd2bc53cd 100644
--- a/moment_kinetics/src/communication.jl
+++ b/moment_kinetics/src/communication.jl
@@ -513,6 +513,10 @@ end
 
     import MPI: Buffer
     function Buffer(A::DebugMPISharedArray)
+        A.is_initialized .= 1
+        A.is_read .= true
+        A.is_written .= true
+        A.accessed[] = true
         return Buffer(A.data)
     end