packet pool: add option to control global/per-device packet pool (LCI_USE_GLOBAL_PACKET_POOL)

lci/api/lci.h

@@ -624,6 +624,12 @@ extern bool LCI_UCX_USE_TRY_LOCK;
  */
 extern bool LCI_UCX_PROGRESS_FOCUSED;
 
+/**
+ * @ingroup LCI_DEVICE
+ * @brief Whether to use global packet pool.
+ */
+extern bool LCI_USE_GLOBAL_PACKET_POOL;
+
 /**
  * @ingroup LCI_COMM
  * @brief Try_lock mode of network backend.

lci/runtime/device.c

@@ -60,7 +60,12 @@ LCI_error_t LCI_device_init(LCI_device_t* device_ptr)
   LCII_bq_init(&device->bq);
   LCIU_spinlock_init(&device->bq_spinlock);
 
-  device->heap = &g_heap;
+  if (LCI_USE_GLOBAL_PACKET_POOL) {
+    device->heap = &g_heap;
+  } else {
+    device->heap = LCIU_malloc(sizeof(LCII_packet_heap_t));
+    initialize_packet_heap(device->heap);
+  }
   LCI_memory_register(device, device->heap->address, device->heap->length,
                       &device->heap_segment);
 
@@ -89,6 +94,10 @@ LCI_error_t LCI_device_free(LCI_device_t* device_ptr)
         LCII_endpoint_get_recv_posted(device->endpoint_progress);
   }
   LCI_memory_deregister(&device->heap_segment);
+  if (!LCI_USE_GLOBAL_PACKET_POOL) {
+    finalize_packet_heap(device->heap);
+    LCIU_free(device->heap);
+  }
   LCII_matchtable_free(&device->mt);
   LCM_archive_fini(&(device->ctx_archive));
   LCII_bq_fini(&device->bq);

lci/runtime/env.c

@@ -34,6 +34,7 @@ LCI_API bool LCI_OFI_CXI_TRY_NO_HACK;
 LCI_API uint64_t LCI_BACKEND_TRY_LOCK_MODE;
 LCI_API bool LCI_UCX_USE_TRY_LOCK;
 LCI_API bool LCI_UCX_PROGRESS_FOCUSED;
+LCI_API bool LCI_USE_GLOBAL_PACKET_POOL;
 LCI_API LCI_device_t LCI_UR_DEVICE;
 LCI_API LCI_endpoint_t LCI_UR_ENDPOINT;
 LCI_API LCI_comp_t LCI_UR_CQ;
@@ -129,6 +130,7 @@ void LCII_env_init(int num_proc, int rank)
   }
   LCI_UCX_USE_TRY_LOCK = LCIU_getenv_or("LCI_UCX_USE_TRY_LOCK", 0);
   LCI_UCX_PROGRESS_FOCUSED = LCIU_getenv_or("LCI_UCX_PROGRESS_FOCUSED", 0);
+  LCI_USE_GLOBAL_PACKET_POOL = LCIU_getenv_or("LCI_USE_GLOBAL_PACKET_POOL", 1);
   if (LCI_UCX_PROGRESS_FOCUSED) LCI_UCX_USE_TRY_LOCK = true;
   LCII_env_init_cq_type();
   LCII_env_init_rdv_protocol();

lci/runtime/lci.c

@@ -79,7 +79,7 @@ LCI_error_t LCI_initialize()
   }
   // initialize global data structure
   LCIS_server_init(&g_server);
-  initialize_packet_heap(&g_heap);
+  if (LCI_USE_GLOBAL_PACKET_POOL) initialize_packet_heap(&g_heap);
   // UR objects
   LCI_device_init(&LCI_UR_DEVICE);
   LCI_queue_create(LCI_UR_DEVICE, &LCI_UR_CQ);
@@ -111,7 +111,7 @@ LCI_error_t LCI_finalize()
   LCI_queue_free(&LCI_UR_CQ);
   LCI_device_free(&LCI_UR_DEVICE);
   LCIS_server_fina(g_server);
-  finalize_packet_heap(&g_heap);
+  if (LCI_USE_GLOBAL_PACKET_POOL) finalize_packet_heap(&g_heap);
   if (LCI_USE_DREG) {
 #ifdef LCI_COMPILE_DREG
     LCII_ucs_cleanup();

lci/runtime/lcii.h

@@ -66,6 +66,8 @@ struct LCII_packet_heap_t {
   int total_recv_posted;  // for debugging purpose
 };
 typedef struct LCII_packet_heap_t LCII_packet_heap_t;
+extern void initialize_packet_heap(LCII_packet_heap_t* heap);
+extern void finalize_packet_heap(LCII_packet_heap_t* heap);
 
 extern LCIS_server_t g_server;
 extern LCII_packet_heap_t g_heap;

lci/runtime/packet_pool.h

@@ -10,7 +10,7 @@
 #include <assert.h>
 
 #define MAX_NPOOLS 272
-#define MAX_LOCAL_POOL 32  // align to a cache line.
+#define MAX_LOCAL_POOL 128
 
 extern int LCII_pool_nkey;
 extern int32_t LCII_tls_pool_metadata[MAX_NPOOLS][MAX_LOCAL_POOL];