test: allow skipping device API put value

Author: 马函廷

.github/scripts/ci/run_unit_test.sh

@@ -62,12 +62,14 @@ build_suite() {
 
 run_device_api() {
   local suite_dir="$PROJECT_ROOT/test/unittest/device_api"
+  declare -p FLAGCX_CI_MPI_ENV_ARGS >/dev/null 2>&1 || FLAGCX_CI_MPI_ENV_ARGS=()
   local -a common_env=(
     -x FLAGCX_USE_HETERO_COMM=1
     -x FLAGCX_MEM_ENABLE=1
     -x FLAGCX_VMM_ENABLE=0
     -x FLAGCX_P2P_DISABLE=1
     -x LD_LIBRARY_PATH
+    "${FLAGCX_CI_MPI_ENV_ARGS[@]}"
   )
   local -a flags=(-b 1M -e 4M -f 2 -R 2)
 

.github/scripts/set_env/cuda.sh

@@ -25,6 +25,7 @@ FLAGCX_CI_TEST_MAKE_ARGS=(USE_NVIDIA=1)
 FLAGCX_CI_INTRA_NP=8
 FLAGCX_CI_NODE_NP=4
 FLAGCX_CI_RUNNER_NP=8
+FLAGCX_CI_MPI_ENV_ARGS=()
 export NP=8
 
 # Two logical four-GPU nodes on the eight-GPU CUDA runner.
@@ -48,6 +49,8 @@ flagcx_ci_configure_suite() {
     device_api)
       FLAGCX_CI_PROJECT_MAKE_ARGS+=(COMPILE_KERNEL=1 FORCE_DEFAULT_PATH=1)
       FLAGCX_CI_TEST_MAKE_ARGS+=(FORCE_DEFAULT_PATH=1)
+      export FLAGCX_SKIP_PUT_VALUE=1
+      FLAGCX_CI_MPI_ENV_ARGS+=(-x FLAGCX_SKIP_PUT_VALUE)
       ;;
   esac
 }

test/unittest/device_api/test_device_api.cpp

@@ -29,6 +29,7 @@
 #include "tools.h"
 
 #include <algorithm>
+#include <cstdlib>
 #include <cstdio>
 #include <cstring>
 #include <unistd.h>
@@ -117,6 +118,10 @@ int main(int argc, char *argv[]) {
   int numWarmupIters = args.getWarmupIters();
   int localRegister = args.getLocalRegister();
   uint64_t splitMask = args.getSplitMask();
+  const bool skipPutValue = [] {
+    const char *value = std::getenv("FLAGCX_SKIP_PUT_VALUE");
+    return value != nullptr && value[0] != '\0' && value[0] != '0';
+  }();
 
   if (stepFactor <= 1) {
     printf("Error: stepFactor must be > 1, got %d\n", stepFactor);
@@ -276,16 +281,23 @@ int main(int argc, char *argv[]) {
     MPI_Barrier(MPI_COMM_WORLD);
 
     // --- K3: PutValue ---
-    FLAGCXCHECK(devHandle->deviceMemset((char *)recvBuff + putValBase, 0,
-                                        (size_t)totalProcs * sizeof(uint64_t),
-                                        flagcxMemDevice, NULL));
-    FLAGCXCHECK(flagcxInterTestPutValue(recvMem, devComm, stream, putValBase));
-    FLAGCXCHECK(devHandle->streamSynchronize(stream));
-    FLAGCXCHECK(devHandle->deviceMemcpy(
-        (char *)hostBuff + putValBase, (char *)recvBuff + putValBase,
-        (size_t)totalProcs * sizeof(uint64_t), flagcxMemcpyDeviceToHost, NULL));
-    bool k3Ok = verifyPutValue(hostBuff, putValBase, totalProcs, proc);
-    printResult("K3 PutValue", k3Ok, proc);
+    if (skipPutValue) {
+      if (proc == 0 && color == 0)
+        printf("  %-30s SKIPPED\n", "K3 PutValue");
+    } else {
+      FLAGCXCHECK(devHandle->deviceMemset(
+          (char *)recvBuff + putValBase, 0,
+          (size_t)totalProcs * sizeof(uint64_t), flagcxMemDevice, NULL));
+      FLAGCXCHECK(
+          flagcxInterTestPutValue(recvMem, devComm, stream, putValBase));
+      FLAGCXCHECK(devHandle->streamSynchronize(stream));
+      FLAGCXCHECK(devHandle->deviceMemcpy(
+          (char *)hostBuff + putValBase, (char *)recvBuff + putValBase,
+          (size_t)totalProcs * sizeof(uint64_t), flagcxMemcpyDeviceToHost,
+          NULL));
+      bool k3Ok = verifyPutValue(hostBuff, putValBase, totalProcs, proc);
+      printResult("K3 PutValue", k3Ok, proc);
+    }
     MPI_Barrier(MPI_COMM_WORLD);
 
     // --- K4: Get ---