From 74a4f4394690dc872afb1f93e3d49c97a35f46f0 Mon Sep 17 00:00:00 2001
From: Peter Boyle <paboyle@ed.ac.uk>
Date: Tue, 28 Jan 2025 15:22:46 +0000
Subject: [PATCH] Optional host buffer bounce for no CUDA aware MPI

---
 Grid/communicator/SharedMemoryMPI.cc          | 28 +++++++++++++++
 Grid/stencil/Stencil.h                        | 28 ++++++++++++++-
 configure.ac                                  | 12 +++++++
 systems/Aurora/benchmarks/bench1.pbs          | 32 ++++++++---------
 systems/Aurora/benchmarks/gpu_tile_compact.sh | 34 -------------------
 systems/Aurora/config-command                 |  2 +-
 systems/Aurora/sourceme.sh                    |  1 +
 7 files changed, 84 insertions(+), 53 deletions(-)
 delete mode 100755 systems/Aurora/benchmarks/gpu_tile_compact.sh

diff --git a/Grid/communicator/SharedMemoryMPI.cc b/Grid/communicator/SharedMemoryMPI.cc
index ec6a5003..2642c0bd 100644
--- a/Grid/communicator/SharedMemoryMPI.cc
+++ b/Grid/communicator/SharedMemoryMPI.cc
@@ -42,6 +42,11 @@ Author: Christoph Lehner <christoph@lhnr.de>
 #ifdef ACCELERATOR_AWARE_MPI
 #define GRID_SYCL_LEVEL_ZERO_IPC
 #define SHM_SOCKETS
+#else
+#undef NUMA_PLACE_HOSTBUF
+#ifdef NUMA_PLACE_HOSTBUF
+#include <numaif.h>
+#endif 
 #endif 
 #include <syscall.h>
 #endif
@@ -537,7 +542,30 @@ void GlobalSharedMemory::SharedMemoryAllocate(uint64_t bytes, int flags)
   // Each MPI rank should allocate our own buffer
   ///////////////////////////////////////////////////////////////////////////////////////////////////////////
 #ifndef ACCELERATOR_AWARE_MPI
+  printf("Host buffer allocate for GPU non-aware MPI\n");
   HostCommBuf= malloc(bytes);
