GPU-aware load balancer

.github/workflows/reconverse-cuda.yaml

@@ -170,6 +170,13 @@ jobs:
         run: |
           make -j4 -C netlrts-linux-x86_64-cuda/tests/charm++/cuda/d2dtest \
             GPU=cuda CUDATOOLKIT_HOME=/usr OPTS="-g"
+      - name: gpumigrate compiles
+        # Device-state migration. Running it needs a GPU, but compiling it
+        # gates the PUPMode::DEVICE surface that application pup routines see,
+        # and the nocopydevice entry method the migration payload travels on.
+        run: |
+          make -j4 -C netlrts-linux-x86_64-cuda/tests/charm++/cuda/gpumigrate \
+            OPTS="-g"
 
       - name: legacy buildold path
         run: |

.gitignore

@@ -14,6 +14,9 @@ configure~
 *.def.h
 *.ci.stamp
 
+# Ignore crash backtrace dumps written beside a crashing binary
+*.btr
+
 # Ignore build artifacts
 config_opts.sh
 smart-build.log

CMakeLists.txt

@@ -755,6 +755,11 @@ if(BUILD_CUDA OR BUILD_HIP)
       message(WARNING "CUPTI library not found. GPU load balancing will not be available.")
     else()
       message(STATUS "Found CUPTI: ${CUPTI_LIBRARY}")
+      # Global, not per-target: the GPU load fields live in LDObjData, which
+      # ck-core, ck-ldb and hybridapi all see. A per-target definition would
+      # give them different layouts for the same struct.
+      add_compile_definitions(HAPI_CUPTI_LB)
+      set(HAPI_CUPTI_LB ON)
     endif()
 
     add_library(hybridapi ${hybridAPI-cxx-sources})

cmake/converse.cmake

@@ -327,6 +327,9 @@ if(NOT RECONVERSE)
     # be shadowed.
     list(APPEND conv-core-h-to-install src/conv-core/converse.h)
     list(APPEND conv-core-h-to-install src/conv-core/charm-config.h)
+    # Classic seed-balancer header: ck.h includes it when CMK_USE_SHMEM is on.
+    # Reconverse builds get reconverse's own cldb.h from cmake/reconverse.
+    list(APPEND conv-core-h-to-install src/conv-ldb/cldb.h)
 endif()
 if(RECONVERSE)
     # Reconverse provides its own conv-rdma.h, which is the authority on the

doc/charm++/manual.rst

@@ -2679,6 +2679,51 @@ infrastructure:
   options to point to the include and library directories used,
   respectively. (``+balancer ScotchLB``)
 
+The following centralized balancer targets GPU applications, and is
+available on CUDA builds only:
+
+- **GreedyRefineCentralGPULB**: Balances at two granularities at once.
+  Objects are assigned across *GPU groups* — the sets of PEs that share a
+  device — by their measured GPU load, and then across the PEs within a
+  group by their host load. This matches the usual arrangement in which
+  each process owns a GPU: which device an object runs on is decided by
+  its device work, and which PE drives it by its host work.
+  (``+balancer GreedyRefineCentralGPULB``)
+
+  The GPU load it reads is measured by CUPTI, per object, and normalized
+  to seconds of whole-device occupancy, so it estimates what an object
+  would cost on a *different* device rather than only what it cost where
+  it ran. Measurement follows ``LBTurnInstrumentOn()`` /
+  ``LBTurnInstrumentOff()``, so an application that instruments a window
+  around its own ``AtSync`` pays the tracing cost only inside it. See
+  :numref:`lbFramework` for the instrumentation calls.
+
+  Like the other greedy-refine balancers it accepts a *tolerance*: how far
+  above the maximum load plain greedy would produce it may go, in exchange
+  for migrating fewer objects. ``+LBGreedyRefineTolerance 1.1`` allows a
+  maximum load 10% higher than greedy's, and an object then stays where it
+  is for as long as its PE remains under that target.
+
+  The default is ``1.1``. Passing ``0`` or less asks the balancer to search
+  for a tolerance instead of being told one: each PE builds a candidate
+  assignment from a different parameter pair and the best is chosen, which
+  costs a reduction per balancing step and is bounded by the PE count.
+  ``+LBDebug 1`` reports the migration count and the achieved maximum load
+  against greedy's, separately for the GPU and host dimensions, which is
+  what says whether the tolerance is buying anything.
+
+.. note::
+
+   Migrating a GPU application requires the runtime to be built with
+   ``CMK_GLOBAL_LOCATION_UPDATE`` (see `Global Location Update`_). A device
+   zerocopy send is addressed to the PE the sender believes hosts the
+   target, so without it a send issued around a migration arrives at a PE
+   that no longer hosts the element, and the receive path aborts rather
+   than silently reading the wrong buffer (``ckrdmadevice.C``). Note the
+   restriction that comes with the option: migrations must happen at load
+   balancing steps, so a GPU application must migrate through ``AtSync()``
+   rather than ``migrateMe()``.
+
 In distributed approaches, the strategy executes across multiple PEs,
 providing scalable computational and communication performance.
 

