Add API for NPU Device Selection in the DML EP

include/onnxruntime/core/providers/dml/dml_provider_factory.h

@@ -30,6 +30,23 @@
 extern "C" {
 #endif
 
+enum class OrtDmlPerformancePreference {
+  Default,
+  HighPerformance,  // Default
+  LowPower
+};
+
+enum class OrtDmlDeviceFilter {
+  None,
+  Gpu,  // Default
+  Npu
+};
+
+struct OrtDmlDeviceOptions {
+  OrtDmlPerformancePreference p;
+  OrtDmlDeviceFilter f;
+};
+
 /**
  * [[deprecated]]
  * This export is deprecated.
@@ -82,6 +99,7 @@
   ORT_API2_STATUS(SessionOptionsAppendExecutionProvider_DML1, _In_ OrtSessionOptions* options,
                   _In_ IDMLDevice* dml_device, _In_ ID3D12CommandQueue* cmd_queue);
 
+
   /**
    * CreateGPUAllocationFromD3DResource
    * This API creates a DML EP resource based on a user-specified D3D12 resource.
@@ -94,11 +112,37 @@
    */
   ORT_API2_STATUS(FreeGPUAllocation, _In_ void* dml_resource);
 
+
   /**
    * GetD3D12ResourceFromAllocation
    * This API gets the D3D12 resource when an OrtValue has been allocated by the DML EP.
    */
   ORT_API2_STATUS(GetD3D12ResourceFromAllocation, _In_ OrtAllocator* provider, _In_ void* dml_resource, _Out_ ID3D12Resource** d3d_resource);
+
+
+    //ORT_API2_STATUS(FreeGPUAllocation_2, _In_ void* dml_resource);
+
+  //ORT_API2_STATUS(GetNPUDevice, _Out_ int* device_id);
+
+  //ORT_API2_STATUS(GetMLDevice, _In_ bool npu_first, _Out_ int* device_ids);
+
+  // enum class Ordering {
+  //     SAME,       // DXGI/DXCore as is today...
+  //     NONE,       // Take the first one
+  //     GPUS_FIRST,
+  //     NPUS_FIRST
+  // };
+
+  // Ordering foo[3][3] = {
+  //                // Default,                HighPerformance,      LowPower
+  //     /*None*/    { Ordering::GPUS_FIRST,   Ordering::GPUS_FIRST, Ordering::NPUS_FIRST       },
+  //     /*Gpu*/     { Ordering::SAME,         Ordering::SAME      , Ordering::SAME             },
+  //     /*Npu*/     { Ordering::NONE,         Ordering::NONE      , Ordering::NONE             },
+
+  // }
+
+  // null means default
+  ORT_API2_STATUS(OrtSessionOptionsAppendExecutionProvider_DML2, _In_ OrtSessionOptions* options, OrtDmlDeviceOptions* device_opts);
 };
 
 #ifdef __cplusplus

onnxruntime/core/providers/dml/dml_provider_factory.cc

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <dxcore.h>
+
 #include <DirectML.h>
 #ifndef _GAMING_XBOX
 #include <dxgi1_4.h>
@@ -172,6 +174,41 @@
   return CreateExecutionProviderFactory_DML(dml_device.Get(), cmd_queue.Get());
 }
 
+std::shared_ptr<IExecutionProviderFactory> DMLProviderFactoryCreator::CreateDXCore(ComPtr<IDXCoreAdapter> dxcore_device) {
+  // Create D3D12 Device from DXCore Adapter
+  ComPtr<ID3D12Device> d3d12_device;
+  ORT_THROW_IF_FAILED(D3D12CreateDevice(dxcore_device.Get(), D3D_FEATURE_LEVEL_11_0, IID_GRAPHICS_PPV_ARGS(d3d12_device.ReleaseAndGetAddressOf())));
+
+  // Create DML Device from D3D12 device
+  D3D12_COMMAND_QUEUE_DESC cmd_queue_desc = {};
+  cmd_queue_desc.Type = D3D12_COMMAND_LIST_TYPE_DIRECT;
+  cmd_queue_desc.Flags = D3D12_COMMAND_QUEUE_FLAG_DISABLE_GPU_TIMEOUT;
+
+  ComPtr<ID3D12CommandQueue> cmd_queue;
+  ORT_THROW_IF_FAILED(d3d12_device->CreateCommandQueue(&cmd_queue_desc, IID_GRAPHICS_PPV_ARGS(cmd_queue.ReleaseAndGetAddressOf())));
+
+  DML_CREATE_DEVICE_FLAGS flags = DML_CREATE_DEVICE_FLAG_NONE;
+
+  // In debug builds, enable the DML debug layer if the D3D12 debug layer is also enabled
+#if _DEBUG && !_GAMING_XBOX
+  ComPtr<ID3D12DebugDevice> debug_device;
+  (void)d3d12_device->QueryInterface(IID_PPV_ARGS(&debug_device));  // ignore failure
+  const bool is_d3d12_debug_layer_enabled = (debug_device != nullptr);
+
+  if (is_d3d12_debug_layer_enabled) {
+    flags |= DML_CREATE_DEVICE_FLAG_DEBUG;
+  }
+#endif
+
+  ComPtr<IDMLDevice> dml_device;
+  ORT_THROW_IF_FAILED(DMLCreateDevice1(d3d12_device.Get(),
+                                   flags,
+                                   DML_FEATURE_LEVEL_5_0,
+                                   IID_PPV_ARGS(&dml_device)));
+
+  return CreateExecutionProviderFactory_DML(dml_device.Get(), cmd_queue.Get());
+}
+
 }  // namespace onnxruntime
 
 // [[deprecated]]
