Add code for creating Generic ML devices to NPU starter sample (#599)

* Add flag to force generic ML only device * Improve naming of device iterated on in NPU sample
microsoft · Jun 25, 2024 · 46f3e14 · 46f3e14
1 parent 372a622
commit 46f3e14
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/Samples/DirectMLNpuInference/main.cpp b/Samples/DirectMLNpuInference/main.cpp
@@ -16,6 +16,8 @@ using Microsoft::WRL::ComPtr;
 void InitializeDirectML(ID3D12Device1** d3dDeviceOut, ID3D12CommandQueue** commandQueueOut, IDMLDevice** dmlDeviceOut) {
     // Whether to skip adapters which support Graphics in order to target NPU for testing
     bool forceComputeOnlyDevice = true;
+    bool forceGenericMLDevice = false;
+
     ComPtr<IDXCoreAdapterFactory> factory;
     HMODULE dxCoreModule = LoadLibraryW(L"DXCore.dll");
     if (dxCoreModule)
@@ -37,11 +39,23 @@ void InitializeDirectML(ID3D12Device1** d3dDeviceOut, ID3D12CommandQueue** comma
         THROW_IF_FAILED(factory->CreateAdapterList(ARRAYSIZE(dxGUIDs), dxGUIDs, IID_PPV_ARGS(&adapterList)));
         for (uint32_t i = 0, adapterCount = adapterList->GetAdapterCount(); i < adapterCount; i++)
         {
-            ComPtr<IDXCoreAdapter> nextGpuAdapter;
-            THROW_IF_FAILED(adapterList->GetAdapter(static_cast<uint32_t>(i), IID_PPV_ARGS(&nextGpuAdapter)));
-            if (!forceComputeOnlyDevice || !nextGpuAdapter->IsAttributeSupported(DXCORE_ADAPTER_ATTRIBUTE_D3D12_GRAPHICS))
+            ComPtr<IDXCoreAdapter> currentGpuAdapter;
+            THROW_IF_FAILED(adapterList->GetAdapter(static_cast<uint32_t>(i), IID_PPV_ARGS(&currentGpuAdapter)));
+
+            if (!forceComputeOnlyDevice && !forceGenericMLDevice)
+            {
+                // No device restrictions
+                adapter = std::move(currentGpuAdapter);
+                break;
+            }
+            else if (forceComputeOnlyDevice && currentGpuAdapter->IsAttributeSupported(DXCORE_ADAPTER_ATTRIBUTE_D3D12_CORE_COMPUTE))
+            {
+                adapter = std::move(currentGpuAdapter);
+                break;
+            }
+            else if (forceGenericMLDevice && currentGpuAdapter->IsAttributeSupported(DXCORE_ADAPTER_ATTRIBUTE_D3D12_GENERIC_ML))
             {
-                adapter = std::move(nextGpuAdapter);
+                adapter = std::move(currentGpuAdapter);
                 break;
             }
         }