+#ifdef NUMA_PLACE_HOSTBUF
+  int numa;
+  char *numa_name=(char *)getenv("MPI_BUF_NUMA");
+  if(numa_name) {
+    page_size = sysconf(_SC_PAGESIZE);
+    numa = atoi(numa_name);
+    unsigned long page_count = bytes/page_size;
+    std::vector<void *> pages(pcount);
+    std::vector<int>    nodes(pcount,numa);
+    std::vector<int>    status(pcount,-1);
+    for(unsigned long p=0;p<page_count;p++){
+      pages[p] = HostCommBuf + p*page_size;
+    }
+    int ret = move_pages(0,
+			 page_count,
+			 &pages[0],
+			 &nodes[0],
+			 &status[0],
+			 MPOL_MF_MOVE);
+    printf("Host buffer move to numa domain %d : move_pages returned %d\n",numa,ret);
+    if (ret) perror(" move_pages failed for reason:");
+#endif  
 #endif  
   ShmCommBuf = acceleratorAllocDevice(bytes);
   if (ShmCommBuf == (void *)NULL ) {
diff --git a/Grid/stencil/Stencil.h b/Grid/stencil/Stencil.h
index 917a10be..a768f344 100644
--- a/Grid/stencil/Stencil.h
+++ b/Grid/stencil/Stencil.h
@@ -368,6 +368,7 @@ public:
     //    accelerator_barrier();     // All kernels should ALREADY be complete
     //    _grid->StencilBarrier();   // Everyone is here, so noone running slow and still using receive buffer
                                // But the HaloGather had a barrier too.
+#ifdef ACCELERATOR_AWARE_MPI
     for(int i=0;i<Packets.size();i++){
       _grid->StencilSendToRecvFromBegin(MpiReqs,
 					Packets[i].send_buf,
@@ -376,6 +377,23 @@ public:
 					Packets[i].from_rank,Packets[i].do_recv,
 					Packets[i].xbytes,Packets[i].rbytes,i);
     }
+#else
+#warning "Using COPY VIA HOST BUFFERS IN STENCIL"
+    for(int i=0;i<Packets.size();i++){
+      // Introduce a host buffer with a cheap slab allocator and zero cost wipe all
+      Packets[i].host_send_buf = _grid->HostBufferMalloc(Packets[i].xbytes);
+      Packets[i].host_recv_buf = _grid->HostBufferMalloc(Packets[i].rbytes);
+      if ( Packets[i].do_send ) {
+	acceleratorCopyFromDevice(Packets[i].send_buf, Packets[i].host_send_buf,Packets[i].xbytes);
+      }
+      _grid->StencilSendToRecvFromBegin(MpiReqs,
+					Packets[i].host_send_buf,
+					Packets[i].to_rank,Packets[i].do_send,
+					Packets[i].host_recv_buf,
+					Packets[i].from_rank,Packets[i].do_recv,
+					Packets[i].xbytes,Packets[i].rbytes,i);
+    }
+#endif
     // Get comms started then run checksums
     // Having this PRIOR to the dslash seems to make Sunspot work... (!)
     for(int i=0;i<Packets.size();i++){
@@ -393,8 +411,16 @@ public:
     else DslashLogFull();
     //    acceleratorCopySynchronise();// is in the StencilSendToRecvFromComplete
     //    accelerator_barrier(); 
+#ifndef ACCELERATOR_AWARE_MPI
+#warning "Using COPY VIA HOST BUFFERS IN STENCIL"
+    for(int i=0;i<Packets.size();i++){
+      if ( Packets[i].do_recv ) {
+	acceleratorCopyToDevice(Packets[i].host_recv_buf, Packets[i].recv_buf,Packets[i].rbytes);
+      }
+    }
+    _grid->HostBufferFreeAll();
+#endif    // run any checksums
     _grid->StencilBarrier(); 
-    // run any checksums
     for(int i=0;i<Packets.size();i++){
       if ( Packets[i].do_recv )
 	FlightRecorder::recvLog(Packets[i].recv_buf,Packets[i].rbytes,Packets[i].from_rank);
diff --git a/configure.ac b/configure.ac
index 776ca264..0b71b834 100644
--- a/configure.ac
+++ b/configure.ac
@@ -240,6 +240,18 @@ case ${ac_SFW_FP16} in
 esac
 
 
+############### MPI BOUNCE TO HOST
+AC_ARG_ENABLE([accelerator-aware-mpi],
+    [AS_HELP_STRING([--enable-accelerator-aware-mpi=yes|no],[run mpi transfers from device])],
+    [ac_ACCELERATOR_AWARE_MPI=${enable_accelerator_aware_mpi}], [ac_ACCELERATOR_AWARE_MPI=yes])
+
+case ${ac_ACCELERATOR_AWARE_MPI} in
+    yes)
+      AC_DEFINE([ACCELERATOR_CSHIFT],[1],[ Cshift runs on device])
+      AC_DEFINE([ACCELERATOR_AWARE_MPI],[1],[ Stencil can use device pointers]);;
+    *);;
+esac
+
 ############### SYCL/CUDA/HIP/none
 AC_ARG_ENABLE([accelerator],
     [AS_HELP_STRING([--enable-accelerator=cuda|sycl|hip|none],[enable none,cuda,sycl,hip acceleration])],
diff --git a/systems/Aurora/benchmarks/bench1.pbs b/systems/Aurora/benchmarks/bench1.pbs
index e85dc09e..14dad48a 100644
--- a/systems/Aurora/benchmarks/bench1.pbs
+++ b/systems/Aurora/benchmarks/bench1.pbs
@@ -1,6 +1,7 @@
 #!/bin/bash
 
-#PBS -q EarlyAppAccess
+##PBS -q EarlyAppAccess
+#PBS -q debug
 #PBS -l select=1
 #PBS -l walltime=00:20:00
 #PBS -A LatticeQCD_aesp_CNDA
@@ -12,27 +13,24 @@ source ../sourceme.sh
 cp $PBS_NODEFILE nodefile
 
 export OMP_NUM_THREADS=4
-export MPIR_CVAR_CH4_OFI_ENABLE_GPU_PIPELINE=1
-unset MPIR_CVAR_CH4_OFI_GPU_PIPELINE_D2H_ENGINE_TYPE
-unset MPIR_CVAR_CH4_OFI_GPU_PIPELINE_H2D_ENGINE_TYPE
-unset MPIR_CVAR_GPU_USE_IMMEDIATE_COMMAND_LIST
+export MPICH_OFI_NIC_POLICY=GPU
+
+#export MPIR_CVAR_CH4_OFI_ENABLE_GPU_PIPELINE=1
+#unset MPIR_CVAR_CH4_OFI_GPU_PIPELINE_D2H_ENGINE_TYPE
+#unset MPIR_CVAR_CH4_OFI_GPU_PIPELINE_H2D_ENGINE_TYPE
+#unset MPIR_CVAR_GPU_USE_IMMEDIATE_COMMAND_LIST
 #export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_D2H_ENGINE_TYPE=0
 #export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_H2D_ENGINE_TYPE=0
 #export MPIR_CVAR_GPU_USE_IMMEDIATE_COMMAND_LIST=1
-export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_BUFFER_SZ=1048576
-export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_THRESHOLD=131072
-export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_NUM_BUFFERS_PER_CHUNK=16
-export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_MAX_NUM_BUFFERS=16
-export MPICH_OFI_NIC_POLICY=GPU
+#export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_BUFFER_SZ=1048576
+#export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_THRESHOLD=131072
+#export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_NUM_BUFFERS_PER_CHUNK=16
+#export MPIR_CVAR_CH4_OFI_GPU_PIPELINE_MAX_NUM_BUFFERS=16
 
 CMD="mpiexec -np 12 -ppn 12  -envall \
-	     ./Benchmark_dwf_fp32 --mpi 2.1.2.3 --grid 32.32.64.48 \
-		--shm-mpi 1 --shm 2048 --device-mem 32000 --accelerator-threads 32 --debug-signals"
+	     ./gpu_tile.sh ./Benchmark_dwf_fp32 --mpi 2.1.2.3 --grid 32.32.64.96 \
+		--shm-mpi 1 --shm 2048 --device-mem 32000 --accelerator-threads 8 "
 
-#for f in 1 2 3 4 5 6 7 8
-for f in 1
-do
 echo $CMD
-$CMD | tee 1node.32.32.64.48.dwf.hbm.$f
-done
+$CMD
 
diff --git a/systems/Aurora/benchmarks/gpu_tile_compact.sh b/systems/Aurora/benchmarks/gpu_tile_compact.sh
deleted file mode 100755
index 099a0ded..00000000
--- a/systems/Aurora/benchmarks/gpu_tile_compact.sh
+++ /dev/null
@@ -1,34 +0,0 @@
-#!/bin/bash
-
-#export NUMA_MAP=(2 2 2 3 3 3 2 2 2 3 3 3 )
-#export NUMA_MAP=(0 0 1 1 0 0 1 1 0 0 1 1);
-#export  GPU_MAP=(0.0 0.1 3.0 3.1 1.0 1.1 4.0 4.1 2.0 2.1 5.0 5.1)
-
-export NUMA_MAP=(0 0 0 0 0 0 1 1 1 1 1 1 );
-export  GPU_MAP=(0.0 1.0 2.0 3.0 4.0 5.0 0.1 1.1 2.1 3.1 4.1 5.1 )
-
-export NUMA=${NUMA_MAP[$PALS_LOCAL_RANKID]}
-export gpu_id=${GPU_MAP[$PALS_LOCAL_RANKID]}
-  
-unset EnableWalkerPartition
-export EnableImplicitScaling=0
-export ZE_AFFINITY_MASK=$gpu_id
-export ONEAPI_DEVICE_FILTER=gpu,level_zero
-
-export SYCL_PI_LEVEL_ZERO_DEVICE_SCOPE_EVENTS=0
-export SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=1
-export SYCL_PI_LEVEL_ZERO_USE_COPY_ENGINE=0:5
-#export SYCL_PI_LEVEL_ZERO_USE_COPY_ENGINE=0:2
-export SYCL_PI_LEVEL_ZERO_USE_COPY_ENGINE_FOR_D2D_COPY=1
-#export SYCL_PI_LEVEL_ZERO_USM_RESIDENT=1
-
-echo "rank $PALS_RANKID ; local rank $PALS_LOCAL_RANKID ; ZE_AFFINITY_MASK=$ZE_AFFINITY_MASK ; NUMA $NUMA "
-
-if [ $PALS_RANKID = "0" ]
-then
-#    numactl -m $NUMA -N $NUMA onetrace --chrome-device-timeline  "$@"
-#    numactl -m $NUMA -N $NUMA unitrace --chrome-kernel-logging --chrome-mpi-logging --chrome-sycl-logging --demangle "$@"
-    numactl -m $NUMA -N $NUMA  "$@"
-else 
-    numactl -m $NUMA -N $NUMA  "$@"
-fi
diff --git a/systems/Aurora/config-command b/systems/Aurora/config-command
index 423da64e..64bef44b 100644
--- a/systems/Aurora/config-command
+++ b/systems/Aurora/config-command
@@ -17,7 +17,7 @@ export CXXFLAGS="-O3 -fiopenmp -fsycl-unnamed-lambda -fsycl -Wno-tautological-co
 	--with-lime=$CLIME \
 	--enable-shm=nvlink \
 	--enable-accelerator=sycl \
-	--enable-accelerator-aware-mpi=yes\
+	--enable-accelerator-aware-mpi=no\
 	--enable-unified=no \
 	MPICXX=mpicxx \
 	CXX=icpx 
diff --git a/systems/Aurora/sourceme.sh b/systems/Aurora/sourceme.sh
index 4cdc950f..89126f5b 100644
--- a/systems/Aurora/sourceme.sh
+++ b/systems/Aurora/sourceme.sh
@@ -2,6 +2,7 @@
 #module load mpich/icc-all-debug-pmix-gpu/52.2
 #module load mpich-config/mode/deterministic
 #module load intel_compute_runtime/release/821.35
+module load pti-gpu
 
 source ~/spack/share/spack/setup-env.sh 
 spack load c-lime