@@ -196,6 +233,18 @@
   return nullptr;
 }
 
+bool isHardwareAdapter(ComPtr<IDXCoreAdapter> adapter) {
+    bool isHardware{ false };
+    THROW_IF_FAILED(adapter->GetProperty(
+        DXCoreAdapterProperty::IsHardware,
+        &isHardware));
+    return isHardware;
+}
+
+bool supportsGraphics(ComPtr<IDXCoreAdapter> adapter) {
+    return adapter->IsAttributeSupported(DXCORE_ADAPTER_ATTRIBUTE_D3D12_GRAPHICS);
+}
+
 ORT_API_STATUS_IMPL(CreateGPUAllocationFromD3DResource, _In_ ID3D12Resource* d3d_resource, _Out_ void** dml_resource) {
   API_IMPL_BEGIN
 #ifdef USE_DML
@@ -216,6 +265,128 @@
   API_IMPL_END
 }
 
+ORT_API_STATUS_IMPL(FreeGPUAllocation_2, _In_ void* ptr) {
+  API_IMPL_BEGIN
+#ifdef USE_DML
+  Dml::FreeGPUAllocation(ptr);
+#endif  // USE_DML
+  return nullptr;
+  API_IMPL_END
+}
+
+ORT_API_STATUS_IMPL(OrtSessionOptionsAppendExecutionProvider_DML2, _In_ OrtSessionOptions* options, OrtDmlDeviceOptions* device_opts) {
+API_IMPL_BEGIN
+  //options->provider_factories.push_back(onnxruntime::DMLProviderFactoryCreator::Create(device_id));
+
+    OrtDmlPerformancePreference p = device_opts->p;
+    OrtDmlDeviceFilter f = device_opts->f;
+
+    // Create DXCore Adapter Factory
+    ComPtr<IDXCoreAdapterFactory> adapterFactory;
+    ORT_THROW_IF_FAILED(::DXCoreCreateAdapterFactory(adapterFactory.GetAddressOf()));
+
+    // Get a list of all the adapters that support compute
+    ComPtr<IDXCoreAdapterList> d3D12CoreComputeAdapters;
+    GUID attributes[]{ DXCORE_ADAPTER_ATTRIBUTE_D3D12_CORE_COMPUTE };
+    ORT_THROW_IF_FAILED(
+        adapterFactory->CreateAdapterList(_countof(attributes),
+            attributes,
+            d3D12CoreComputeAdapters.GetAddressOf()));
+
+    const uint32_t count{ d3D12CoreComputeAdapters->GetAdapterCount() };
+    int compute_only_device_id = -1;
+
+    // will hold all the adapters in the order specified by the device options
+    ComPtr<IDXCoreAdapter>* ordered_adapters =
+        (ComPtr<IDXCoreAdapter>*) malloc(count * sizeof(ComPtr<IDXCoreAdapter>));
+    int ordered_adapter_count = 0;
+
+    // Used in the case that both GPU and NPU adapters are considered
+    ComPtr<IDXCoreAdapter>* gpu_adapters =
+        (ComPtr<IDXCoreAdapter>*) malloc(count * sizeof(ComPtr<IDXCoreAdapter>));
+    int gpu_adapter_count = 0;
+    ComPtr<IDXCoreAdapter>* npu_adapters =
+        (ComPtr<IDXCoreAdapter>*) malloc(count * sizeof(ComPtr<IDXCoreAdapter>));
+    int npu_adapter_count = 0;
+
+    // Iterate through all compute capable adapters
+    for (uint32_t i = 0; i < count; ++i)
+    {
+        ComPtr<IDXCoreAdapter> candidateAdapter;
+        ORT_THROW_IF_FAILED(
+            d3D12CoreComputeAdapters->GetAdapter(i, candidateAdapter.GetAddressOf()));
+
+        // Only considering hardware adapters
+        if (!isHardwareAdapter(candidateAdapter))
+            continue;
+
+        if (f == OrtDmlDeviceFilter::Gpu) // consider GPUs only
+        {
+            if (supportsGraphics(candidateAdapter)) {
+                ordered_adapters[ordered_adapter_count] = candidateAdapter;
+                ordered_adapter_count++;
+            }
+        }
+        else if(f == OrtDmlDeviceFilter::Npu) // consider NPUs only
+        {
+            if (!supportsGraphics(candidateAdapter)) {
+                ordered_adapters[ordered_adapter_count] = candidateAdapter;
+                ordered_adapter_count++;
+            }
+        }
+        else // consider both GPUs and NPUs
+        {
+            if (supportsGraphics(candidateAdapter)) {
+                gpu_adapters[gpu_adapter_count] = candidateAdapter;
+                gpu_adapter_count++;
+            } else {
+                npu_adapters[npu_adapter_count] = candidateAdapter;
+                npu_adapter_count++;
+            }
+        }
+    }
+
+    if (f == OrtDmlDeviceFilter::None) // considering both NPUs and GPUs
+    {
+        // order the adapters based on the performance preference
+        ComPtr<IDXCoreAdapter>* high_pri_adapters;
+        ComPtr<IDXCoreAdapter>* low_pri_adapters;
+        int num_high_pri_adapters;
+        int num_low_pri_adapters;
+        if (p == OrtDmlPerformancePreference::LowPower) // NPUs first
+        {
+            high_pri_adapters = npu_adapters;
+            num_high_pri_adapters = npu_adapter_count;
+            low_pri_adapters = gpu_adapters;
+            num_low_pri_adapters = gpu_adapter_count;
+        }
+        else // GPUs first
+        {
+            high_pri_adapters = gpu_adapters;
+            num_high_pri_adapters = gpu_adapter_count;
+            low_pri_adapters = npu_adapters;
+            num_low_pri_adapters = npu_adapter_count;
+        }
+
+        for (int i = 0; i < num_high_pri_adapters; i++) {
+            ordered_adapters[ordered_adapter_count] = high_pri_adapters[i];
+            ordered_adapter_count++;
+        }
+        for (int i = 0; i < num_low_pri_adapters; i++) {
+            ordered_adapters[ordered_adapter_count] = low_pri_adapters[i];
+            ordered_adapter_count++;
+        }
+    }
+
+    // check if ordered_adapter count is 0 (no adapters?)
+    ComPtr<IDXCoreAdapter> highest_pri_adapter = ordered_adapters[0];
+    // return the create function for a dxcore device
+    options->provider_factories.push_back(onnxruntime::DMLProviderFactoryCreator::CreateDXCore(
+        highest_pri_adapter));
+  return nullptr;
+API_IMPL_END
+}
+
 ORT_API_STATUS_IMPL(GetD3D12ResourceFromAllocation, _In_ OrtAllocator* ort_allocator, _In_ void* allocation, _Out_ ID3D12Resource** d3d_resource) {
   API_IMPL_BEGIN
 #ifdef USE_DML
@@ -238,7 +409,8 @@
   &OrtSessionOptionsAppendExecutionProviderEx_DML,
   &CreateGPUAllocationFromD3DResource,
   &FreeGPUAllocation,
-  &GetD3D12ResourceFromAllocation
+  &GetD3D12ResourceFromAllocation,
+  &OrtSessionOptionsAppendExecutionProvider_DML2
 };
 
 const OrtDmlApi* GetOrtDmlApi(_In_ uint32_t /*version*/) NO_EXCEPTION {

onnxruntime/core/providers/dml/dml_provider_factory_creator.h

@@ -10,11 +10,14 @@
 #include "core/providers/providers.h"
 #include "core/providers/dml/dml_provider_factory.h"
 
+#include <dxcore.h>
+
 namespace onnxruntime {
 
 struct DMLProviderFactoryCreator {
   static std::shared_ptr<IExecutionProviderFactory> Create(int device_id);
   static std::shared_ptr<IExecutionProviderFactory> Create(int device_id, bool skip_software_device_check);
+  static std::shared_ptr<IExecutionProviderFactory> CreateDXCore(Microsoft::WRL::ComPtr<IDXCoreAdapter> dxcore_device);
   static Microsoft::WRL::ComPtr<ID3D12Device> CreateD3D12Device(int device_id, bool skip_software_device_check);
 };
 }  // namespace onnxruntime