nvapi: Add NvAPI_GetLogicalGPUFromDisplay

Also add a few missing nullptr checks.
jp7677 · May 27, 2024 · 2b67af4 · 2b67af4
1 parent b7e018f
commit 2b67af4
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 6 deletions.
diff --git a/src/nvapi.cpp b/src/nvapi.cpp
@@ -89,6 +89,9 @@ extern "C" {
         if (nvapiAdapterRegistry == nullptr)
             return ApiNotInitialized(n);
 
+        if (hPhysicalGpu == nullptr)
+            return InvalidArgument(n);
+
         NvapiAdapter* adapter = nullptr;
         for (auto i = 0U; i < nvapiAdapterRegistry->GetAdapterCount(); i++)
             if (nvapiAdapterRegistry->GetAdapter(i)->GetBoardId() == gpuId)
@@ -130,7 +133,7 @@ extern "C" {
         if (nvapiAdapterRegistry == nullptr)
             return ApiNotInitialized(n);
 
-        if (hPhysicalGPU == nullptr)
+        if (hPhysicalGPU == nullptr || pLogicalGPU == nullptr)
             return InvalidArgument(n);
 
         auto adapter = reinterpret_cast<NvapiAdapter*>(hPhysicalGPU);
@@ -142,13 +145,31 @@ extern "C" {
         return Ok(n);
     }
 
+    NvAPI_Status __cdecl NvAPI_GetLogicalGPUFromDisplay(NvDisplayHandle hNvDisp, NvLogicalGpuHandle* pLogicalGPU) {
+        constexpr auto n = __func__;
+
+        if (nvapiAdapterRegistry == nullptr)
+            return ApiNotInitialized(n);
+
+        if (hNvDisp == nullptr || pLogicalGPU == nullptr)
+            return InvalidArgument(n);
+
+        auto output = reinterpret_cast<NvapiOutput*>(hNvDisp);
+        if (!nvapiAdapterRegistry->IsOutput(output))
+            return ExpectedDisplayHandle(n);
+
+        *pLogicalGPU = reinterpret_cast<NvLogicalGpuHandle>(output->GetParent());
+
+        return Ok(n);
+    }
+
     NvAPI_Status __cdecl NvAPI_GetPhysicalGPUsFromLogicalGPU(NvLogicalGpuHandle hLogicalGPU, NvPhysicalGpuHandle hPhysicalGPU[NVAPI_MAX_PHYSICAL_GPUS], NvU32* pGpuCount) {
         constexpr auto n = __func__;
 
         if (nvapiAdapterRegistry == nullptr)
             return ApiNotInitialized(n);
 
-        if (hLogicalGPU == nullptr)
+        if (hLogicalGPU == nullptr || hPhysicalGPU == nullptr || pGpuCount == nullptr)
             return InvalidArgument(n);
 
         auto adapter = reinterpret_cast<NvapiAdapter*>(hLogicalGPU);
@@ -209,7 +230,7 @@ extern "C" {
         if (nvapiAdapterRegistry == nullptr)
             return ApiNotInitialized(n);
 
-        if (NvDispHandle == nullptr)
+        if (NvDispHandle == nullptr || szDisplayName == nullptr)
             return InvalidArgument(n);
 
         auto output = reinterpret_cast<NvapiOutput*>(NvDispHandle);
@@ -227,7 +248,7 @@ extern "C" {
         if (nvapiAdapterRegistry == nullptr)
             return ApiNotInitialized(n);
 
-        if (szDisplayName == nullptr)
+        if (szDisplayName == nullptr || pNvDispHandle == nullptr)
             return InvalidArgument(n);
 
         auto output = nvapiAdapterRegistry->FindOutput(szDisplayName);

diff --git a/src/nvapi_interface.cpp b/src/nvapi_interface.cpp
@@ -130,6 +130,7 @@ extern "C" {
         INSERT_AND_RETURN_WHEN_EQUALS(NvAPI_GetDisplayDriverVersion)
         INSERT_AND_RETURN_WHEN_EQUALS(NvAPI_GetPhysicalGPUsFromDisplay)
         INSERT_AND_RETURN_WHEN_EQUALS(NvAPI_GetLogicalGPUFromPhysicalGPU)
+        INSERT_AND_RETURN_WHEN_EQUALS(NvAPI_GetLogicalGPUFromDisplay)
         INSERT_AND_RETURN_WHEN_EQUALS(NvAPI_GetPhysicalGPUsFromLogicalGPU)
         INSERT_AND_RETURN_WHEN_EQUALS(NvAPI_EnumNvidiaDisplayHandle)
         INSERT_AND_RETURN_WHEN_EQUALS(NvAPI_EnumNvidiaUnAttachedDisplayHandle)

diff --git a/tests/nvapi_sysinfo_topo.cpp b/tests/nvapi_sysinfo_topo.cpp
@@ -106,6 +106,35 @@ TEST_CASE("Topology methods succeed", "[.sysinfo-topo]") {
         REQUIRE(logicalhandle == reinterpret_cast<NvLogicalGpuHandle>(handles[1]));
     }
 
+    SECTION("GetLogicalGPUFromDisplay succeeds") {
+        NvLogicalGpuHandle logicalhandles[NVAPI_MAX_LOGICAL_GPUS]{};
+        NvU32 count;
+        REQUIRE(NvAPI_EnumLogicalGPUs(logicalhandles, &count) == NVAPI_OK);
+
+        NvDisplayHandle handle1 = nullptr;
+        NvDisplayHandle handle2 = nullptr;
+        NvDisplayHandle handle3 = nullptr;
+        REQUIRE(NvAPI_EnumNvidiaDisplayHandle(0U, &handle1) == NVAPI_OK);
+        REQUIRE(NvAPI_EnumNvidiaDisplayHandle(1U, &handle2) == NVAPI_OK);
+        REQUIRE(NvAPI_EnumNvidiaDisplayHandle(2U, &handle3) == NVAPI_OK);
+
+        NvLogicalGpuHandle logicalhandle1;
+        REQUIRE(NvAPI_GetLogicalGPUFromDisplay(handle1, &logicalhandle1) == NVAPI_OK);
+        REQUIRE(logicalhandle1 == reinterpret_cast<NvLogicalGpuHandle>(logicalhandles[0]));
+
+        NvLogicalGpuHandle logicalhandle2;
+        REQUIRE(NvAPI_GetLogicalGPUFromDisplay(handle2, &logicalhandle2) == NVAPI_OK);
+        REQUIRE(logicalhandle2 == reinterpret_cast<NvLogicalGpuHandle>(logicalhandles[0]));
+
+        NvLogicalGpuHandle logicalhandle3;
+        REQUIRE(NvAPI_GetLogicalGPUFromDisplay(handle3, &logicalhandle3) == NVAPI_OK);
+        REQUIRE(logicalhandle3 == reinterpret_cast<NvLogicalGpuHandle>(logicalhandles[1]));
+
+        auto invalid = reinterpret_cast<NvDisplayHandle>(0x1);
+        NvLogicalGpuHandle logicalhandle4;
+        REQUIRE(NvAPI_GetLogicalGPUFromDisplay(invalid, &logicalhandle4) == NVAPI_EXPECTED_DISPLAY_HANDLE);
+    }
+
     SECTION("GetPhysicalGPUsFromLogicalGPU succeeds") {
         NvLogicalGpuHandle handles[NVAPI_MAX_LOGICAL_GPUS]{};
         NvU32 count;
@@ -208,9 +237,9 @@ TEST_CASE("Topology methods succeed", "[.sysinfo-topo]") {
         REQUIRE(NvAPI_GetAssociatedNvidiaDisplayName(handle3, name3) == NVAPI_OK);
         REQUIRE_THAT(name3, Equals("Output3"));
 
+        auto invalid = reinterpret_cast<NvDisplayHandle>(0x1);
         NvAPI_ShortString name4;
-        NvDisplayHandle handle4 = nullptr;
-        REQUIRE(NvAPI_GetAssociatedNvidiaDisplayName(handle4, name4) == NVAPI_INVALID_ARGUMENT);
+        REQUIRE(NvAPI_GetAssociatedNvidiaDisplayName(invalid, name4) == NVAPI_EXPECTED_DISPLAY_HANDLE);
     }
 
     SECTION("GetAssociatedNvidiaDisplayHandle succeeds") {