gt: link gtl on cray amd systems, more iters

Option to disable managed tests, TEST_MANAGED=OFF cmake var

.clang-format

@@ -0,0 +1,118 @@
+---
+Language:        Cpp
+# BasedOnStyle:  Mozilla
+AccessModifierOffset: -2
+AlignAfterOpenBracket: Align
+AlignConsecutiveAssignments: false
+AlignConsecutiveDeclarations: false
+AlignEscapedNewlines: Right
+AlignOperands:   true
+AlignTrailingComments: true
+AllowAllParametersOfDeclarationOnNextLine: false
+AllowShortBlocksOnASingleLine: false
+AllowShortCaseLabelsOnASingleLine: true
+AllowShortFunctionsOnASingleLine: Inline
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AlwaysBreakAfterDefinitionReturnType: None
+AlwaysBreakAfterReturnType: None
+AlwaysBreakBeforeMultilineStrings: false
+AlwaysBreakTemplateDeclarations: Yes
+BinPackArguments: true
+BinPackParameters: true
+BraceWrapping:   
+  AfterClass:      true
+  AfterControlStatement: false
+  AfterEnum:       true
+  AfterFunction:   true
+  AfterNamespace:  true
+  AfterObjCDeclaration: false
+  AfterStruct:     true
+  AfterUnion:      true
+  AfterExternBlock: false
+  BeforeCatch:     false
+  BeforeElse:      false
+  IndentBraces:    false
+  SplitEmptyFunction: true
+  SplitEmptyRecord: false
+  SplitEmptyNamespace: true
+BreakBeforeBinaryOperators: None
+BreakBeforeBraces: Custom
+BreakBeforeInheritanceComma: false
+BreakInheritanceList: BeforeComma
+BreakBeforeTernaryOperators: true
+BreakConstructorInitializersBeforeComma: false
+BreakConstructorInitializers: BeforeColon
+BreakAfterJavaFieldAnnotations: false
+BreakStringLiterals: true
+ColumnLimit:     80
+CommentPragmas:  '^ IWYU pragma:'
+CompactNamespaces: false
+ConstructorInitializerAllOnOneLineOrOnePerLine: true
+ConstructorInitializerIndentWidth: 2
+ContinuationIndentWidth: 2
+Cpp11BracedListStyle: true
+DerivePointerAlignment: false
+DisableFormat:   false
+ExperimentalAutoDetectBinPacking: false
+FixNamespaceComments: true
+ForEachMacros:   
+  - foreach
+  - Q_FOREACH
+  - BOOST_FOREACH
+IncludeBlocks:   Preserve
+IncludeCategories: 
+  - Regex:           '^"(llvm|llvm-c|clang|clang-c)/'
+    Priority:        2
+  - Regex:           '^(<|"(gtest|gmock|isl|json)/)'
+    Priority:        3
+  - Regex:           '.*'
+    Priority:        1
+IncludeIsMainRegex: '(Test)?$'
+IndentCaseLabels: true
+IndentPPDirectives: None
+IndentWidth:     2
+IndentWrappedFunctionNames: false
+JavaScriptQuotes: Leave
+JavaScriptWrapImports: true
+KeepEmptyLinesAtTheStartOfBlocks: true
+MacroBlockBegin: ''
+MacroBlockEnd:   ''
+MaxEmptyLinesToKeep: 1
+NamespaceIndentation: None
+ObjCBinPackProtocolList: Auto
+ObjCBlockIndentWidth: 2
+ObjCSpaceAfterProperty: true
+ObjCSpaceBeforeProtocolList: false
+PenaltyBreakAssignment: 2
+PenaltyBreakBeforeFirstCallParameter: 19
+PenaltyBreakComment: 300
+PenaltyBreakFirstLessLess: 120
+PenaltyBreakString: 1000
+PenaltyBreakTemplateDeclaration: 10
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 200
+PointerAlignment: Left
+ReflowComments:  true
+SortIncludes:    true
+SortUsingDeclarations: true
+SpaceAfterCStyleCast: false
+SpaceAfterTemplateKeyword: true
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeCpp11BracedList: false
+SpaceBeforeCtorInitializerColon: true
+SpaceBeforeInheritanceColon: true
+SpaceBeforeParens: ControlStatements
+SpaceBeforeRangeBasedForLoopColon: true
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles:  false
+SpacesInContainerLiterals: true
+SpacesInCStyleCastParentheses: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Standard:        Cpp11
+TabWidth:        8
+UseTab:          Never
+...
+

CMakeLists.txt

@@ -3,34 +3,59 @@ cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
 # create project
 project(mpi-daxpy-test)
 
+option(TEST_MANAGED "Test managed memory" ON)
+
 # add dependencies
 include(cmake/CPM.cmake)
 CPMFindPackage(NAME gtensor
-               GITHUB_REPOSITORY bd4/gtensor
+               GITHUB_REPOSITORY wdmapp/gtensor
-               GIT_TAG "pr/sycl-include-refactor"
+               GIT_TAG "main"
                OPTIONS "GTENSOR_ENABLE_BLAS ON")
 
+set(MPI_CXX_SKIP_MPICXX ON)
 find_package(MPI REQUIRED)
 
-add_executable(mpi_daxpy_gt)
-target_sources(mpi_daxpy_gt PRIVATE mpi_daxpy_gt.cc)
-target_link_libraries(mpi_daxpy_gt gtensor::gtensor)
-target_link_libraries(mpi_daxpy_gt gtensor::blas)
-target_link_libraries(mpi_daxpy_gt MPI::MPI_CXX)
-
-add_executable(mpi_stencil_gt)
-target_sources(mpi_stencil_gt PRIVATE mpi_stencil_gt.cc)
-target_link_libraries(mpi_stencil_gt gtensor::gtensor)
-target_link_libraries(mpi_stencil_gt MPI::MPI_CXX)
-
 if ("${GTENSOR_DEVICE}" STREQUAL "cuda")
   enable_language(CUDA)
+endif()
+
+add_executable(mpi_daxpy_gt)
+target_sources(mpi_daxpy_gt PRIVATE mpi_daxpy_gt.cc)
+target_link_libraries(mpi_daxpy_gt PRIVATE gtensor::gtensor)
+target_link_libraries(mpi_daxpy_gt PRIVATE gtensor::blas)
+target_link_libraries(mpi_daxpy_gt PRIVATE MPI::MPI_CXX)
+
+add_executable(mpi_stencil_gt)
+target_sources(mpi_stencil_gt PRIVATE mpi_stencil_gt.cc)
+target_link_libraries(mpi_stencil_gt PRIVATE gtensor::gtensor)
+target_link_libraries(mpi_stencil_gt PRIVATE MPI::MPI_CXX)
+
+add_executable(mpi_stencil2d_gt)
+target_sources(mpi_stencil2d_gt PRIVATE mpi_stencil2d_gt.cc)
+target_link_libraries(mpi_stencil2d_gt PRIVATE gtensor::gtensor)
+target_link_libraries(mpi_stencil2d_gt PRIVATE MPI::MPI_CXX)
+#target_compile_features(mpi_stencil2d_gt PRIVATE cxx_std_17)
+if (TEST_MANAGED)
+  message(STATUS "${PROJECT_NAME}: Enabling managed memory")
+  target_compile_definitions(mpi_stencil2d_gt PRIVATE TEST_MANAGED)
+endif()
+if (GTENSOR_DEVICE STREQUAL "hip" AND DEFINED ENV{PE_MPICH_GTL_DIR_amd_gfx90a})
+  message(STATUS "${PROJECT_NAME}: Linking gtl libs for HIP backend")
+  target_link_options(mpi_stencil2d_gt PRIVATE
+                      $ENV{PE_MPICH_GTL_DIR_amd_gfx90a}
+                      $ENV{PE_MPICH_GTL_LIBS_amd_gfx90a})
+endif()
+
+if ("${GTENSOR_DEVICE}" STREQUAL "cuda") 
   set_source_files_properties(mpi_daxpy_gt.cc
                               TARGET_DIRECTORY mpi_daxpy_gt
                               PROPERTIES LANGUAGE CUDA)
   set_source_files_properties(mpi_stencil_gt.cc
                               TARGET_DIRECTORY mpi_stencil_gt
                               PROPERTIES LANGUAGE CUDA)
