Support for prepared transfer API in nixlbench

benchmark/nixlbench/src/utils/utils.cpp

@@ -86,6 +86,7 @@ NB_ARG_BOOL(recreate_xfer,
             false,
             "Recreate xfer each iteration (default: false for all backends, true for GUSLI)");
 NB_ARG_BOOL(reregister_mem, false, "Register and deregister memory on every iteration");
+NB_ARG_BOOL(prepped_xfer, false, "Use prepared transfer API (prepare+make)");
 NB_ARG_INT32(pipeline_depth, 1, "Number of transfer requests in flight simultaneously");
 NB_ARG_INT32(large_blk_iter_ftr,
              16,
@@ -276,6 +277,7 @@ std::string xferBenchConfig::filepath = "";
 std::string xferBenchConfig::filenames = "";
 bool xferBenchConfig::storage_enable_direct = false;
 bool xferBenchConfig::reregister_mem = false;
+bool xferBenchConfig::prepped_xfer = false;
 int xferBenchConfig::pipeline_depth = 1;
 long xferBenchConfig::page_size = sysconf(_SC_PAGESIZE);
 std::string xferBenchConfig::obj_access_key = "";
@@ -494,6 +496,7 @@ xferBenchConfig::loadParams(void) {
     storage_enable_direct = NB_ARG(storage_enable_direct);
     recreate_xfer = NB_ARG(recreate_xfer);
     reregister_mem = NB_ARG(reregister_mem);
+    prepped_xfer = NB_ARG(prepped_xfer);
     pipeline_depth = NB_ARG(pipeline_depth);
     if (pipeline_depth < 1) {
         std::cerr << "pipeline_depth must be >= 1" << std::endl;
@@ -516,6 +519,12 @@ xferBenchConfig::loadParams(void) {
                   << " Setting recreate_xfer to true." << std::endl;
         recreate_xfer = true;
     }
+    if (prepped_xfer && reregister_mem) {
+        std::cerr << "prepped_xfer is incompatible with reregister_mem: the prepared "
+                     "descriptor list handles pin the registration."
+                  << std::endl;
+        return -1;
+    }
 
     // Validate runtime configuration
     if (!isStorageBackend()) {
@@ -649,6 +658,8 @@ xferBenchConfig::printConfig() {
                     std::to_string(recreate_xfer));
         printOption("Re-register memory each iteration (--reregister_mem=[0,1])",
                     std::to_string(reregister_mem));
+        printOption("Prepped xfer (prep+make) (--prepped_xfer=[0,1])",
+                    std::to_string(prepped_xfer));
         printOption("Pipeline depth (--pipeline_depth=N)", std::to_string(pipeline_depth));
         printOption("Use hugepages (--use_hugepages=[0,1])", std::to_string(use_hugepages));
 

benchmark/nixlbench/src/utils/utils.h

@@ -175,6 +175,7 @@ class xferBenchConfig {
     static int posix_kernel_queue_size;
     static bool storage_enable_direct;
     static bool reregister_mem;
+    static bool prepped_xfer;
     static int pipeline_depth;
     static int gds_batch_pool_size;
     static int gds_batch_limit;

benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp

@@ -29,6 +29,7 @@
 #include <filesystem>
 #include <iomanip>
 #include <memory>
+#include <numeric>
 #include <sstream>
 #include "utils/neuron.h"
 #include "utils/utils.h"
@@ -1418,6 +1419,9 @@ struct slotState {
     bool in_flight = false;
     bool registered = false;
     nixlTime::us_t post_ts = 0;
+    nixlDlistH *prep_local_dlist = nullptr;
+    nixlDlistH *prep_remote_dlist = nullptr;
+    std::vector<int> indices;
 };
 
 // Register memory (if --reregister_mem) and create the XferReq for a slot
@@ -1444,10 +1448,36 @@ prepareSlot(nixlAgent *agent,
     }
 
     if (!slot.req) {
-        nixl_xfer_dlist_t ld(GET_SEG_TYPE(true));
-        nixl_xfer_dlist_t rd(GET_SEG_TYPE(false));
-        prepareTransferDescriptors(ld, rd, slot.local_iov, slot.remote_iov);
-        nixl_status_t rc = agent->createXferReq(op, ld, rd, target, slot.req, &params);
+        nixl_status_t rc;
+        if (xferBenchConfig::prepped_xfer) {
+            if (!slot.prep_local_dlist) {
+                nixl_xfer_dlist_t ld(GET_SEG_TYPE(true));
+                nixl_xfer_dlist_t rd(GET_SEG_TYPE(false));
+                prepareTransferDescriptors(ld, rd, slot.local_iov, slot.remote_iov);
+                rc = agent->prepXferDlist(NIXL_INIT_AGENT, ld, slot.prep_local_dlist, &params);
+                if (rc != NIXL_SUCCESS) {
+                    return rc;
+                }
+                rc = agent->prepXferDlist(target, rd, slot.prep_remote_dlist, &params);
+                if (rc != NIXL_SUCCESS) {
+                    return rc;
+                }
+                slot.indices.resize(ld.descCount());
+                std::iota(slot.indices.begin(), slot.indices.end(), 0);
+            }
+            rc = agent->makeXferReq(op,
+                                    slot.prep_local_dlist,
+                                    slot.indices,
+                                    slot.prep_remote_dlist,
+                                    slot.indices,
+                                    slot.req,
+                                    &params);
+        } else {
+            nixl_xfer_dlist_t ld(GET_SEG_TYPE(true));
+            nixl_xfer_dlist_t rd(GET_SEG_TYPE(false));
+            prepareTransferDescriptors(ld, rd, slot.local_iov, slot.remote_iov);
+            rc = agent->createXferReq(op, ld, rd, target, slot.req, &params);
+        }
         if (rc != NIXL_SUCCESS) {
             return rc;
         }
@@ -1499,6 +1529,14 @@ cleanupSlots(nixlAgent *agent, nixlBackendH *backend_engine, std::vector<slotSta
             agent->releaseXferReq(slot.req);
             slot.req = nullptr;
         }
+        if (slot.prep_local_dlist) {
+            agent->releasedDlistH(slot.prep_local_dlist);
+            slot.prep_local_dlist = nullptr;
+        }
+        if (slot.prep_remote_dlist) {
+            agent->releasedDlistH(slot.prep_remote_dlist);
+            slot.prep_remote_dlist = nullptr;
+        }
         if (xferBenchConfig::reregister_mem && slot.registered) {
             deregisterIterationMem(agent, slot.local_iov, slot.remote_iov, backend_engine);
             slot.registered = false;