From 30d07163adcf1ce8bf6172e5f2be5f0d812d2d22 Mon Sep 17 00:00:00 2001
From: Krzysztof Bogacki <krzysztof.bogacki@leancode.pl>
Date: Sat, 25 Jun 2022 22:00:45 +0200
Subject: tests: Check behavior of NvAPI_GPU_GetGpuCoreCount

---
 tests/nvapi_sysinfo.cpp     | 21 +++++++++++++++++++++
 tests/nvapi_sysinfo_mocks.h |  1 +
 tests/nvapi_system.cpp      | 10 ++++++++++
 3 files changed, 32 insertions(+)

diff --git a/tests/nvapi_sysinfo.cpp b/tests/nvapi_sysinfo.cpp
index 87b54bb..fd990b9 100644
--- a/tests/nvapi_sysinfo.cpp
+++ b/tests/nvapi_sysinfo.cpp
@@ -747,6 +747,23 @@ TEST_CASE("Sysinfo methods succeed", "[.sysinfo]") {
         ALLOW_CALL(*nvml, DeviceGetHandleByPciBusId_v2(_, _))
             .RETURN(NVML_SUCCESS);
 
+        SECTION("GetGpuCoreCount returns OK") {
+            auto cores = 1536U;
+            ALLOW_CALL(*nvml, DeviceGetNumGpuCores(_, _)) // NOLINT(bugprone-use-after-move)
+                .LR_SIDE_EFFECT(*_2 = cores)
+                .RETURN(NVML_SUCCESS);
+
+            SetupResourceFactory(std::move(dxgiFactory), std::move(vulkan), std::move(nvml), std::move(lfx));
+            REQUIRE(NvAPI_Initialize() == NVAPI_OK);
+
+            NvPhysicalGpuHandle handle;
+            REQUIRE(NvAPI_SYS_GetPhysicalGpuFromDisplayId(0, &handle) == NVAPI_OK);
+
+            NvU32 count;
+            REQUIRE(NvAPI_GPU_GetGpuCoreCount(handle, &count) == NVAPI_OK);
+            REQUIRE(count == cores);
+        }
+
         SECTION("GetPCIIdentifiers returns OK and has subsystem ID when NVML is available") {
             auto id = 0x88161043;
             ALLOW_CALL(*nvml, DeviceGetPciInfo_v3(_, _)) // NOLINT(bugprone-use-after-move)
@@ -1141,6 +1158,8 @@ TEST_CASE("Sysinfo methods succeed", "[.sysinfo]") {
         }
 
         SECTION("NVML depending methods return no-implementation when NVML is not available") {
+            NvU32 cores;
+            REQUIRE(NvAPI_GPU_GetGpuCoreCount(handle, &cores) == NVAPI_NO_IMPLEMENTATION);
             NV_GPU_DYNAMIC_PSTATES_INFO_EX info;
             info.version = NV_GPU_DYNAMIC_PSTATES_INFO_EX_VER;
             REQUIRE(NvAPI_GPU_GetDynamicPstatesInfoEx(handle, &info) == NVAPI_NO_IMPLEMENTATION);
@@ -1177,6 +1196,8 @@ TEST_CASE("Sysinfo methods succeed", "[.sysinfo]") {
         }
 
         SECTION("NVML depending methods return handle-invalidated when NVML is available but without suitable adapter") {
+            NvU32 cores;
+            REQUIRE(NvAPI_GPU_GetGpuCoreCount(handle, &cores) == NVAPI_HANDLE_INVALIDATED);
             NV_GPU_DYNAMIC_PSTATES_INFO_EX info;
             info.version = NV_GPU_DYNAMIC_PSTATES_INFO_EX_VER;
             REQUIRE(NvAPI_GPU_GetDynamicPstatesInfoEx(handle, &info) == NVAPI_HANDLE_INVALIDATED);
diff --git a/tests/nvapi_sysinfo_mocks.h b/tests/nvapi_sysinfo_mocks.h
index 5cf900a..2c85ee9 100644
--- a/tests/nvapi_sysinfo_mocks.h
+++ b/tests/nvapi_sysinfo_mocks.h
@@ -79,6 +79,7 @@ class NvmlMock : public trompeloeil::mock_interface<dxvk::Nvml> {
     IMPLEMENT_CONST_MOCK2(DeviceGetPerformanceState);
     IMPLEMENT_CONST_MOCK2(DeviceGetUtilizationRates);
     IMPLEMENT_CONST_MOCK3(DeviceGetVbiosVersion);
+    IMPLEMENT_CONST_MOCK2(DeviceGetNumGpuCores);
     IMPLEMENT_CONST_MOCK2(DeviceGetBusType);
     IMPLEMENT_CONST_MOCK2(DeviceGetDynamicPstatesInfo);
 };
diff --git a/tests/nvapi_system.cpp b/tests/nvapi_system.cpp
index ca4daa9..28d5f11 100644
--- a/tests/nvapi_system.cpp
+++ b/tests/nvapi_system.cpp
@@ -95,6 +95,7 @@ TEST_CASE("Sysinfo methods succeed against local system", "[system]") {
     GETNVAPIPROCADDR(GPU_GetPhysicalFrameBufferSize);
     GETNVAPIPROCADDR(GPU_GetAdapterIdFromPhysicalGpu);
     GETNVAPIPROCADDR(GPU_GetArchInfo);
+    GETNVAPIPROCADDR(GPU_GetGpuCoreCount);
     GETNVAPIPROCADDR(GPU_CudaEnumComputeCapableGpus);
     GETNVAPIPROCADDR(GPU_GetVbiosVersionString);
     GETNVAPIPROCADDR(GPU_GetDynamicPstatesInfoEx);
@@ -118,6 +119,7 @@ TEST_CASE("Sysinfo methods succeed against local system", "[system]") {
     REQUIRE(nvAPI_GPU_GetPhysicalFrameBufferSize);
     REQUIRE(nvAPI_GPU_GetAdapterIdFromPhysicalGpu);
     REQUIRE(nvAPI_GPU_GetArchInfo);
+    REQUIRE(nvAPI_GPU_GetGpuCoreCount);
     REQUIRE(nvAPI_GPU_CudaEnumComputeCapableGpus);
     REQUIRE(nvAPI_GPU_GetVbiosVersionString);
     REQUIRE(nvAPI_GPU_GetDynamicPstatesInfoEx);
@@ -210,6 +212,14 @@ TEST_CASE("Sysinfo methods succeed against local system", "[system]") {
         else
             std::cout << "N/A" << std::endl;
 
+        NvU32 gpuCoreCount;
+        result = nvAPI_GPU_GetGpuCoreCount(handle, &gpuCoreCount);
+        std::cout << "    GPU core count:             ";
+        if (result == NVAPI_OK)
+            std::cout << std::dec << gpuCoreCount << std::endl;
+        else
+            std::cout << "N/A" << std::endl;
+
         auto computeGpusIt = std::find_if(
             std::begin(computeGpus),
             std::end(computeGpus),
-- 
cgit v1.2.3