+  set_source_files_properties(mpi_stencil2d_gt.cc
+                              TARGET_DIRECTORY mpi_stencil2d_gt
+                              PROPERTIES LANGUAGE CUDA)
 else()
   set_source_files_properties(mpi_daxpy_gt.cc
                               TARGET_DIRECTORY mpi_daxpy_gt
@@ -38,4 +63,21 @@ else()
   set_source_files_properties(mpi_stencil_gt.cc
                               TARGET_DIRECTORY mpi_stencil_gt
                               PROPERTIES LANGUAGE CXX)
+  set_source_files_properties(mpi_stencil2d_gt.cc
+                              TARGET_DIRECTORY mpi_stencil2d_gt
+                              PROPERTIES LANGUAGE CXX)
+endif()
+
+if ("${GTENSOR_DEVICE}" STREQUAL "sycl")
+  add_executable(mpi_stencil2d_sycl)
+  target_sources(mpi_stencil2d_sycl PRIVATE mpi_stencil2d_sycl.cc)
+  target_link_libraries(mpi_stencil2d_sycl MPI::MPI_CXX)
+  target_compile_options(mpi_stencil2d_sycl PRIVATE -fsycl -x c++)
+  target_link_options(mpi_stencil2d_sycl PRIVATE -fsycl)
+
+  add_executable(mpi_stencil2d_sycl_oo)
+  target_sources(mpi_stencil2d_sycl_oo PRIVATE mpi_stencil2d_sycl_oo.cc)
+  target_link_libraries(mpi_stencil2d_sycl_oo MPI::MPI_CXX)
+  target_compile_options(mpi_stencil2d_sycl_oo PRIVATE -fsycl -x c++)
+  target_link_options(mpi_stencil2d_sycl_oo PRIVATE -fsycl)
 endif()

mpi_daxpy_gt.cc

@@ -73,7 +73,7 @@ int main(int argc, char **argv) {
     device_id = gt::backend::clib::device_get();
     vendor_id = gt::backend::clib::device_get_vendor_id(device_id);
 
-    gt::blas::handle_t* h = gt::blas::create();
+    gt::blas::handle_t h;
 
     gt::copy(x, d_x);
     gt::copy(y, d_y);

mpi_stencil2d_gt.cc

@@ -0,0 +1,734 @@
+/*
+ * Test GPU aware MPI on different platforms using a distributed
+ * 1d stencil on a 2d array. The exchange in second (non-contiguous)
+ * direction forces use of staging buffers, which replicates what
+ * is needed for all but the innermost dimension exchanges in the
+ * GENE fusion code.
+ *
+ * Takes optional command line arg for size of each dimension of the domain
+ * n_global, in 1024 increments. Default is 8 * 1024 (so 256K plus ghost points
+ * in size for doulbles per array), which should fit on any system but may not
+ * be enough to tax larger HPC GPUs and MPI impelmentations.
+ *
+ * There will be four exchange buffers of size 2 * n_global, i.e. 128K each
+ * by default.
+ *
+ * Gtensor is used so a single source can be used for all platforms.
+ */
+
+#include <cmath>
+#include <mpi.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+
+#include "gtensor/gtensor.h"
+#include "gtensor/reductions.h"
+
+using namespace gt::placeholders;
+
+const double PI = 3.141592653598793;
+
+inline void check(const char* fname, int line, int mpi_rval)
+{
+  if (mpi_rval != MPI_SUCCESS) {
+    printf("%s:%d error %d\n", fname, line, mpi_rval);
+    exit(2);
+  }
+}
+
+#define CHECK(x) check(__FILE__, __LINE__, (x))
+
+// little hack to make code parameterizable on managed vs device memory
+namespace gt
+{
+
+namespace ext
+{
+namespace detail
+{
+
+template <typename T, gt::size_type N, typename S = gt::space::device>
+struct gthelper
+{
+  using gtensor = gt::gtensor<T, N, S>;
+};
+
+#ifdef GTENSOR_HAVE_DEVICE
+
+template <typename T, gt::size_type N>
+struct gthelper<T, N, gt::space::managed>
+{
+  using gtensor = gt::gtensor_container<gt::space::managed_vector<T>, N>;
+};
+#endif
+
+} // namespace detail
+
+template <typename T, gt::size_type N, typename S = gt::space::device>
+using gtensor2 = typename detail::gthelper<T, N, S>::gtensor;
+
+} // namespace ext
+
+} // namespace gt
+
+static const gt::gtensor<double, 1> stencil5 = {1.0 / 12.0, -2.0 / 3.0, 0.0,
+                                                2.0 / 3.0, -1.0 / 12.0};
+
+/*
+ * Return unevaluated expression that calculates the 1d stencil in the
+ * first dimension of a 2d array.
+ *
+ * Size of the result will be size of z with minus 4 in second dimension.
+ */
+template <typename S>
+inline auto stencil2d_1d_5_d0(const gt::ext::gtensor2<double, 2, S>& z,
+                              const gt::gtensor<double, 1>& stencil)
+{
+  return stencil(0) * z.view(_s(0, -4), _all) +
+         stencil(1) * z.view(_s(1, -3), _all) +
+         stencil(2) * z.view(_s(2, -2), _all) +
+         stencil(3) * z.view(_s(3, -1), _all) +
+         stencil(4) * z.view(_s(4, _), _all);
+}
+
+/*
+ * Return unevaluated expression that calculates the 1d stencil in the
+ * second dimension of a 2d array.
+ *
+ * Size of the result will be size of z with minus 4 in second dimension.
+ */
+template <typename S>
+inline auto stencil2d_1d_5_d1(const gt::ext::gtensor2<double, 2, S>& z,
+                              const gt::gtensor<double, 1>& stencil)
+{
+  return stencil(0) * z.view(_all, _s(0, -4)) +
+         stencil(1) * z.view(_all, _s(1, -3)) +
+         stencil(2) * z.view(_all, _s(2, -2)) +
+         stencil(3) * z.view(_all, _s(3, -1)) +
+         stencil(4) * z.view(_all, _s(4, _));
+}
+
+void set_rank_device(int n_ranks, int rank)
+{
+  int n_devices, device, ranks_per_device;
+
+  n_devices = gt::backend::clib::device_get_count();
+
+  if (n_ranks > n_devices) {
+    if (n_ranks % n_devices != 0) {
+      printf(
+        "ERROR: Number of ranks (%d) not a multiple of number of GPUs (%d)\n",
+        n_ranks, n_devices);
+      exit(EXIT_FAILURE);
+    }
+    ranks_per_device = n_ranks / n_devices;
+    device = rank / ranks_per_device;
+  } else {
+    ranks_per_device = 1;
+    device = rank;
+  }
+
+  gt::backend::clib::device_set(device);
+}
+
+// exchange in first dimension, staging into contiguous buffers on device
+template <typename S>
+void boundary_exchange_x(MPI_Comm comm, int world_size, int rank,
+                         gt::ext::gtensor2<double, 2, S>& d_z, int n_bnd,
+                         bool stage_host = false)
+{
+  auto buf_shape = gt::shape(n_bnd, d_z.shape(1));
+  gt::gtensor_device<double, 2> sbuf_l(buf_shape);
+  gt::gtensor_device<double, 2> sbuf_r(buf_shape);
+  gt::gtensor_device<double, 2> rbuf_r(buf_shape);
+  gt::gtensor_device<double, 2> rbuf_l(buf_shape);
+
+  gt::shape_type<2> host_buf_shape;
+  if (stage_host) {
+    host_buf_shape = buf_shape;
+  } else {
+    host_buf_shape = {0, 0};
+  }
+  gt::gtensor<double, 2> h_sbuf_l(host_buf_shape);
+  gt::gtensor<double, 2> h_sbuf_r(host_buf_shape);
+  gt::gtensor<double, 2> h_rbuf_r(host_buf_shape);
+  gt::gtensor<double, 2> h_rbuf_l(host_buf_shape);
+
+  MPI_Request req_l[2];
+  MPI_Request req_r[2];
+
+  int rank_l = rank - 1;
+  int rank_r = rank + 1;
+
+  // start async copy of ghost points into send buffers
+  if (rank_l >= 0) {
+    sbuf_l = d_z.view(_s(n_bnd, 2 * n_bnd), _all);
+    if (stage_host) {
+      gt::copy(sbuf_l, h_sbuf_l);
+    }
+  }
+  if (rank_r <= world_size) {
+    sbuf_r = d_z.view(_s(-2 * n_bnd, -n_bnd), _all);
+    if (stage_host) {
+      gt::copy(sbuf_r, h_sbuf_r);
+    }
+  }
+
+  // initiate async recv
+  if (rank_l >= 0) {
+    double* rbuf_l_data = nullptr;
+    if (stage_host) {
+      rbuf_l_data = h_rbuf_l.data();
+    } else {
+      rbuf_l_data = rbuf_l.data().get();
+    }
+    CHECK(MPI_Irecv(rbuf_l_data, rbuf_l.size(), MPI_DOUBLE, rank_l, 123, comm,
+                    &req_l[0]));
+  }
+
+  if (rank_r < world_size) {
+    double* rbuf_r_data = nullptr;
+    if (stage_host) {
+      rbuf_r_data = h_rbuf_r.data();
+    } else {
+      rbuf_r_data = rbuf_r.data().get();
+    }
+    CHECK(MPI_Irecv(rbuf_r_data, rbuf_r.size(), MPI_DOUBLE, rank_r, 456, comm,
+                    &req_r[0]));
+  }
+
+  // wait for send buffer fill
+  gt::synchronize();
+
+  // initiate async sends
+  if (rank_l >= 0) {
+    double* sbuf_l_data = nullptr;
+    if (stage_host) {
+      sbuf_l_data = h_sbuf_l.data();
+    } else {
+      sbuf_l_data = sbuf_l.data().get();
+    }
+    CHECK(MPI_Isend(sbuf_l_data, sbuf_l.size(), MPI_DOUBLE, rank_l, 456, comm,
+                    &req_l[1]));
+  }
+
+  if (rank_r < world_size) {
+    double* sbuf_r_data = nullptr;
+    if (stage_host) {
+      sbuf_r_data = h_sbuf_r.data();
+    } else {
+      sbuf_r_data = sbuf_r.data().get();
+    }
+    CHECK(MPI_Isend(sbuf_r_data, sbuf_r.size(), MPI_DOUBLE, rank_r, 123, comm,
+                    &req_r[1]));
+  }
+
+  // wait for send/recv to complete, then copy data back into main data array
+  int mpi_rval;
+  if (rank_l >= 0) {
+    MPI_Status status[2];
+    mpi_rval = MPI_Waitall(2, req_l, status);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("send_l error: %d (%d %d)\n", mpi_rval, status[0].MPI_ERROR,
+             status[1].MPI_ERROR);
+    }
+    if (stage_host) {
+      gt::copy(h_rbuf_l, rbuf_l);
+    }
+    d_z.view(_s(0, n_bnd), _all) = rbuf_l;
+  }
+  if (rank_r < world_size) {
+    MPI_Status status[2];
+    mpi_rval = MPI_Waitall(2, req_r, status);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("send_r error: %d (%d %d)\n", mpi_rval, status[0].MPI_ERROR,
+             status[1].MPI_ERROR);
+    }
+    if (stage_host) {
+      gt::copy(h_rbuf_r, rbuf_r);
+    }
+    d_z.view(_s(-n_bnd, _), _all) = rbuf_r;
+  }
+
+  gt::synchronize();
+}
+
+// exchange in second dimension, optional staging into device buffer
+template <typename S>
+void boundary_exchange_y(MPI_Comm comm, int world_size, int rank,
+                         gt::ext::gtensor2<double, 2, S>& d_z, int n_bnd,
+                         bool stage_device)
+{
+  gt::shape_type<2> buf_shape;
+  if (stage_device) {
+    buf_shape = gt::shape(d_z.shape(0), n_bnd);
+  } else {
+    buf_shape = {0, 0};
+  }
+
+  gt::gtensor_device<double, 2> sbuf_l(buf_shape);
+  gt::gtensor_device<double, 2> sbuf_r(buf_shape);
+  gt::gtensor_device<double, 2> rbuf_r(buf_shape);
+  gt::gtensor_device<double, 2> rbuf_l(buf_shape);
+
+  MPI_Request req_l[2];
+  MPI_Request req_r[2];
+
+  int rank_l = rank - 1;
+  int rank_r = rank + 1;
+
+  auto sv_l = gt::view_strided(d_z, _all, _s(n_bnd, 2 * n_bnd));
+  auto sv_r = gt::view_strided(d_z, _all, _s(-2 * n_bnd, -n_bnd));
+  auto rv_l = gt::view_strided(d_z, _all, _s(0, n_bnd));
+  auto rv_r = gt::view_strided(d_z, _all, _s(-n_bnd, _));
+
+  // start async copy of ghost points into send buffers
+  if (stage_device) {
+    if (rank_l >= 0) {
+      sbuf_l = sv_l;
+    }
+    if (rank_r <= world_size) {
+      sbuf_r = sv_r;
+    }
+  }
+
+  // initiate async recv
+  if (rank_l >= 0) {
+    double* rbuf_l_data = nullptr;
+    if (stage_device) {
+      rbuf_l_data = rbuf_l.data().get();
+    } else {
+      rbuf_l_data = rv_l.data().get();
+    }
+    CHECK(MPI_Irecv(rbuf_l_data, rv_l.size(), MPI_DOUBLE, rank_l, 123, comm,
+                    &req_l[0]));
+  }
+
+  if (rank_r < world_size) {
+    double* rbuf_r_data = nullptr;
+    if (stage_device) {
+      rbuf_r_data = rbuf_r.data().get();
+    } else {
+      rbuf_r_data = rv_r.data().get();
+    }
+    CHECK(MPI_Irecv(rbuf_r_data, rv_r.size(), MPI_DOUBLE, rank_r, 456, comm,
+                    &req_r[0]));
+  }
+
+  // wait for send buffer fill
+  // if (stage_device) {
+  gt::synchronize();
+  // }
+
+  // initiate async sends
+  if (rank_l >= 0) {
+    double* sbuf_l_data = nullptr;
+    if (stage_device) {
+      sbuf_l_data = sbuf_l.data().get();
+    } else {
+      sbuf_l_data = sv_l.data().get();
+    }
+    CHECK(MPI_Isend(sbuf_l_data, sv_l.size(), MPI_DOUBLE, rank_l, 456, comm,
+                    &req_l[1]));
+  }
+
+  if (rank_r < world_size) {
+    double* sbuf_r_data = nullptr;
+    if (stage_device) {
+      sbuf_r_data = sbuf_r.data().get();
+    } else {
+      sbuf_r_data = sv_r.data().get();
+    }
+    CHECK(MPI_Isend(sbuf_r_data, sv_r.size(), MPI_DOUBLE, rank_r, 123, comm,
+                    &req_r[1]));
+  }
+
+  // wait for send/recv to complete, then copy data back into main data array
+  int mpi_rval;
+  if (rank_l >= 0) {
+    MPI_Status status[2];
+    mpi_rval = MPI_Waitall(2, req_l, status);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("send_l error: %d (%d %d)\n", mpi_rval, status[0].MPI_ERROR,
+             status[1].MPI_ERROR);
+    }
+    if (stage_device) {
+      gt::copy(rbuf_l, rv_l);
+    }
+  }
+  if (rank_r < world_size) {
+    MPI_Status status[2];
+    mpi_rval = MPI_Waitall(2, req_r, status);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("send_r error: %d (%d %d)\n", mpi_rval, status[0].MPI_ERROR,
+             status[1].MPI_ERROR);
+    }
+    if (stage_device) {
+      gt::copy(rbuf_r, rv_r);
+    }
+  }
+
+  gt::synchronize();
+}
+
+template <int Dim, typename S>
+void print_test_name(bool use_buffers)
+{
+  if constexpr (std::is_same<S, gt::space::device>::value) {
+    printf("TEST dim:%d, device , buf:%d", Dim, use_buffers);
+  } else {
+    printf("TEST dim:%d, managed, buf:%d", Dim, use_buffers);
+  }
+}
+
+template <typename S, int Dim>
+void test_deriv(int device_id, uint32_t vendor_id, int world_size,
+                int world_rank, int n_global, int n_iter, bool use_buffers,
+                int n_warmup = 5)
+{
+  // Note: domain will be n_global x n_global plus ghost points in one dimension
+
+  int n_sten = 5;
+  int n_bnd = (n_sten - 1) / 2;
+
+  const int n_local = n_global / world_size;
+
+  int nx_local, ny_local;
+  int nx_local_ghost, ny_local_ghost;
+  int nx_bnd, ny_bnd;
+
+  if constexpr (Dim == 0) {
+    nx_bnd = n_bnd;
+    ny_bnd = 0;
+    nx_local = n_local;
+    nx_local_ghost = n_local + 2 * n_bnd;
+    ny_local = n_global;
+    ny_local_ghost = n_global;
+  } else {
+    nx_bnd = 0;
+    ny_bnd = n_bnd;
+    nx_local = n_global;
+    nx_local_ghost = n_global;
+    ny_local = n_local;
+    ny_local_ghost = n_local + 2 * n_bnd;
+  }
+
+  gt::shape_type<2> z_shape(nx_local_ghost, ny_local_ghost);
+  gt::shape_type<2> dz_shape(nx_local, ny_local);
+
+  auto h_z = gt::empty<double>(z_shape);
+  gt::ext::gtensor2<double, 2, S> d_z(z_shape);
+
+  auto h_dz_numeric = gt::empty<double>(dz_shape);
+  auto h_dz_actual = gt::empty<double>(dz_shape);
+  gt::ext::gtensor2<double, 2, S> d_dz_numeric(dz_shape);
+
+  double ln = 8;
+  double delta = ln / n_global;
+  double ln_local = ln / world_size;
+  double scale = n_global / ln;
+  auto fn = [](double x, double y) { return x * x * x + y * y; };
+  auto fn_dzdx = [](double x, double y) { return 3 * x * x; };
+  auto fn_dzdy = [](double x, double y) { return 2 * y; };
+
+  struct timespec start, end;
+  double iter_time = 0.0;
+  double total_time = 0.0;
+
+  double x_start = 0, y_start = 0;
+  if constexpr (Dim == 0) {
+    x_start = world_rank * ln_local;
+  } else {
+    y_start = world_rank * ln_local;
+  }
+  for (int j = 0; j < ny_local; j++) {
+    double ytmp = y_start + j * delta;
+    for (int i = 0; i < nx_local; i++) {
+      double xtmp = x_start + i * delta;
+      h_z(i + nx_bnd, j + ny_bnd) = fn(xtmp, ytmp);
+      if constexpr (Dim == 0) {
+        h_dz_actual(i, j) = fn_dzdx(xtmp, ytmp);
+      } else {
+        h_dz_actual(i, j) = fn_dzdy(xtmp, ytmp);
+      }
+    }
+  }
+
+  // fill boundary points on ends
+  if constexpr (Dim == 0) {
+    if (world_rank == 0) {
+      for (int j = 0; j < ny_local; j++) {
+        double ytmp = j * delta;
+        for (int i = 0; i < nx_bnd; i++) {
+          double xtmp = (i - nx_bnd) * delta;
+          h_z(i, j) = fn(xtmp, ytmp);
+        }
+      }
+    }
+    if (world_rank == world_size - 1) {
+      for (int j = 0; j < ny_local; j++) {
+        double ytmp = j * delta;
+        for (int i = 0; i < nx_bnd; i++) {
+          double xtmp = ln + i * delta;
+          h_z(nx_bnd + nx_local + i, j) = fn(xtmp, ytmp);
+        }
+      }
+    }
+  } else {
+    if (world_rank == 0) {
+      for (int j = 0; j < ny_bnd; j++) {
+        double ytmp = (j - ny_bnd) * delta;
+        for (int i = 0; i < nx_local; i++) {
+          double xtmp = i * delta;
+          h_z(i, j) = fn(xtmp, ytmp);
+        }
+      }
+    }
+    if (world_rank == world_size - 1) {
+      for (int j = 0; j < ny_bnd; j++) {
+        double ytmp = ln + j * delta;
+        for (int i = 0; i < nx_local; i++) {
+          double xtmp = i * delta;
+          h_z(i, ny_bnd + ny_local + j) = fn(xtmp, ytmp);
+        }
+      }
+    }
+  }
+
+  /*
+  for (int i = 0; i < 5; i++) {
+    printf("%d row1-l %f\n", world_rank, h_z(1, i));
+  }
+  for (int i = 0; i < 5; i++) {
+    printf("%d row1-r %f\n", world_rank, h_z(1, n_local_with_ghost - 1 - i));
+  }
+  */
+
+  gt::copy(h_z, d_z);
+  // gt::synchronize();
+
+  for (int i = 0; i < n_warmup + n_iter; i++) {
+    clock_gettime(CLOCK_MONOTONIC, &start);
+    if constexpr (Dim == 0) {
+      boundary_exchange_x<S>(MPI_COMM_WORLD, world_size, world_rank, d_z, n_bnd,
+                             use_buffers);
+    } else {
+      boundary_exchange_y<S>(MPI_COMM_WORLD, world_size, world_rank, d_z, n_bnd,
+                             use_buffers);
+    }
+    clock_gettime(CLOCK_MONOTONIC, &end);
+    iter_time =
+      ((end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) * 1.0e-9);
+
+    if (i >= n_warmup) {
+      total_time += iter_time;
+    }
+
+    // do some calculation, to try to more closely simulate what happens in GENE
+    if constexpr (Dim == 0) {
+      d_dz_numeric = stencil2d_1d_5_d0<S>(d_z, stencil5) * scale;
+    } else {
+      d_dz_numeric = stencil2d_1d_5_d1<S>(d_z, stencil5) * scale;
+    }
+    gt::synchronize();
+  }
+#ifdef DEBUG
+  printf("%d/%d exchange time %0.8f ms\n", world_rank, world_size,
+         total_time / n_iter * 1000);
+#endif
+
+  gt::copy(d_dz_numeric, h_dz_numeric);
+
+  /*
+  for (int i = 0; i < 5; i++) {
+    printf("%d la %f\n%d ln %f\n", world_rank, h_dzdx_actual(8, i), world_rank,
+           h_dzdx_numeric(8, i));
+  }
+  for (int i = 0; i < 5; i++) {
+    int idx = n_local - 1 - i;
+    printf("%d ra %f\n%d rn %f\n", world_rank, h_dzdx_actual(8, idx),
+           world_rank, h_dzdx_numeric(8, idx));
+  }
+  */
+
+  double err_norm = std::sqrt(gt::sum_squares(h_dz_numeric - h_dz_actual));
+
+#ifdef DEBUG
+  printf("%d/%d [%d:0x%08x] err_norm = %.8f\n", world_rank, world_size,
+         device_id, vendor_id, err_norm);
+#endif
+
+  double time_sum;
+  MPI_Reduce(&total_time, &time_sum, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
+
+  double err_sum;
+  MPI_Reduce(&err_norm, &err_sum, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
+
+  if (world_rank == 0) {
+    print_test_name<Dim, S>(use_buffers);
+    printf("; %0.8f, err=%0.8f\n", time_sum, err_sum);
+  }
+}
+
+template <typename S, int Dim>
+void test_sum(int device_id, uint32_t vendor_id, int world_size, int world_rank,
+              int n_global, int n_iter, int n_warmup = 5)
+{
+  // Note: domain will be n_global x n_global plus ghost points in one dimension
+  const int n_local = n_global / world_size;
+
+  int nx_local, ny_local;
+
+  struct timespec start, end;
+  double iter_time = 0.0;
+  double total_time = 0.0;
+
+  if constexpr (Dim == 0) {
+    nx_local = n_local;
+    ny_local = n_global;
+  } else {
+    nx_local = n_global;
+    ny_local = n_local;
+  }
+
+  gt::shape_type<2> z_shape(nx_local, ny_local);
+
+  gt::ext::gtensor2<double, 2, S> d_z(z_shape, PI / world_size);
+
+  // reduction test
+  gt::shape_type<1> sum_shape;
+  if constexpr (Dim == 0) {
+    sum_shape = gt::shape(d_z.shape(0));
+  } else {
+    sum_shape = gt::shape(d_z.shape(1));
+  }
+  gt::ext::gtensor2<double, 1, S> d_sum(sum_shape);
+  gt::gtensor<double, 1> h_sum(sum_shape);
+  for (int i = 0; i < n_warmup + n_iter; i++) {
+    if constexpr (Dim == 0) {
+      gt::sum_axis_to(d_sum, d_z, 0);
+      gt::synchronize();
+
+      clock_gettime(CLOCK_MONOTONIC, &start);
+      CHECK(MPI_Allreduce(MPI_IN_PLACE, d_sum.data().get(), d_sum.size(),
+                          MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD));
+
+      clock_gettime(CLOCK_MONOTONIC, &end);
+    } else {
+      gt::sum_axis_to(d_sum, d_z, 1);
+      gt::synchronize();
+
+      clock_gettime(CLOCK_MONOTONIC, &start);
+      CHECK(MPI_Allreduce(MPI_IN_PLACE, d_sum.data().get(), d_sum.size(),
+                          MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD));
+      clock_gettime(CLOCK_MONOTONIC, &end);
+    }
+    iter_time =
+      ((end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) * 1.0e-9);
+
+    if (i >= n_warmup) {
+      total_time += iter_time;
+    }
+  }
+#ifdef DEBUG
+  printf("%d/%d allreduce time %0.8f ms\n", world_rank, world_size,
+         total_time / n_iter * 1000);
+#endif
+
+  gt::copy(d_sum, h_sum);
+
+  double time_sum;
+  MPI_Reduce(&total_time, &time_sum, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
+
+  if (world_rank == 0) {
+    print_test_name<Dim, S>(false);
+    printf("; allreduce=%0.8f\n", time_sum);
+  }
+}
+
+int main(int argc, char** argv)
+{
+  using S = gt::space::managed;
+
+  // Note: domain will be n_global x n_global plus ghost points in one dimension
+  int n_global = 8 * 1024;
+  int n_iter = 1000;
+  int n_warmup = 10;
+
+  if (argc > 1) {
+    n_global = std::atoi(argv[1]) * 1024;
+  }
+  if (argc > 2) {
+    n_iter = std::atoi(argv[2]);
+  }
+
+  int world_size, world_rank, device_id;
+  uint32_t vendor_id;
+
+  CHECK(MPI_Init(NULL, NULL));
+
+  CHECK(MPI_Comm_size(MPI_COMM_WORLD, &world_size));
+  CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &world_rank));
+
+  if (n_global % world_size != 0) {
+    printf("%d nmpi (%d) must be divisor of domain size (%d), exiting\n",
+           world_rank, world_size, n_global);
+    exit(1);
+  }
+
+  const int n_local = n_global / world_size;
+
+  set_rank_device(world_size, world_rank);
+  device_id = gt::backend::clib::device_get();
+  vendor_id = gt::backend::clib::device_get_vendor_id(device_id);
+
+  if (world_rank == 0) {
+    printf("n procs    = %d\n", world_size);
+    printf("n_global   = %d\n", n_global);
+    printf("n_local    = %d\n", n_local);
+    printf("n_iter     = %d\n", n_iter);
+    printf("n_warmup   = %d\n", n_warmup);
+  }
+
+  fflush(stdout);
+
+  test_deriv<gt::space::device, 0>(device_id, vendor_id, world_size, world_rank,
+                                   n_global, n_iter, true, 5);
+  test_deriv<gt::space::device, 0>(device_id, vendor_id, world_size, world_rank,
+                                   n_global, n_iter, false, 5);
+#ifdef TEST_MANAGED
+  test_deriv<gt::space::managed, 0>(device_id, vendor_id, world_size,
+                                    world_rank, n_global, n_iter, true, 5);
+  test_deriv<gt::space::managed, 0>(device_id, vendor_id, world_size,
+                                    world_rank, n_global, n_iter, false, 5);
+#endif
+
+  test_deriv<gt::space::device, 1>(device_id, vendor_id, world_size, world_rank,
+                                   n_global, n_iter, true, 5);
+  test_deriv<gt::space::device, 1>(device_id, vendor_id, world_size, world_rank,
+                                   n_global, n_iter, false, 5);
+#ifdef TEST_MANAGED
+  test_deriv<gt::space::managed, 1>(device_id, vendor_id, world_size,
+                                    world_rank, n_global, n_iter, true, 5);
+  test_deriv<gt::space::managed, 1>(device_id, vendor_id, world_size,
+                                    world_rank, n_global, n_iter, false, 5);
+#endif
+
+  test_sum<gt::space::device, 0>(device_id, vendor_id, world_size, world_rank,
+                                 n_global, n_iter, 5);
+#ifdef TEST_MANAGED
+  test_sum<gt::space::managed, 0>(device_id, vendor_id, world_size, world_rank,
+                                  n_global, n_iter, 5);
+#endif
+  test_sum<gt::space::device, 1>(device_id, vendor_id, world_size, world_rank,
+                                 n_global, n_iter, 5);
+#ifdef TEST_MANAGED
+  test_sum<gt::space::managed, 1>(device_id, vendor_id, world_size, world_rank,
+                                  n_global, n_iter, 5);
+#endif
+
+  MPI_Finalize();
+
+  return EXIT_SUCCESS;
+}

mpi_stencil2d_sycl.cc

@@ -0,0 +1,547 @@
+/*
+ * Test GPU aware MPI on different platforms using a distributed
+ * 1d stencil on a 2d array. The exchange in second (non-contiguous)
+ * direction forces use of staging buffers, which replicates what
+ * is needed for all but the innermost dimension exchanges in the
+ * GENE fusion code.
+ *
+ * Takes optional command line arg for size of each dimension of the domain
+ * n_global, in 1024 increments. Default is 8 * 1024 (so 256K plus ghost points
+ * in size for doulbles per array), which should fit on any system but may not
+ * be enough to tax larger HPC GPUs and MPI impelmentations.
+ *
+ * There will be four exchange buffers of size 2 * n_global, i.e. 128K each
+ * by default.
+ */
+
+#include <cassert>
+#include <cmath>
+#include <iostream>
+#include <mpi.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+
+#include "sycl/sycl.hpp"
+
+inline void check(const char* fname, int line, int mpi_rval)
+{
+  if (mpi_rval != MPI_SUCCESS) {
+    printf("%s:%d error %d\n", fname, line, mpi_rval);
+    exit(2);
+  }
+}
+
+#define CHECK(x) check(__FILE__, __LINE__, (x))
+
+static constexpr double stencil5[] = {1.0 / 12.0, -2.0 / 3.0, 0.0, 2.0 / 3.0,
+                                      -1.0 / 12.0};
+
+inline constexpr std::size_t idx2(int n, int row, int col)
+{
+  return row + col * n;
+}
+
+/*
+ * Calculate 1d stencil of second dimension of 2d array on GPU. Out array must
+ * be contiguous column major nrows x ncols array, while in array must be
+ * (nrows)x(ncols+4) to accomodate 2 ghost points in each direction for the
+ * second dimension.
+ *
+ * Returns sycl event, async with respect to host.
+ */
+auto stencil2d_1d_5(sycl::queue& q, int out_nrows, int out_ncols, double* out2d,
+                    const double* in2d, double scale)
+{
+  // Note: swap index order; SYCL is row-major oriented, and this example
+  // is col-major
+  int in_nrows = out_nrows + 4;
+  auto range = sycl::range<2>(out_ncols, out_nrows);
+  auto e = q.submit([&](sycl::handler& cgh) {
+    cgh.parallel_for(range, [=](sycl::item<2> item) {
+      int row = item.get_id(1);
+      int col = item.get_id(0);
+      int out_idx = idx2(out_nrows, row, col);
+      int in_base_idx = idx2(in_nrows, row, col);
+      out2d[out_idx] = (stencil5[0] * in2d[in_base_idx + 0] +
+                        stencil5[1] * in2d[in_base_idx + 1] +
+                        stencil5[2] * in2d[in_base_idx + 2] +
+                        stencil5[3] * in2d[in_base_idx + 3] +
+                        stencil5[4] * in2d[in_base_idx + 4]) *
+                       scale;
+    });
+  });
+  return e;
+}
+
+/*
+ * Copy slice of first dimension of in array into contiguous
+ * buffer out. In has dimension nrows x ncols, buf has dimension (end - start) x
+ * ncols.
+ */
+auto buf_from_view(sycl::queue& q, int ncols, int buf_nrows, double* buf,
+                   int in_nrows, double* in, int start, int end)
+{
+  assert(buf_nrows >= end - start);
+  // Note: reverse index order because SYCL is row-major
+  auto range = sycl::range<2>(ncols, end - start);
+  auto e = q.submit([&](sycl::handler& cgh) {
+    cgh.parallel_for(range, [=](sycl::item<2> item) {
+      int row = item.get_id(1);
+      int col = item.get_id(0);
+      buf[idx2(buf_nrows, row, col)] = in[idx2(in_nrows, start + row, col)];
+    });
+  });
+  return e;
+}
+
+/*
+ * Copy contiguous buffer into first dimension of array as a slice. Out has
+ * dimension nrows x ncols, buf has dimension (end - start) * ncols.
+ */
+auto buf_to_view(sycl::queue& q, int ncols, int out_nrows, double* out,
+                 int buf_nrows, double* buf, int start, int end)
+{
+  assert(buf_nrows >= end - start);
+  // Note: reverse index order because SYCL is row-major
+  auto range = sycl::range<2>(ncols, end - start);
+  auto e = q.submit([&](sycl::handler& cgh) {
+    cgh.parallel_for(range, [=](sycl::item<2> item) {
+      int row = item.get_id(1);
+      int col = item.get_id(0);
+      out[idx2(out_nrows, start + row, col)] = buf[idx2(buf_nrows, row, col)];
+    });
+  });
+  return e;
+}
+
+void test_buf_view(sycl::queue& q, const int n)
+{
+  int n_bnd = 2;
+  int n_with_ghost = n + 2 * n_bnd;
+  double* data = sycl::malloc_host<double>(n_with_ghost * n, q);
+  double* buf = sycl::malloc_host<double>(n_bnd * n, q);
+  double* buf2 = sycl::malloc_host<double>(n_bnd * n, q);
+
+  for (int j = 0; j < n; j++) {
+    for (int i = 0; i < n_with_ghost; i++) {
+      data[idx2(n_with_ghost, i, j)] = (i - n_bnd) + j / 1000.0;
+    }
+    buf2[idx2(n_bnd, 0, j)] = 100.0 + j;
+    buf2[idx2(n_bnd, 1, j)] = 100.0 + j + 0.1;
+  }
+
+  for (int j = 0; j < n; j++) {
+    for (int i = 0; i < n; i++) {
+      printf("data[%d, %d] = %f\n", i, j, data[idx2(n_with_ghost, i, j)]);
+    }
+  }
+  for (int j = 0; j < n; j++) {
+    for (int i = 0; i < n_bnd; i++) {
+      printf("buf2[%d, %d] = %f\n", i, j, buf2[idx2(n_bnd, i, j)]);
+    }
+  }
+
+  buf_from_view(q, n, n_bnd, buf, n_with_ghost, data, 0, n_bnd).wait();
+  for (int j = 0; j < n; j++) {
+    for (int i = 0; i < n_bnd; i++) {
+      printf("buf[%d, %d] = %f\n", i, j, buf[idx2(n_bnd, i, j)]);
+    }
+  }
+
+  buf_to_view(q, n, n_with_ghost, data, n_bnd, buf2, n - n_bnd, n).wait();
+
+  for (int j = 0; j < n; j++) {
+    for (int i = 0; i < n; i++) {
+      printf("data[%d, %d] = %f\n", i, j, data[idx2(n_with_ghost, i, j)]);
+    }
+  }
+}
+
+/*
+ * Calculate the norm of the difference of two arrays, as sqrt of sum of
+ * squared distances.
+ */
+double diff_norm(sycl::queue& q, int size, double* d_a, double* d_b)
+{
+  double result = 0.0;
+  sycl::buffer<double> result_buf{&result, 1};
+  {
+    sycl::range<1> range(size);
+    auto e = q.submit([&](sycl::handler& cgh) {
+      auto reducer = sycl::reduction(result_buf, cgh, 0.0, std::plus<>{});
+      cgh.parallel_for(range, reducer, [=](sycl::id<1> idx, auto& r) {
+        double diff = d_a[idx] - d_b[idx];
+        r.combine(diff * diff);
+      });
+    });
+    e.wait();
+  }
+  return std::sqrt(result_buf.get_host_access()[0]);
+}
+
+sycl::queue get_rank_queue(int n_ranks, int rank)
+{
+  int n_devices, device_idx, ranks_per_device;
+
+  sycl::platform p{sycl::default_selector_v};
+  auto devices = p.get_devices();
+  n_devices = devices.size();
+
+  if (n_ranks > n_devices) {
+    if (n_ranks % n_devices != 0) {
+      printf(
+        "ERROR: Number of ranks (%d) not a multiple of number of GPUs (%d)\n",
+        n_ranks, n_devices);
+      exit(EXIT_FAILURE);
+    }
+    ranks_per_device = n_ranks / n_devices;
+    device_idx = rank / ranks_per_device;
+  } else {
+    ranks_per_device = 1;
+    device_idx = rank;
+  }
+
+  // printf("n_devices = %d\n", n_devices);
+  // printf("device_idx = %d\n", device_idx);
+
+  return sycl::queue{devices[device_idx], sycl::property::queue::in_order()};
+}
+
+// exchange in first dimension, staging into contiguous buffers on device
+void boundary_exchange_x(MPI_Comm comm, int world_size, int rank,
+                         sycl::queue& q, int n_global, int n_local, int n_bnd,
+                         double* d_z, bool stage_host = false)
+{
+  int n_local_with_ghost = n_local + 2 * n_bnd;
+  int buf_size = n_bnd * n_global;
+  static double* sbuf_l = nullptr;
+  static double* sbuf_r = nullptr;
+  static double* rbuf_l = nullptr;
+  static double* rbuf_r = nullptr;
+
+  if (sbuf_l == nullptr) {
+    sbuf_l = sycl::malloc_device<double>(buf_size, q);
+    sbuf_r = sycl::malloc_device<double>(buf_size, q);
+    rbuf_l = sycl::malloc_device<double>(buf_size, q);
+    rbuf_r = sycl::malloc_device<double>(buf_size, q);
+  }
+
+  static double* h_sbuf_l = nullptr;
+  static double* h_sbuf_r = nullptr;
+  static double* h_rbuf_l = nullptr;
+  static double* h_rbuf_r = nullptr;
+  if (stage_host && h_sbuf_l == nullptr) {
+    h_sbuf_l = sycl::malloc_host<double>(buf_size, q);
+    h_sbuf_r = sycl::malloc_host<double>(buf_size, q);
+    h_rbuf_l = sycl::malloc_host<double>(buf_size, q);
+    h_rbuf_r = sycl::malloc_host<double>(buf_size, q);
+  }
+
+  MPI_Request req_l[2];
+  MPI_Request req_r[2];
+
+  int rank_l = rank - 1;
+  int rank_r = rank + 1;
+
+  // start async copy of ghost points into send buffers
+  if (rank_l >= 0) {
+    // printf("rank_l = %d\n", rank_l); fflush(nullptr);
+    // sbuf_l = d_z.view(_all, _s(n_bnd, 2 * n_bnd));
+    auto e = buf_from_view(q, n_global, n_bnd, sbuf_l, n_local_with_ghost, d_z,
+                           n_bnd, 2 * n_bnd);
+    if (stage_host) {
+      q.copy(sbuf_l, h_sbuf_l, buf_size, e);
+      /*
+      for (int i = 0; i < n_bnd; i++) {
+        for (int j = 0; j < n_global; j++) {
+          int idx = idx2(n_global, j, i);
+          printf("sbuf_l[%d, %d] = %f\n", j, i, h_sbuf_l[idx]);
+          fflush(nullptr);
+        }
+      }
+      */
+    }
+  }
+  if (rank_r < world_size) {
+    // printf("rank_r = %d\n", rank_r); fflush(nullptr);
+    // sbuf_r = d_z.view(_all, _s(-2 * n_bnd, -n_bnd));
+    auto e = buf_from_view(q, n_global, n_bnd, sbuf_r, n_local_with_ghost, d_z,
+                           n_local, n_local + n_bnd);
+    if (stage_host) {
+      q.copy(sbuf_r, h_sbuf_r, buf_size, e);
+      /*
+      for (int i = 0; i < n_bnd; i++) {
+        for (int j = 0; j < n_global; j++) {
+          int idx = idx2(n_global, j, i);
+          printf("sbuf_r[%d, %d] = %f\n", j, i, h_sbuf_r[idx]);
+          fflush(nullptr);
+        }
+      }
+      */
+    }
+  }
+
+  // initiate async recv
+  if (rank_l >= 0) {
+    double* rbuf_l_data = nullptr;
+    if (stage_host) {
+      rbuf_l_data = h_rbuf_l;
+    } else {
+      rbuf_l_data = rbuf_l;
+    }
+    MPI_Irecv(rbuf_l_data, buf_size, MPI_DOUBLE, rank_l, 123, comm, &req_l[0]);
+  }
+
+  if (rank_r < world_size) {
+    double* rbuf_r_data = nullptr;
+    if (stage_host) {
+      rbuf_r_data = h_rbuf_r;
+    } else {
+      rbuf_r_data = rbuf_r;
+    }
+    MPI_Irecv(rbuf_r_data, buf_size, MPI_DOUBLE, rank_r, 456, comm, &req_r[0]);
+  }
+
+  // wait for send buffer fill
+  q.wait();
+
+  // initiate async sends
+  if (rank_l >= 0) {
+    double* sbuf_l_data = nullptr;
+    if (stage_host) {
+      sbuf_l_data = h_sbuf_l;
+    } else {
+      sbuf_l_data = sbuf_l;
+    }
+    MPI_Isend(sbuf_l_data, buf_size, MPI_DOUBLE, rank_l, 456, comm, &req_l[1]);
+  }
+
+  if (rank_r < world_size) {
+    double* sbuf_r_data = nullptr;
+    if (stage_host) {
+      sbuf_r_data = h_sbuf_r;
+    } else {
+      sbuf_r_data = sbuf_r;
+    }
+    MPI_Isend(sbuf_r_data, buf_size, MPI_DOUBLE, rank_r, 123, comm, &req_r[1]);
+  }
+
+  // wait for send/recv to complete, then copy data back into main data array
+  int mpi_rval;
+  if (rank_l >= 0) {
+    mpi_rval = MPI_Waitall(2, req_l, MPI_STATUSES_IGNORE);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("send_l error: %d\n", mpi_rval);
+    }
+    if (stage_host) {
+      /*
+      for (int i = 0; i < n_bnd; i++) {
+        for (int j = 0; j < n_global; j++) {
+          int idx = idx2(n_global, j, i);
+          printf("rbuf_l[%d, %d] = %f\n", j, i, h_rbuf_l[idx]);
+          fflush(nullptr);
+        }
+      }
+      */
+      q.copy(h_rbuf_l, rbuf_l, buf_size);
+    }
+    // d_z.view(_all, _s(0, n_bnd)) = rbuf_l;
+    buf_to_view(q, n_global, n_local_with_ghost, d_z, n_bnd, rbuf_l, 0, n_bnd);
+  }
+  if (rank_r < world_size) {
+    mpi_rval = MPI_Waitall(2, req_r, MPI_STATUSES_IGNORE);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("send_r error: %d\n", mpi_rval);
+    }
+    if (stage_host) {
+      /*
+      for (int i = 0; i < n_bnd; i++) {
+        for (int j = 0; j < n_global; j++) {
+          int idx = idx2(n_global, j, i);
+          printf("rbuf_r[%d, %d] = %f\n", j, i, h_rbuf_r[idx]);
+          fflush(nullptr);
+        }
+      }
+      */
+      q.copy(h_rbuf_r, rbuf_r, buf_size);
+    }
+    // d_z.view(_all, _s(-n_bnd, _)) = rbuf_r;
+    buf_to_view(q, n_global, n_local_with_ghost, d_z, n_bnd, rbuf_r,
+                n_local + n_bnd, n_local + 2 * n_bnd);
+  }
+
+  q.wait();
+}
+
+int main(int argc, char** argv)
+{
+  // sycl::queue q2{};
+  // test_buf_view(q2, 6);
+  // return EXIT_SUCCESS;
+
+  // Note: domain will be n_global x n_global plus ghost points in one dimension
+  int n_global = 8 * 1024;
+  bool stage_host = false;
+  int n_iter = 100;
+  int n_warmup = 5;
+
+  if (argc > 1) {
+    n_global = std::atoi(argv[1]) * 1024;
+  }
+  if (argc > 2) {
+    if (argv[2][0] == '1') {
+      stage_host = true;
+    }
+  }
+  if (argc > 3) {
+    n_iter = std::atoi(argv[3]);
+  }
+
+  int n_sten = 5;
+  int n_bnd = (n_sten - 1) / 2;
+  int world_size, world_rank, device_id;
+  uint32_t vendor_id;
+
+  MPI_Init(NULL, NULL);
+
+  MPI_Comm_size(MPI_COMM_WORLD, &world_size);
+  MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
+
+  if (n_global % world_size != 0) {
+    printf("%d nmpi (%d) must be divisor of domain size (%d), exiting\n",
+           world_rank, world_size, n_global);
+    exit(1);
+  }
+
+  const int n_local = n_global / world_size;
+  const int n_local_with_ghost = n_local + 2 * n_bnd;
+
+  sycl::queue q = get_rank_queue(world_size, world_rank);
+  auto dev = q.get_device();
+
+  if (dev.has(sycl::aspect::ext_intel_pci_address)) {
+    auto BDF = dev.get_info<sycl::ext::intel::info::device::pci_address>();
+    auto UUID = dev.get_info<sycl::ext::intel::info::device::uuid>();
+    uint32_t uuid_first = UUID[3] | (UUID[2] << 8) | (UUID[1] << 16) | (UUID[0] << 24);
+    std::cout << world_rank << " " << BDF << "(" << std::hex << uuid_first
+              << ")" << std::endl;
+  }
+
+  if (world_rank == 0) {
+    printf("n procs    = %d\n", world_size);
+    printf("rank       = %d\n", world_rank);
+    printf("n_global   = %d\n", n_global);
+    printf("n_local    = %d\n", n_local);
+    printf("n_iter     = %d\n", n_iter);
+    printf("n_warmup   = %d\n", n_warmup);
+    printf("stage_host = %d\n", stage_host);
+  }
+
+  int z_size = n_local_with_ghost * n_global;
+  int dzdx_size = n_local * n_global;
+
+  double* h_z = sycl::malloc_host<double>(z_size, q);
+  double* d_z = sycl::malloc_device<double>(z_size, q);
+
+  double* h_dzdx_numeric = sycl::malloc_host<double>(dzdx_size, q);
+  double* h_dzdx_actual = sycl::malloc_host<double>(dzdx_size, q);
+  double* d_dzdx_numeric = sycl::malloc_device<double>(dzdx_size, q);
+
+  double lx = 8;
+  double dx = lx / n_global;
+  double lx_local = lx / world_size;
+  double scale = n_global / lx;
+  auto fn = [](double x, double y) { return x * x * x + y * y; };
+  auto fn_dzdx = [](double x, double y) { return 3 * x * x; };
+
+  struct timespec start, end;
+  double iter_time = 0.0;
+  double total_time = 0.0;
+
+  double x_start = world_rank * lx_local;
+  for (int j = 0; j < n_global; j++) {
+    double ytmp = j * dx;
+    for (int i = 0; i < n_local; i++) {
+      double xtmp = x_start + i * dx;
+      h_z[idx2(n_local_with_ghost, i + n_bnd, j)] = fn(xtmp, ytmp);
+      h_dzdx_actual[idx2(n_local, i, j)] = fn_dzdx(xtmp, ytmp);
+    }
+  }
+
+  // fill boundary points on ends
+  if (world_rank == 0) {
+    for (int j = 0; j < n_global; j++) {
+      double ytmp = j * dx;
+      for (int i = 0; i < n_bnd; i++) {
+        double xtmp = (i - n_bnd) * dx;
+        h_z[idx2(n_local_with_ghost, i, j)] = fn(xtmp, ytmp);
+      }
+    }
+  }
+  if (world_rank == world_size - 1) {
+    for (int j = 0; j < n_global; j++) {
+      double ytmp = j * dx;
+      for (int i = 0; i < n_bnd; i++) {
+        double xtmp = lx + i * dx;
+        h_z[idx2(n_local_with_ghost, n_bnd + n_local + i, j)] = fn(xtmp, ytmp);
+      }
+    }
+  }
+
+  /*
+  for (int i = 0; i < 5; i++) {
+    int idx = idx2(n_global, 1, i);
+    printf("%d row1-l %f\n", world_rank, h_z[idx]);
+  }
+  for (int i = 0; i < 5; i++) {
+    int idx = idx2(n_global, 1, n_local_with_ghost - 1 - i);
+    printf("%d row1-r %f\n", world_rank, h_z[idx]);
+  }
+  */
+
+  q.copy(h_z, d_z, z_size);
+
+  for (int i = 0; i < n_warmup + n_iter; i++) {
+    clock_gettime(CLOCK_MONOTONIC, &start);
+    boundary_exchange_x(MPI_COMM_WORLD, world_size, world_rank, q, n_global,
+                        n_local, n_bnd, d_z, stage_host);
+    clock_gettime(CLOCK_MONOTONIC, &end);
+    iter_time =
+      ((end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) * 1.0e-9);
+
+    if (i >= n_warmup) {
+      total_time += iter_time;
+    }
+
+    // do some calculation, to try to more closely simulate what happens in GENE
+    auto e = stencil2d_1d_5(q, n_local, n_global, d_dzdx_numeric, d_z, scale);
+    e.wait();
+  }
+  printf("%d/%d exchange time %0.8f ms\n", world_rank, world_size,
+         total_time / n_iter * 1000);
+
+  q.copy(d_dzdx_numeric, h_dzdx_numeric, dzdx_size).wait();
+
+  /*
+  for (int i = 0; i < 5; i++) {
+    int idx = idx2(n_global, 8, i);
+    printf("%d la %f\n%d ln %f\n", world_rank, h_dzdx_actual[idx], world_rank,
+           h_dzdx_numeric[idx]);
+  }
+  for (int i = 0; i < 5; i++) {
+    int idx = idx2(n_global, 8, n_local - 1 - i);
+    printf("%d ra %f\n%d rn %f\n", world_rank, h_dzdx_actual[idx], world_rank,
+           h_dzdx_numeric[idx]);
+  }
+  */
+
+  double err_norm = diff_norm(q, dzdx_size, h_dzdx_numeric, h_dzdx_actual);
+
+  printf("%d/%d [%d:0x%08x] err_norm = %.8f\n", world_rank, world_size,
+         device_id, vendor_id, err_norm);
+
+  MPI_Finalize();
+
+  return EXIT_SUCCESS;
+}

