TheRock/rocm-systems/projects/hip-tests/catch/multiproc/hipDeviceTotalMemMproc.cc

/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#include <hip_test_common.hh>
#ifdef __linux__
#include <unistd.h>
#include <sys/wait.h>

/**
 * @addtogroup hipDeviceTotalMem hipDeviceTotalMem
 * @{
 * @ingroup DriverTest
 */

#define MAX_SIZE 30
#define VISIBLE_DEVICE 0

/**
 * Fetches Gpu device count
 */
static void getDeviceCount(int* pdevCnt) {
  int fd[2], val = 0;
  pid_t childpid;

  // create pipe descriptors
  pipe(fd);

  // disable visible_devices env from shell
#ifdef __HIP_PLATFORM_NVIDIA__
  unsetenv("CUDA_VISIBLE_DEVICES");
#else
  unsetenv("ROCR_VISIBLE_DEVICES");
  unsetenv("HIP_VISIBLE_DEVICES");
#endif

  childpid = fork();

  if (childpid > 0) {  // Parent
    close(fd[1]);
    // parent will wait to read the device cnt
    read(fd[0], &val, sizeof(val));

    // close the read-descriptor
    close(fd[0]);

    // wait for child exit
    wait(NULL);

    *pdevCnt = val;
  } else if (!childpid) {  // Child
    int devCnt = 1;
    // writing only, no need for read-descriptor
    close(fd[0]);

    HIP_CHECK(hipGetDeviceCount(&devCnt));
    // send the value on the write-descriptor:
    write(fd[1], &devCnt, sizeof(devCnt));

    // close the write descriptor:
    close(fd[1]);
    exit(0);
  } else {  // failure
    *pdevCnt = 0;
    return;
  }
}


/**
 * Func tries to fetch total memory of masked devices and returns pass/fail.
 */
static bool getTotalMemoryOfMaskedDevices(int actualNumGPUs) {
  bool testResult = true;
  int fd[2];

  pipe(fd);
  pid_t cPid;
  cPid = fork();
  if (cPid == 0) {  // child
    hipError_t err;
    char visibleDeviceString[MAX_SIZE] = {};
    snprintf(visibleDeviceString, MAX_SIZE, "%d", VISIBLE_DEVICE);

    // disable visible_devices env from shell
#ifdef __HIP_PLATFORM_NVIDIA__
    unsetenv("CUDA_VISIBLE_DEVICES");
    setenv("CUDA_VISIBLE_DEVICES", visibleDeviceString, 1);
    HIP_CHECK(hipInit(0));
#else
    unsetenv("ROCR_VISIBLE_DEVICES");
    unsetenv("HIP_VISIBLE_DEVICES");
    setenv("ROCR_VISIBLE_DEVICES", visibleDeviceString, 1);
    setenv("HIP_VISIBLE_DEVICES", visibleDeviceString, 1);
#endif

    for (int count = 1; count < actualNumGPUs; count++) {
      size_t totMem;
      err = hipDeviceTotalMem(&totMem, count);
      if (err == hipSuccess) {
        testResult &= false;
      } else {
        printf("hipDeviceTotalMem: Error Code Returned: '%s'(%d)\n", hipGetErrorString(err), err);
      }
    }
    close(fd[0]);
    printf("testResult = %d \n", testResult);
    write(fd[1], &testResult, sizeof(testResult));
    close(fd[1]);
    exit(0);

  } else if (cPid > 0) {  // parent
    close(fd[1]);
    read(fd[0], &testResult, sizeof(testResult));
    close(fd[0]);
    wait(NULL);

  } else {
    printf("fork() failed\n");
    HIP_ASSERT(false);
  }

  return testResult;
}


/**
 * Test Description
 * ------------------------
 *  - Check that total memory is returned correctly when
 *    the devices are masked.
 * Test source
 * ------------------------
 *  - unit/multiproc/hipDeviceTotalMemMproc.cc
 * Test requirements
 * ------------------------
 *  - Multi-device test
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Unit_hipDeviceTotalMem_MaskedDevices) {
  int count = -1;
  constexpr int ReqGPUs = 2;
  bool ret;

  getDeviceCount(&count);

  if (count >= ReqGPUs) {
    ret = getTotalMemoryOfMaskedDevices(count);
    REQUIRE(ret == true);
  } else {
    HIP_SKIP_TEST(HipTest::SkipReason::kFewerThanTwoGpus);
  }
}

/**
 * End doxygen group hipDeviceTotalMem.
 * @}
 */

#endif