src/arch/common/conv-mach-cuda.sh

@@ -1,4 +1,9 @@
 BUILD_CUDA=1
-CMK_INCDIR="-I$CUDA_DIR/include $CMK_INCDIR "
-CMK_LIBDIR="-L$CUDA_DIR/lib64 $CMK_LIBDIR "
-CMK_LIBS="-lhybridapi -lcudart -lrt $CMK_LIBS "
+CMK_INCDIR="-I$CUDA_DIR/include -I$CUDA_DIR/extras/CUPTI/include $CMK_INCDIR "
+CMK_LIBDIR="-L$CUDA_DIR/lib64 -L$CUDA_DIR/extras/CUPTI/lib64 $CMK_LIBDIR "
+CMK_LIBS="-lhybridapi -lcudart -lcupti -lrt $CMK_LIBS "
+# HAPI_CUPTI_LB turns on per-object GPU load measurement. It has to be a global
+# define rather than one confined to hybridapi: the fields it guards live in
+# LDObjData, which ck-core and ck-ldb see too, and a partial definition would
+# give them different layouts for the same struct.
+CMK_DEFS="$CMK_DEFS -DHAPI_CUPTI_LB "

src/arch/cuda/hybridAPI/devicemanager.h

@@ -19,8 +19,23 @@ struct DeviceManager {
   // Buddy allocator for communication buffer
   buddy::allocator* comm_buffer;
 
+  // Device properties needed to estimate how many SMs a kernel occupies while
+  // it runs, which is what turns a raw kernel timeline into a load. Filled
+  // lazily by hapiPopulateDeviceProps, because device_managers is not
+  // populated yet when a DeviceManager is constructed.
+  int multi_processor_count;
+  int max_threads_per_sm;
+  int max_blocks_per_sm;
+  int max_registers_per_sm;
+  int max_shared_mem_per_sm;
+  int warp_size;
+  bool props_initialized;
+
   DeviceManager(int local_index_, int global_index_) :
-    local_index(local_index_), global_index(global_index_), comm_buffer(nullptr) {
+    local_index(local_index_), global_index(global_index_), comm_buffer(nullptr),
+    multi_processor_count(0), max_threads_per_sm(0), max_blocks_per_sm(0),
+    max_registers_per_sm(0), max_shared_mem_per_sm(0), warp_size(0),
+    props_initialized(false) {
 #if CMK_SMP
     lock = CmiCreateLock();
 #endif

src/arch/cuda/hybridAPI/gpumanager.h

@@ -13,6 +13,12 @@
 
 #include <unordered_map>
 #include <queue>
+#include <mutex>
+#include <atomic>
+
+#if CMK_CUDA && CMK_LBDB_ON
+#include "lbdb.h"  // for LBKernelRecord, LDObjKey and GpuObjectTokenTable
+#endif
 
 // Local rank of the logical node (process) that the given PE belongs to,
 // within its physical node: the logical node index modulo the number of
@@ -55,7 +61,7 @@ struct hapi_ipc_event_shared {
   pthread_mutex_t lock;
 };
 
-#ifdef HAPI_CUPTI_LB
+#if CMK_CUDA && CMK_LBDB_ON
 struct CuptiBufferItem {
   uint8_t* buffer;
   size_t validSize;
@@ -183,14 +189,89 @@ struct GPUManager {
   std::vector<hapi_ipc_device_info> hapi_ipc_device_infos;
 
   //CUPTI load balancing
-#ifdef HAPI_CUPTI_LB
-  std::unordered_map<uint32_t, uint64_t> cupti_correlation_db_;//correlationID -> ObjectID
-
-  std::unordered_map<uint64_t, uint64_t> cupti_obj_gpu_times_;//objectID -> accumulated GPU time in ns
-
+#if CMK_CUDA && CMK_LBDB_ON
+  // Runtime correlation ID -> process-local full-object token.
+  std::unordered_map<uint32_t, uint64_t> cupti_object_correlation_db_;
+
+  GpuObjectTokenTable cupti_object_tokens_;
+  std::mutex cupti_object_token_lock_;
+
+  // Full LB object identity -> attributed kernel records.
+  std::unordered_map<LDObjKey, std::vector<LBKernelRecord>, LDObjKeyHash>
+      cupti_obj_kernel_records_;
+
+  // Kernels that could not be attributed to any object (launched outside a
+  // migratable entry method, or with no correlation record). They occupy SMs
+  // and so take part in the sweep-line as contention, but receive no load.
+  // Kept separate rather than under a sentinel object ID, because 0 is a
+  // perfectly valid chare element ID.
+  std::vector<LBKernelRecord> cupti_unattributed_kernels_;
+
+  // Full object identity -> SM-utilization-normalized GPU load in seconds.
+  std::unordered_map<LDObjKey, double, LDObjKeyHash> cupti_obj_norm_load_;
+
+  // Previous round's count of kernels whose correlations had not been parsed
+  // yet, used to size the parked vector.
+  uint32_t cupti_pending_hint_ = 0;
+
+  // Written by CUPTI's buffer-completed callback, which may run on a
+  // CUPTI-owned thread. That thread exists in non-SMP builds too, where the
+  // Converse locks compile out, so this needs a real mutex rather than a
+  // CmiNodeLock.
   std::queue<CuptiBufferItem> cupti_buffer_queue_;
+  std::mutex cupti_queue_lock_;
 
   bool cupti_initialized_;
+  // Whether activity tracing is currently running. Separate from
+  // cupti_initialized_: the buffer callbacks are registered once, but tracing
+  // itself is switched on and off as the application asks for it.
+  //
+  // Atomic because the entry-method hooks read it on every invocation from
+  // every PE thread while another thread may be switching tracing on or off.
+  std::atomic<bool> cupti_tracing_active_{false};
+  // Serializes hapiCuptiStartTracing/hapiCuptiStopTracing. This state lives in
+  // the node-wide GPUManager, but the switch is reached per-PE through
+  // LBDatabase::TurnStatsOn/Off, so every PE thread calls in. Without this,
+  // several threads enable or disable the same CUPTI activity kinds and flush
+  // concurrently, which corrupts CUPTI's internal buffer bookkeeping and shows
+  // up later as heap corruption in an unrelated allocation. Must NOT be the
+  // same mutex as cupti_queue_lock_: the flush in the stop path invokes the
+  // buffer-completed callback, which takes that one.
+  std::mutex cupti_tracing_lock_;
+  // PEs of this process whose instrumentation is currently on. Tracing runs
+  // while this is non-zero; see hapiCuptiStartTracing.
+  int cupti_tracing_users_ = 0;
+  // Bumped every time CUPTI is detached. Detaching clears CUPTI's
+  // external-correlation stack for every PE, but the counters that keep the
+  // entry-method push/pop hooks paired are per-PE, and only the PE that ran the
+  // detach could reset its own. Each PE compares this against the generation it
+  // last saw and zeroes its counter when they differ.
+  uint64_t cupti_generation_ = 0;
+  // Serializes turning this round's raw CUPTI records into cupti_obj_norm_load_,
+  // and makes that work happen exactly once per LB round no matter how many PE
+  // threads ask for it. Doing it with "rank 0 works between two CmiNodeBarrier
+  // calls" does not hold: both barriers sit behind #if CMK_SMP, which is 0 in
+  // the multicore build even though a process really does run many PE threads,
+  // so the barriers vanish and the other ranks read cupti_obj_norm_load_ while
+  // rank 0 is rebuilding it. A lock also cannot deadlock the way a spin barrier
+  // can: a PE waiting here waits on a PE that is running, not on one that has
+  // yet to arrive.
+  std::mutex cupti_prepare_lock_;
+  // Set once this round's loads are built; cleared by hapiClearCuptiData.
+  bool cupti_loads_ready_ = false;
+  // Epoch the built loads correspond to, so a second caller in the same round
+  // reads what the first built instead of rebuilding it.
+  uint64_t cupti_loads_epoch_ = 0;
+  // Arrival gate for the per-round load build; see hapiCuptiArrive. A load
+  // balancer's per-PE barrier fires when THAT PE's objects are at AtSync, and
+  // the first PE to fire would otherwise flush, drain and clear the
+  // process-wide CUPTI records on the spot. Any PE whose objects were still
+  // finishing kernels at that instant has them dropped from the round --
+  // consistently the last PE of the process, whose objects then read as zero
+  // GPU load every step.
+  std::mutex cupti_arrive_lock_;
+  uint64_t cupti_arrive_epoch_ = 0;
+  int cupti_arrive_count_ = 0;
 #endif
 
   void init() {
@@ -234,6 +315,15 @@ struct GPUManager {
     hapi_ipc_event_pool_size_pe = -1;
     hapi_ipc_event_pool_size_total = -1;
 
+#if CMK_CUDA && CMK_LBDB_ON
+    // CUPTI load balancing
+    cupti_initialized_ = false;
+    cupti_tracing_active_.store(false, std::memory_order_relaxed);
+    cupti_generation_ = 0;
+    cupti_loads_ready_ = false;
+    cupti_loads_epoch_ = 0;
+#endif
+
     // Allocate host/device buffers array (both user and system-addressed)
     host_buffers_ = new void*[NUM_BUFFERS*2];
     device_buffers_ = new void*[NUM_BUFFERS*2];

src/arch/cuda/hybridAPI/hapi.h

@@ -2,11 +2,13 @@
 #define __HAPI_H_
 #include "hapi_portable.h"
 
-/* HAPI_CUPTI_LB: per-object GPU time attribution (CUPTI activity tracing,
- * event-based kernel timing, launch wrappers) feeding GPU-aware load
- * balancing. Dormant -- nothing defines it -- until the GPU-LB series
- * (plan item 11) enables it together with its ck-ldb counterparts
- * (setObjGPUTime and friends, LBHasBalancersRegistered). */
+/* HAPI_CUPTI_LB: per-object GPU time attribution via CUPTI activity tracing,
+ * feeding GPU-aware load balancing. Defined for a CUDA build in which CUPTI
+ * was found; without it the entry points below are no-ops and every object
+ * reads zero GPU load, so a GPU-aware balancer falls back to the host
+ * dimension. The data structures the attribution produces are declared under
+ * CMK_CUDA rather than under this macro, so that LDObjData has one layout
+ * across every translation unit either way. */
 
 /* See hapi_functions.h for the majority of function declarations provided
  * by the Hybrid API. */
@@ -263,29 +265,50 @@ static inline hapiError_t hapiFreeHost_Pool(void* ptr, bool pool) {
   return hapiFreeHost(ptr, pool);
 }
 
-void hapiRecordTime(hapiStream_t stream, hapiEvent_t start);
-#ifdef HAPI_CUPTI_LB
+#if CMK_CUDA && CMK_LBDB_ON
 void hapiCuptiInit();
 void hapiCuptiFinalize();
+
+// Stamp/unstamp the running migratable object onto every kernel launched
+// inside the scope. Called around every entry method, so both are no-ops
+// unless tracing is running.
 uint64_t hapiCuptiPushObjCorrelation();
 void hapiCuptiPopObjCorrelation();
+
+// Epoch a load balancer passes to hapiPrepareCuptiLoads: larger than any
+// epoch an application-level sampler will use, so an LB round always rebuilds
+// rather than reading a sampler's older loads.
+#define HAPI_CUPTI_EPOCH_LB_ROUND UINT64_MAX
+
+// Flush, parse and normalize the CUPTI records accumulated since the last
+// hapiClearCuptiData, once per epoch however many PE threads call it. The
+// result is GPUManager::cupti_obj_norm_load_, which every PE of the process
+// then reads; see the comment on GPUManager::cupti_prepare_lock_.
+void hapiPrepareCuptiLoads(uint64_t epoch = HAPI_CUPTI_EPOCH_LB_ROUND);
 void hapiProcessCuptiBuffers();
+void hapiNormalizeCuptiLoads();
 void hapiClearCuptiData();
-#endif
 
-#ifdef HAPI_CUPTI_LB
-#define HAPI_LAUNCH_KERNEL_WRAPPER(call, stream)\
-    hapiEvent_t start;\
-    hapiEventCreate(&start);\
-    hapiEventRecord(start, stream);\
-    call;\
-    hapiRecordTime(stream, start);
-#else
-#define HAPI_LAUNCH_KERNEL_WRAPPER(call, stream)\
-    call;
+// Arrival gate in front of hapiPrepareCuptiLoads for a load-balancing round.
+// Returns true to exactly one caller per epoch, once `expected` callers have
+// arrived: the records are process-wide, so the drain has to wait for the last
+// PE of the process rather than run on the first.
+bool hapiCuptiArrive(uint64_t epoch, int expected);
+
+// Start/stop CUPTI activity tracing. Tracing is the dominant cost of GPU load
+// measurement, so an application that instruments a window rather than the
+// whole run pays for it only inside that window. Reached per-PE through
+// LBDatabase::TurnStatsOn/Off; the process traces while any of its PEs wants
+// instrumentation.
+void hapiCuptiStartTracing();
+void hapiCuptiStopTracing();
+bool hapiCuptiTracingActive();
 #endif
 
-#ifdef HAPI_CUPTI_LB
+// Attribute one kernel launch to the running object. The runtime already
+// brackets every entry method this way (see CkCallstackPush/Pop), so an
+// application needs this only for a launch it makes outside one.
+#if CMK_CUDA && CMK_LBDB_ON
 #define CUPTI_LAUNCH_WRAPPER(call)\
   hapiCuptiPushObjCorrelation();\
   call;\

src/arch/cuda/hybridAPI/hapi_functions.h

@@ -49,11 +49,6 @@ AMPI_CUSTOM_FUNC(void, hapiAddCallback, hapiStream_t, void*, void*)
 // AMPI_CUSTOM_FUNC(cudaError_t, hapiMemcpyAsync, void*, const void*, size_t, enum cudaMemcpyKind, cudaStream_t)
 // AMPI_CUSTOM_FUNC(cudaError_t, hapiMemcpy2DAsync, void*, size_t, const void*, size_t, size_t, size_t, enum cudaMemcpyKind, cudaStream_t)
 
-// Kernel launch wrapper
-#ifdef HAPI_CUPTI_LB  /* pairs with HAPI_LAUNCH_KERNEL_WRAPPER; lands with plan item 11 */
-AMPI_CUSTOM_FUNC(hapiError_t, hapiLaunchKernel, const void*, dim3, dim3, void**, size_t, hapiStream_t)
-#endif
-
 // Explicit memory allocations using pinned memory pool.
 AMPI_CUSTOM_FUNC(hapiError_t, hapiPoolMalloc, void**, size_t)
 AMPI_CUSTOM_FUNC(hapiError_t, hapiPoolFree, void*)
@@ -63,6 +58,9 @@ AMPI_CUSTOM_FUNC(void, hapiErrorDie, hapiError_t, const char*, const char*, int)
 
 // Returns the GPU device index this PE is mapped to (set during hapiMapping).
 AMPI_CUSTOM_FUNC(uint64_t, hapiMyDevice, void)
+// SM count of that device, or 0 if it has not been queried yet. A GPU-aware
+// balancer needs it to compare devices of different sizes.
+AMPI_CUSTOM_FUNC(int, hapiMyDeviceTotalSMs, void)
 
 #ifdef HAPI_INSTRUMENT_WRS
 AMPI_CUSTOM_FUNC(void, hapiInitInstrument, int n_chares, char n_types)