mpi_stencil2d_sycl_oo.cc

@@ -0,0 +1,705 @@
+/*
+ * Test GPU aware MPI on different platforms using a distributed
+ * 1d stencil on a 2d array. The exchange in second (non-contiguous)
+ * direction forces use of staging buffers, which replicates what
+ * is needed for all but the innermost dimension exchanges in the
+ * GENE fusion code.
+ *
+ * Takes optional command line arg for size of each dimension of the domain
+ * n_global, in 1024 increments. Default is 8 * 1024 (so 256K plus ghost points
+ * in size for doulbles per array), which should fit on any system but may not
+ * be enough to tax larger HPC GPUs and MPI impelmentations.
+ *
+ * There will be four exchange buffers of size 2 * n_global, i.e. 128K each
+ * by default.
+ *
+ * Modified version that uses minimal owning (array2d) and non-owning (span2d)
+ * classes to make indexing handling less error prone, without using all of
+ * gtensor. Note that the owning class is not trivially copyable and not device
+ * copyable, because it must have a non-trivial destructor.
+ *
+ * TODO: Since no temporaries are used, perhaps a helper that allocates and
+ * returns a span is a simpler option to create this minimal example?
+ */
+
+#include <cassert>
+#include <cmath>
+#include <memory>
+#include <mpi.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+#include <type_traits>
+
+#include "sycl/sycl.hpp"
+
+// #define DEBUG
+
+#ifdef DEBUG
+#define dprintf(...) fprintf(stderr, __VA_ARGS__)
+#else
+#define dprintf(...)                                                           \
+  do {                                                                         \
+  } while (0)
+#endif
+
+constexpr std::size_t idx2(int n, int row, int col)
+{
+  return row + col * n;
+}
+
+template <typename T, sycl::usm::alloc Alloc>
+class span2d
+{
+public:
+  using value_type = T;
+  using pointer = value_type*;
+  using const_pointer = const value_type*;
+  using reference = value_type&;
+  using const_reference = const value_type&;
+  using size_type = std::size_t;
+
+  span2d(T* data, const int nrows, const int ncols)
+    : data_(data), nrows_(nrows), ncols_(ncols)
+  {}
+
+  // use default copy and move ctor. Ideall move ctor would better invalidate
+  // the moved from object, but this is supposed to be a small example...
+  span2d(const span2d& other) = default;
+  span2d& operator=(const span2d& other) = default;
+  span2d(span2d&& other) = default;
+  span2d& operator=(span2d&& other) = default;
+
+  // Note: shallow const
+  reference operator()(int row, int col) const
+  {
+    assert(row < nrows_);
+    assert(col < ncols_);
+    return data_[idx2(nrows_, row, col)];
+  }
+
+  // Note: shallow const
+  reference operator[](size_type i) const
+  {
+    assert(i < (nrows_ * ncols_));
+    return data_[i];
+  }
+
+  int nrows() const { return nrows_; }
+  int ncols() const { return ncols_; }
+  size_type size() const { return nrows_ * ncols_; }
+
+  span2d to_span() { return *this; }
+
+  // Note: shallow const
+  pointer data() const { return data_; }
+
+private:
+  const sycl::usm::alloc alloc_ = Alloc;
+  T* data_;
+  const int nrows_;
+  const int ncols_;
+};
+
+template <typename T>
+auto empty_host(sycl::queue& q, int nrows, int ncols)
+{
+  T* data = sycl::malloc(nrows * ncols, q, sycl::usm::alloc::host);
+  return span2d<T, sycl::usm::alloc::host>(data, nrows, ncols);
+}
+
+template <typename T>
+auto empty_device(sycl::queue& q, int nrows, int ncols)
+{
+  T* data = sycl::malloc(nrows * ncols, q, sycl::usm::alloc::device);
+  return span2d<T, sycl::usm::alloc::device>(data, nrows, ncols);
+}
+
+template <typename T, sycl::usm::alloc Alloc>
+class array2d : public span2d<T, Alloc>
+{
+public:
+  using base_type = span2d<T, Alloc>;
+  using value_type = T;
+  using pointer = value_type*;
+  using const_pointer = const value_type*;
+  using reference = value_type&;
+  using const_reference = const value_type&;
+  using size_type = std::size_t;
+
+  array2d(sycl::queue& q, const int nrows, const int ncols)
+    : base_type(sycl::malloc<value_type>(nrows * ncols, q, Alloc), nrows,
+                ncols),
+      q_(q)
+  {}
+
+  // Results in a double free, why?
+  // ~array2d() { sycl::free(this->data(), q_); }
+
+  // skip these to keep the example simple, pass by reference everywhere
+  array2d(const array2d& other) = delete;
+  array2d& operator=(const array2d& other) = delete;
+  array2d(array2d&& other) = delete;
+  array2d& operator=(array2d&& other) = delete;
+
+  base_type to_span()
+  {
+    return base_type(this->data(), this->nrows(), this->ncols());
+  }
+
+private:
+  sycl::queue& q_;
+};
+
+template <typename SrcArray, typename DestArray>
+auto copy(sycl::queue& q, SrcArray& src, DestArray& dest)
+{
+  static_assert(std::is_same_v<typename SrcArray::value_type,
+                               typename DestArray::value_type>,
+                "value types must match");
+  assert(src.size() == dest.size());
+  return q.copy(src.data(), dest.data(), dest.size());
+}
+
+template <typename Array>
+auto copy_dest_slice(sycl::queue& q, Array& src, Array& dest, int dim,
+                     int start, int end)
+{
+  dprintf("copy dest_slice %d %d %d\n", dim, start, end);
+  auto s_src = src.to_span();
+  auto s_dest = dest.to_span();
+  assert(dim == 0 || dim == 1);
+  if (dim == 0) {
+    assert(src.ncols() == dest.ncols());
+    if (start < 0) {
+      start += dest.nrows();
+    }
+    if (end < 0) {
+      end += dest.nrows();
+    } else if (end == 0 && start > end) {
+      end = dest.nrows();
+    }
+    assert(start < end);
+    auto range = sycl::range<2>(dest.ncols(), end - start);
+    dprintf("d_z < buf range %d - %d (%d, %d)\n", start, end, dest.ncols(),
+            end - start);
+    auto e = q.submit([&](sycl::handler& cgh) {
+      cgh.parallel_for(range, [=](sycl::item<2> item) {
+        int row = item.get_id(1);
+        int col = item.get_id(0);
+        s_dest(start + row, col) = s_src(row, col);
+      });
+    });
+    return e;
+  } else {
+    assert(src.nrows() == dest.nrows());
+    if (start < 0) {
+      start += dest.ncols();
+    }
+    if (end < 0) {
+      end += dest.ncols();
+    } else if (end == 0 && start > end) {
+      end = dest.ncols();
+    }
+    auto range = sycl::range<2>(end - start, dest.nrows());
+    auto e = q.submit([&](sycl::handler& cgh) {
+      cgh.parallel_for(range, [=](sycl::item<2> item) {
+        int row = item.get_id(1);
+        int col = item.get_id(0);
+        s_dest(row, start + col) = s_src(row, col);
+      });
+    });
+    return e;
+  }
+}
+
+template <typename Array>
+auto copy_src_slice(sycl::queue& q, Array& src, Array& dest, int dim, int start,
+                    int end)
+{
+  dprintf("copy  src_slice %d %d %d (%d, %d) -> (%d, %d)\n", dim, start, end,
+          src.nrows(), src.ncols(), dest.nrows(), dest.ncols());
+  assert(dim == 0 || dim == 1);
+  auto s_src = src.to_span();
+  auto s_dest = dest.to_span();
+  if (dim == 0) {
+    assert(src.ncols() == dest.ncols());
+    if (start < 0) {
+      start += src.nrows();
+    }
+    if (end < 0) {
+      end += src.nrows();
+    } else if (end == 0 && start > end) {
+      end = src.nrows();
+    }
+    auto range = sycl::range<2>(dest.ncols(), end - start);
+    dprintf("buf < d_z range %d - %d (%d, %d)\n", start, end, dest.ncols(),
+            end - start);
+    auto e = q.submit([&](sycl::handler& cgh) {
+      cgh.parallel_for(range, [=](sycl::item<2> item) {
+        int row = item.get_id(1);
+        int col = item.get_id(0);
+        s_dest(row, col) = s_src(start + row, col);
+      });
+    });
+    return e;
+  } else {
+    assert(src.nrows() == dest.nrows());
+    if (start < 0) {
+      start += src.ncols();
+    }
+    if (end < 0) {
+      end += src.ncols();
+    } else if (end == 0 && start > end) {
+      end = src.ncols();
+    }
+    auto range = sycl::range<2>(end - start, dest.nrows());
+    auto e = q.submit([&](sycl::handler& cgh) {
+      cgh.parallel_for(range, [=](sycl::item<2> item) {
+        int row = item.get_id(1);
+        int col = item.get_id(0);
+        s_dest(row, col) = s_src(row, start + col);
+      });
+    });
+    return e;
+  }
+}
+
+inline void check(const char* fname, int line, int mpi_rval)
+{
+  if (mpi_rval != MPI_SUCCESS) {
+    printf("%s:%d error %d\n", fname, line, mpi_rval);
+    exit(2);
+  }
+}
+
+#define CHECK(x) check(__FILE__, __LINE__, (x))
+
+static constexpr double stencil5[] = {1.0 / 12.0, -2.0 / 3.0, 0.0, 2.0 / 3.0,
+                                      -1.0 / 12.0};
+
+/*
+ * Calculate 1d stencil of second dimension of 2d array on GPU. Out array must
+ * be contiguous column major nrows x ncols array, while in array must be
+ * (nrows)x(ncols+4) to accomodate 2 ghost points in each direction for the
+ * second dimension.
+ *
+ * Returns sycl event, async with respect to host.
+ */
+template <typename Array>
+auto stencil2d_1d_5(sycl::queue& q, Array& out2d, Array& in2d, double scale)
+{
+  // Note: swap index order; SYCL is row-major oriented, and this example
+  // is col-major
+  auto range = sycl::range<2>(out2d.ncols(), out2d.nrows());
+  auto s_in2d = in2d.to_span();
+  auto s_out2d = out2d.to_span();
+  auto e = q.submit([&](sycl::handler& cgh) {
+    cgh.parallel_for(range, [=](sycl::item<2> item) {
+      int row = item.get_id(1);
+      int col = item.get_id(0);
+      s_out2d(row, col) = (stencil5[0] * s_in2d(row + 0, col) +
+                           stencil5[1] * s_in2d(row + 1, col) +
+                           stencil5[2] * s_in2d(row + 2, col) +
+                           stencil5[3] * s_in2d(row + 3, col) +
+                           stencil5[4] * s_in2d(row + 4, col)) *
+                          scale;
+    });
+  });
+  return e;
+}
+
+/*
+ * Calculate the norm of the difference of two arrays, as sqrt of sum of
+ * squared distances.
+ */
+double diff_norm(sycl::queue& q, std::size_t size, double* d_a, double* d_b)
+{
+  double result = 0.0;
+  sycl::buffer<double> result_buf{&result, 1};
+  {
+    sycl::range<1> range(size);
+    auto e = q.submit([&](sycl::handler& cgh) {
+      auto reducer = sycl::reduction(result_buf, cgh, 0.0, std::plus<>{});
+      cgh.parallel_for(range, reducer, [=](sycl::id<1> idx, auto& r) {
+        double diff = d_a[idx] - d_b[idx];
+        r.combine(diff * diff);
+      });
+    });
+    e.wait();
+  }
+  return std::sqrt(result_buf.get_host_access()[0]);
+}
+
+sycl::queue get_rank_queue(int n_ranks, int rank)
+{
+  int n_devices, device_idx, ranks_per_device;
+
+  sycl::context ctx{};
+  auto devices = ctx.get_devices();
+  n_devices = devices.size();
+
+  if (n_ranks > n_devices) {
+    if (n_ranks % n_devices != 0) {
+      printf(
+        "ERROR: Number of ranks (%d) not a multiple of number of GPUs (%d)\n",
+        n_ranks, n_devices);
+      exit(EXIT_FAILURE);
+    }
+    ranks_per_device = n_ranks / n_devices;
+    device_idx = rank / ranks_per_device;
+  } else {
+    ranks_per_device = 1;
+    device_idx = rank;
+  }
+
+  dprintf("%d: n_devices = %d\n", rank, n_devices);
+  dprintf("%d: device_idx = %d\n", rank, device_idx);
+
+  return sycl::queue{devices[device_idx], sycl::property::queue::in_order()};
+}
+
+// exchange in first dimension, staging into contiguous buffers on device
+template <typename T, sycl::usm::alloc Alloc>
+void boundary_exchange_x(MPI_Comm comm, int world_size, int rank,
+                         sycl::queue& q, int n_bnd, array2d<T, Alloc>& d_z,
+                         bool stage_host = false)
+{
+  static array2d<double, sycl::usm::alloc::device> sbuf_l{q, n_bnd,
+                                                          d_z.ncols()};
+  static array2d<double, sycl::usm::alloc::device> sbuf_r{q, n_bnd,
+                                                          d_z.ncols()};
+  static array2d<double, sycl::usm::alloc::device> rbuf_l{q, n_bnd,
+                                                          d_z.ncols()};
+  static array2d<double, sycl::usm::alloc::device> rbuf_r{q, n_bnd,
+                                                          d_z.ncols()};
+
+  static array2d<double, sycl::usm::alloc::host> h_sbuf_l{q, n_bnd,
+                                                          d_z.ncols()};
+  static array2d<double, sycl::usm::alloc::host> h_sbuf_r{q, n_bnd,
+                                                          d_z.ncols()};
+  static array2d<double, sycl::usm::alloc::host> h_rbuf_l{q, n_bnd,
+                                                          d_z.ncols()};
+  static array2d<double, sycl::usm::alloc::host> h_rbuf_r{q, n_bnd,
+                                                          d_z.ncols()};
+
+  int buf_size = sbuf_l.size();
+
+  MPI_Request req_l[2];
+  MPI_Request req_r[2];
+
+  int rank_l = rank - 1;
+  int rank_r = rank + 1;
+
+  // start async copy of ghost points into send buffers
+  if (rank_l >= 0) {
+    dprintf("%d: rank_l = %d\n", rank, rank_l);
+    fflush(nullptr);
+    // sbuf_l = d_z.view(_all, _s(n_bnd, 2 * n_bnd));
+    auto e = copy_src_slice(q, d_z, sbuf_l, 0, n_bnd, 2 * n_bnd);
+    if (stage_host) {
+      e.wait();
+      copy(q, sbuf_l, h_sbuf_l).wait();
+      for (int i = 0; i < h_sbuf_l.ncols(); i++) {
+        for (int j = 0; j < h_sbuf_l.nrows(); j++) {
+          dprintf("%d: sbuf_l[%d, %d] = %f\n", rank, j, i, h_sbuf_l(j, i));
+          fflush(nullptr);
+        }
+      }
+    }
+  }
+  if (rank_r < world_size) {
+    dprintf("%d: rank_r = %d\n", rank, rank_r);
+    fflush(nullptr);
+    // sbuf_r = d_z.view(_all, _s(-2 * n_bnd, -n_bnd));
+    auto e = copy_src_slice(q, d_z, sbuf_r, 0, -2 * n_bnd, -n_bnd);
+    if (stage_host) {
+      e.wait();
+      copy(q, sbuf_r, h_sbuf_r).wait();
+      for (int i = 0; i < h_sbuf_r.ncols(); i++) {
+        for (int j = 0; j < h_sbuf_r.nrows(); j++) {
+          dprintf("%d: sbuf_r[%d, %d] = %f\n", rank, j, i, h_sbuf_r(j, i));
+          fflush(nullptr);
+        }
+      }
+    }
+  }
+
+  // initiate async recv
+  if (rank_l >= 0) {
+    double* rbuf_l_data = nullptr;
+    if (stage_host) {
+      rbuf_l_data = h_rbuf_l.data();
+    } else {
+      rbuf_l_data = rbuf_l.data();
+    }
+    MPI_Irecv(rbuf_l_data, buf_size, MPI_DOUBLE, rank_l, 123, comm, &req_l[0]);
+  }
+
+  if (rank_r < world_size) {
+    double* rbuf_r_data = nullptr;
+    if (stage_host) {
+      rbuf_r_data = h_rbuf_r.data();
+    } else {
+      rbuf_r_data = rbuf_r.data();
+    }
+    MPI_Irecv(rbuf_r_data, buf_size, MPI_DOUBLE, rank_r, 456, comm, &req_r[0]);
+  }
+
+  // wait for send buffer fill
+  q.wait();
+
+  // initiate async sends
+  if (rank_l >= 0) {
+    double* sbuf_l_data = nullptr;
+    if (stage_host) {
+      sbuf_l_data = h_sbuf_l.data();
+    } else {
+      sbuf_l_data = sbuf_l.data();
+    }
+    MPI_Isend(sbuf_l_data, buf_size, MPI_DOUBLE, rank_l, 456, comm, &req_l[1]);
+  }
+
+  if (rank_r < world_size) {
+    double* sbuf_r_data = nullptr;
+    if (stage_host) {
+      sbuf_r_data = h_sbuf_r.data();
+    } else {
+      sbuf_r_data = sbuf_r.data();
+    }
+    MPI_Isend(sbuf_r_data, buf_size, MPI_DOUBLE, rank_r, 123, comm, &req_r[1]);
+  }
+
+  // wait for send/recv to complete, then copy data back into main data array
+  int mpi_rval;
+  if (rank_l >= 0) {
+    mpi_rval = MPI_Waitall(2, req_l, MPI_STATUSES_IGNORE);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("%d: send_l error: %d\n", rank, mpi_rval);
+    }
+    if (stage_host) {
+#ifdef DEBUG
+      for (int i = 0; i < h_rbuf_l.ncols(); i++) {
+        for (int j = 0; j < h_rbuf_l.nrows(); j++) {
+          dprintf("%d: rbuf_l[%d, %d] = %f\n", rank, j, i, h_rbuf_l(j, i));
+          fflush(nullptr);
+        }
+      }
+#endif
+      copy(q, h_rbuf_l, rbuf_l).wait();
+    }
+    // d_z.view(_all, _s(0, n_bnd)) = rbuf_l;
+    copy_dest_slice(q, rbuf_l, d_z, 0, 0, n_bnd);
+  }
+  if (rank_r < world_size) {
+    mpi_rval = MPI_Waitall(2, req_r, MPI_STATUSES_IGNORE);
+    if (mpi_rval != MPI_SUCCESS) {
+      printf("%d: send_r error: %d\n", rank, mpi_rval);
+    }
+    if (stage_host) {
+#ifdef DEBUG
+      for (int i = 0; i < h_rbuf_r.ncols(); i++) {
+        for (int j = 0; j < h_rbuf_r.nrows(); j++) {
+          dprintf("%d: rbuf_r[%d, %d] = %f\n", rank, j, i, h_rbuf_r(j, i));
+          fflush(nullptr);
+        }
+      }
+#endif
+      copy(q, h_rbuf_r, rbuf_r).wait();
+    }
+    // d_z.view(_all, _s(-n_bnd, _)) = rbuf_r;
+    copy_dest_slice(q, rbuf_r, d_z, 0, -n_bnd, 0);
+  }
+
+  q.wait();
+}
+
+int main(int argc, char** argv)
+{
+  using T = double;
+
+  static_assert(
+    std::is_trivially_copyable_v<span2d<T, sycl::usm::alloc::device>>,
+    "span2d device not trivial");
+  static_assert(std::is_trivially_copyable_v<span2d<T, sycl::usm::alloc::host>>,
+                "span2d host not trivial");
+
+  // Note: domain will be n_global x n_global plus ghost points in one dimension
+  int n_global = 8 * 1024;
+  bool stage_host = false;
+  int n_iter = 100;
+  int n_warmup = 5;
+
+  if (argc > 1) {
+    n_global = std::atoi(argv[1]) * 1024;
+  }
+  if (argc > 2) {
+    if (argv[2][0] == '1') {
+      stage_host = true;
+    }
+  }
+  if (argc > 3) {
+    n_iter = std::atoi(argv[3]);
+  }
+
+#ifdef DEBUG
+  n_global /= 1024;
+  n_iter = 1;
+  n_warmup = 0;
+#endif
+
+  int n_sten = 5;
+  int n_bnd = (n_sten - 1) / 2;
+  int world_size, world_rank, device_id;
+  uint32_t vendor_id;
+
+  MPI_Init(NULL, NULL);
+
+  MPI_Comm_size(MPI_COMM_WORLD, &world_size);
+  MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
+
+  if (n_global % world_size != 0) {
+    printf("%d: nmpi (%d) must be divisor of domain size (%d), exiting\n",
+           world_rank, world_size, n_global);
+    exit(1);
+  }
+
+  const int n_local = n_global / world_size;
+  const int n_local_with_ghost = n_local + 2 * n_bnd;
+
+  sycl::queue q = get_rank_queue(world_size, world_rank);
+
+  vendor_id = q.get_device().get_info<sycl::info::device::vendor_id>();
+
+  if (world_rank == 0) {
+    printf("n procs    = %d\n", world_size);
+    printf("rank       = %d\n", world_rank);
+    printf("n_global   = %d\n", n_global);
+    printf("n_local    = %d\n", n_local);
+    printf("n_iter     = %d\n", n_iter);
+    printf("n_warmup   = %d\n", n_warmup);
+    printf("stage_host = %d\n", stage_host);
+  }
+
+  int z_size = n_local_with_ghost * n_global;
+  int dzdx_size = n_local * n_global;
+
+  array2d<T, sycl::usm::alloc::host> h_z{q, n_local_with_ghost, n_global};
+  array2d<T, sycl::usm::alloc::device> d_z{q, n_local_with_ghost, n_global};
+
+  array2d<T, sycl::usm::alloc::host> h_dzdx_actual{q, n_local, n_global};
+  array2d<T, sycl::usm::alloc::host> h_dzdx_numeric{q, n_local, n_global};
+  array2d<T, sycl::usm::alloc::device> d_dzdx_actual{q, n_local, n_global};
+  array2d<T, sycl::usm::alloc::device> d_dzdx_numeric{q, n_local, n_global};
+
+  double lx = 8;
+  double dx = lx / n_global;
+  double lx_local = lx / world_size;
+  double scale = n_global / lx;
+  auto fn = [](double x, double y) { return x * x * x + y * y; };
+  auto fn_dzdx = [](double x, double y) { return 3 * x * x; };
+
+  struct timespec start, end;
+  double iter_time = 0.0;
+  double total_time = 0.0;
+
+  double x_start = world_rank * lx_local;
+  for (int j = 0; j < h_z.ncols(); j++) {
+    double ytmp = j * dx;
+    for (int i = 0; i < n_local; i++) {
+      double xtmp = x_start + i * dx;
+      h_z(i + n_bnd, j) = fn(xtmp, ytmp);
+      h_dzdx_actual(i, j) = fn_dzdx(xtmp, ytmp);
+    }
+  }
+
+  // fill boundary points on ends
+  if (world_rank == 0) {
+    for (int j = 0; j < h_z.ncols(); j++) {
+      double ytmp = j * dx;
+      for (int i = 0; i < n_bnd; i++) {
+        double xtmp = (i - n_bnd) * dx;
+        h_z(i, j) = fn(xtmp, ytmp);
+      }
+    }
+  }
+  if (world_rank == world_size - 1) {
+    for (int j = 0; j < h_z.ncols(); j++) {
+      double ytmp = j * dx;
+      for (int i = 0; i < n_bnd; i++) {
+        double xtmp = lx + i * dx;
+        h_z(n_bnd + n_local + i, j) = fn(xtmp, ytmp);
+      }
+    }
+  }
+
+#ifdef DEBUG
+  for (int r = 0; r < world_size; r++) {
+    if (r != world_rank) {
+      continue;
+    }
+
+    for (int i = n_bnd; i < 2 * n_bnd; i++) {
+      dprintf("%d: [%d, :]", world_rank, i);
+      for (int j = 0; j < std::min(20, h_z.ncols()); j++) {
+        dprintf(" %f", h_z(i, j));
+      }
+      dprintf("\n");
+    }
+    for (int i = h_z.nrows() - 2 * n_bnd; i < h_z.nrows() - n_bnd; i++) {
+      dprintf("%d: [%d, :]", world_rank, i);
+      for (int j = 0; j < std::min(20, h_z.ncols()); j++) {
+        dprintf(" %f", h_z(i, j));
+      }
+      dprintf("\n");
+    }
+
+    MPI_Barrier(MPI_COMM_WORLD);
+  }
+#endif
+
+  copy(q, h_z, d_z);
+
+  for (int i = 0; i < n_warmup + n_iter; i++) {
+    clock_gettime(CLOCK_MONOTONIC, &start);
+    boundary_exchange_x(MPI_COMM_WORLD, world_size, world_rank, q, n_bnd, d_z,
+                        stage_host);
+    clock_gettime(CLOCK_MONOTONIC, &end);
+    iter_time =
+      ((end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) * 1.0e-9);
+
+    if (i >= n_warmup) {
+      total_time += iter_time;
+    }
+
+    // do some calculation, to try to more closely simulate what happens in GENE
+    auto e = stencil2d_1d_5(q, d_dzdx_numeric, d_z, scale);
+    e.wait();
+  }
+  printf("%d: exchange time %0.8f ms\n", world_rank,
+         total_time / n_iter * 1000);
+
+  copy(q, d_dzdx_numeric, h_dzdx_numeric).wait();
+
+  /*
+  for (int i = 0; i < 5; i++) {
+    int idx = idx2(n_global, 8, i);
+    printf("%d la %f\n%d ln %f\n", world_rank, h_dzdx_actual[idx], world_rank,
+           h_dzdx_numeric[idx]);
+  }
+  for (int i = 0; i < 5; i++) {
+    int idx = idx2(n_global, 8, n_local - 1 - i);
+    printf("%d ra %f\n%d rn %f\n", world_rank, h_dzdx_actual[idx], world_rank,
+           h_dzdx_numeric[idx]);
+  }
+  */
+
+  double err_norm = diff_norm(q, h_dzdx_numeric.size(), h_dzdx_numeric.data(),
+                              h_dzdx_actual.data());
+
+  printf("%d: [0x%08x] err_norm = %.8f\n", world_rank, vendor_id, err_norm);
+
+  MPI_Finalize();
+
+  return EXIT_SUCCESS;
+}

mpi_stencil_gt.cc

@@ -138,6 +138,12 @@ int main(int argc, char **argv) {
   MPI_Comm_size(MPI_COMM_WORLD, &world_size);
   MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
 
+  if (n_global % world_size != 0) {
+    printf("%d nmpi (%d) must be divisor of domain size (%d), exiting\n",
+           world_rank, world_size, n_global);
+    exit(1);
+  }
+
   const int n_local = n_global / world_size;
   const int n_local_with_ghost = n_local + 2 * n_bnd;
 
@@ -176,7 +182,7 @@ int main(int argc, char **argv) {
   }
 
   // fill boundary points on ends
-  if (world_rank == 1) {
+  if (world_rank == 0) {
     for (int i = 0; i < n_bnd; i++) {
       double xtmp = (i - n_bnd) * dx;
       h_y(i) = fn_x_cubed(xtmp);
@@ -190,30 +196,26 @@ int main(int argc, char **argv) {
   }
 
   gt::copy(h_y, d_y);
-  // gt::synchronize();
 
   clock_gettime(CLOCK_MONOTONIC, &start);
   boundary_exchange(MPI_COMM_WORLD, world_size, world_rank, d_y, n_bnd);
-  // gt::synchronize();
   clock_gettime(CLOCK_MONOTONIC, &end);
   seconds = ((end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) * 1.0e-9);
-  printf("%d/%d exchange time %0.4f\n", world_rank, world_size, seconds);
+  printf("%d/%d exchange time %0.8f\n", world_rank, world_size, seconds);
 
   d_dydx_numeric = stencil1d_5(d_y, stencil5) * scale;
-  // gt::synchronize();
 
   gt::copy(d_dydx_numeric, h_dydx_numeric);
-  // gt::synchronize();
 
   /*
   for (int i = 0; i < 5; i++) {
-    printf("{0} l {1}\n{0} l {2}\n", world_rank, h_dydx_actual(i),
+    printf("%d la %f\n%d ln %f\n", world_rank, h_dydx_actual(i),
-               h_dydx_numeric(i));
+           world_rank, h_dydx_numeric(i));
   }
   for (int i = 0; i < 5; i++) {
     int idx = n_local - 1 - i;
-    printf("{0} r {1}\n{0} r {2}\n", world_rank, h_dydx_actual(idx),
+    printf("%d ra %f\n%d rn %f\n", world_rank, h_dydx_actual(idx),
-               h_dydx_numeric(idx));
+           world_rank, h_dydx_numeric(idx));
   }
   */