diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp new file mode 100644 index 0000000000..0f52d4f010 --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp @@ -0,0 +1,499 @@ +#pragma once + +#include "DmlGraphFusionHelper.h" + + +namespace Dml +{ +namespace DmlGraphFusionHelper +{ + Microsoft::WRL::ComPtr + CreateResource( + const ExecutionProviderImpl* provider, + const std::byte* tensorPtr, + size_t tensorByteSize) + { + Microsoft::WRL::ComPtr buffer; + + D3D12_HEAP_PROPERTIES heapProperties = { + D3D12_HEAP_TYPE_DEFAULT, D3D12_CPU_PAGE_PROPERTY_UNKNOWN, D3D12_MEMORY_POOL_UNKNOWN, 0, 0}; + + D3D12_RESOURCE_DESC resourceDesc = {D3D12_RESOURCE_DIMENSION_BUFFER, + 0, + static_cast((tensorByteSize + 3) & ~3), + 1, + 1, + 1, + DXGI_FORMAT_UNKNOWN, + {1, 0}, + D3D12_TEXTURE_LAYOUT_ROW_MAJOR, + D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS}; + + Microsoft::WRL::ComPtr d3dDevice; + ORT_THROW_IF_FAILED(provider->GetD3DDevice(d3dDevice.GetAddressOf())); + + ORT_THROW_IF_FAILED(d3dDevice->CreateCommittedResource( + &heapProperties, + D3D12_HEAP_FLAG_NONE, + &resourceDesc, + D3D12_RESOURCE_STATE_UNORDERED_ACCESS, + nullptr, + IID_GRAPHICS_PPV_ARGS(buffer.GetAddressOf()))); + + ORT_THROW_IF_FAILED(provider->UploadToResource(buffer.Get(), tensorPtr, tensorByteSize)); + + return buffer; + } + + Microsoft::WRL::ComPtr + CreateCpuResource( + const ExecutionProviderImpl* provider, + const std::byte* tensorPtr, + size_t tensorByteSize) + { + Microsoft::WRL::ComPtr buffer; + + D3D12_HEAP_PROPERTIES heapProperties = { + D3D12_HEAP_TYPE_CUSTOM, D3D12_CPU_PAGE_PROPERTY_WRITE_COMBINE, D3D12_MEMORY_POOL_L0, 0, 0}; + + D3D12_RESOURCE_DESC resourceDesc = {D3D12_RESOURCE_DIMENSION_BUFFER, + 0, + static_cast((tensorByteSize + 3) & ~3), + 1, + 1, + 1, + DXGI_FORMAT_UNKNOWN, + {1, 0}, + D3D12_TEXTURE_LAYOUT_ROW_MAJOR, + D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS}; + + Microsoft::WRL::ComPtr d3dDevice; + ORT_THROW_IF_FAILED(provider->GetD3DDevice(d3dDevice.GetAddressOf())); + + ORT_THROW_IF_FAILED(d3dDevice->CreateCommittedResource( + &heapProperties, + D3D12_HEAP_FLAG_NONE, + &resourceDesc, + D3D12_RESOURCE_STATE_UNORDERED_ACCESS, + nullptr, + IID_GRAPHICS_PPV_ARGS(buffer.GetAddressOf()))); + + // Map the buffer and copy the data + void* bufferData = nullptr; + D3D12_RANGE range = {0, tensorByteSize}; + ORT_THROW_IF_FAILED(buffer->Map(0, &range, &bufferData)); + memcpy(bufferData, tensorPtr, tensorByteSize); + buffer->Unmap(0, &range); + + return buffer; + } + + void UnwrapTensor( + Windows::AI::MachineLearning::Adapter::IWinmlExecutionProvider* winmlProvider, + const onnxruntime::Tensor* tensor, + ID3D12Resource** resource, + uint64_t* allocId) + { + IUnknown* allocationUnk = static_cast(const_cast(tensor->DataRaw())); + Microsoft::WRL::ComPtr resourceUnk; + winmlProvider->GetABIDataInterface(false, allocationUnk, &resourceUnk); + + *allocId = winmlProvider->TryGetPooledAllocationId(allocationUnk, 0); + + ORT_THROW_IF_FAILED(resourceUnk->QueryInterface(resource)); + } + + void ProcessInputData( + const ExecutionProviderImpl* providerImpl, + const std::vector& isInputsUploadedByDmlEP, + std::vector& inputEdges, + const gsl::span subGraphInputArgNames, + const std::unordered_map>& initializerNameToInitializerMap, + onnxruntime::Graph& graph, + _Out_ std::vector& inputsUsed, + _Inout_ std::vector& initInputBindings, + _Inout_ std::vector>& initInputResources, + _Inout_ std::vector>& nonOwnedGraphInputsFromInitializers, + _Inout_ std::vector>& initializeResourceRefs, + _Inout_opt_ std::vector>* inputRawData) + { + + const uint32_t fusedNodeInputCount = gsl::narrow_cast(subGraphInputArgNames.size()); + + // Determine the last input which uses an initializer, so initializers can be freed incrementally + // while processing each input in order. + std::map initializerToLastInputIndexMap; + for (uint32_t i = 0; i < fusedNodeInputCount; i++) + { + auto iter = initializerNameToInitializerMap.find(subGraphInputArgNames[i]); + if (iter != initializerNameToInitializerMap.end()) { + initializerToLastInputIndexMap[iter->second.first] = i; + } + } + + // Walk through each graph edge and mark used inputs + inputsUsed.assign(fusedNodeInputCount, false); + for (const DML_INPUT_GRAPH_EDGE_DESC& edge : inputEdges) { + inputsUsed[edge.GraphInputIndex] = true; + } + + for (uint32_t i = 0; i < initInputBindings.size(); i++) + { + bool isInitializerAlreadyRemoved = false; + // If the input isn't actually used by the graph, nothing ever needs to be bound (either for + // initialization or execution). So just throw away the transferred initializer and skip this input. + if (!inputsUsed[i]) + { + auto iter = initializerNameToInitializerMap.find(subGraphInputArgNames[i]); + if(iter != initializerNameToInitializerMap.end() && iter->second.second) + { + graph.RemoveInitializedTensor(subGraphInputArgNames[i]); + } + + if (inputRawData) + { + inputRawData->push_back(std::vector()); + } + + continue; + } + + // Look for the initializer among those transferred from the graph during partitioning + auto iter = initializerNameToInitializerMap.find(subGraphInputArgNames[i]); + if (iter != initializerNameToInitializerMap.end()) + { + std::byte* tensorPtr = nullptr; + size_t tensorByteSize = 0; + std::unique_ptr unpackedTensor; + + //auto& initializer = iter->second; + auto* initializer = iter->second.first; + + // The tensor may be stored as raw data or in typed fields. + if (initializer->has_raw_data()) + { + tensorPtr = (std::byte*)(initializer->raw_data().c_str()); + tensorByteSize = initializer->raw_data().size(); + } + else + { + std::tie(unpackedTensor, tensorByteSize) = Windows::AI::MachineLearning::Adapter::UnpackTensor(*initializer); + tensorPtr = unpackedTensor.get(); + + // Free the initializer if this is the last usage of it. + if (initializerToLastInputIndexMap[initializer] == i) + { + if (iter->second.second) + { + graph.RemoveInitializedTensor(subGraphInputArgNames[i]); + isInitializerAlreadyRemoved = true; + } + } + } + + // Tensor sizes in DML must be a multiple of 4 bytes large. + tensorByteSize = AlignToPow2(tensorByteSize, 4); + + if (inputRawData) + { + inputRawData->push_back(std::vector(tensorPtr, tensorPtr + tensorByteSize)); + } + + if (!isInputsUploadedByDmlEP[i]) + { + // Store the resource to use during execution + ComPtr defaultBuffer = CreateResource(providerImpl, tensorPtr, tensorByteSize); + nonOwnedGraphInputsFromInitializers[i] = defaultBuffer; + initializeResourceRefs.push_back(std::move(defaultBuffer)); + } + else + { + ComPtr initializeInputBuffer; + + // D3D_FEATURE_LEVEL_1_0_CORE doesn't support Custom heaps + if (providerImpl->IsMcdmDevice()) + { + initializeInputBuffer = CreateResource(providerImpl, tensorPtr, tensorByteSize); + } + else + { + initializeInputBuffer = CreateCpuResource(providerImpl, tensorPtr, tensorByteSize); + } + + // Set the binding for operator initialization to the buffer + initInputBindings[i].Buffer = initializeInputBuffer.Get(); + initInputBindings[i].SizeInBytes = tensorByteSize; + initializeResourceRefs.push_back(std::move(initializeInputBuffer)); + } + + // Free the initializer if this is the last usage of it. + if (!isInitializerAlreadyRemoved && initializerToLastInputIndexMap[initializer] == i) + { + if (iter->second.second) + { + graph.RemoveInitializedTensor(subGraphInputArgNames[i]); + } + } + } + else if (inputRawData) + { + inputRawData->push_back(std::vector()); + } + } + } + + std::unordered_map> + GetInitializerToPartitionMap( + const onnxruntime::GraphViewer& graph, + gsl::span> partitions + ) + { + std::unordered_map> initializerPartitionMap; + for (uint32_t partitionIndex = 0; partitionIndex < gsl::narrow_cast(partitions.size()); ++partitionIndex) + { + auto& partition = partitions[partitionIndex]; + + // Skip partitions which have been merged into other partitions + if (partition->GetRootMergedPartition() != partition.get()) + { + continue; + } + + for (const std::string& input : partition->GetInputs()) + { + const onnx::TensorProto* tensor = nullptr; + if (graph.GetInitializedTensor(input, tensor)) + { + initializerPartitionMap[tensor].push_back(partitionIndex); + } + } + } + + return initializerPartitionMap; + } + + void ConvertGraphDesc( + const Dml::GraphDescBuilder::GraphDesc& graphDesc, + _Out_ DML_GRAPH_DESC& dmlGraphDesc, + const uint32_t inputCount, + const uint32_t outputCount, + _Inout_ std::vector& dmlOperatorGraphNodes, + _Inout_ std::vector& dmlGraphNodes, + _Inout_ std::vector& dmlInputEdges, + _Inout_ std::vector& dmlOutputEdges, + _Inout_ std::vector& dmlIntermediateEdges) + { + for (size_t i = 0; i < graphDesc.nodes.size(); ++i) + { + dmlOperatorGraphNodes[i] = DML_OPERATOR_GRAPH_NODE_DESC{graphDesc.nodes[i].op.Get()}; + dmlGraphNodes[i] = DML_GRAPH_NODE_DESC{DML_GRAPH_NODE_TYPE_OPERATOR, &dmlOperatorGraphNodes[i]}; + } + + for (size_t i = 0; i < graphDesc.inputEdges.size(); ++i) + { + dmlInputEdges[i] = DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_INPUT, &graphDesc.inputEdges[i]}; + } + + for (size_t i = 0; i < graphDesc.outputEdges.size(); ++i) + { + dmlOutputEdges[i] = DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_OUTPUT, &graphDesc.outputEdges[i]}; + } + + for (size_t i = 0; i < graphDesc.intermediateEdges.size(); ++i) + { + dmlIntermediateEdges[i] = + DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_INTERMEDIATE, &graphDesc.intermediateEdges[i]}; + } + + dmlGraphDesc.InputCount = inputCount; + dmlGraphDesc.OutputCount = outputCount; + dmlGraphDesc.NodeCount = gsl::narrow_cast(dmlGraphNodes.size()); + dmlGraphDesc.Nodes = dmlGraphNodes.data(); + dmlGraphDesc.InputEdgeCount = gsl::narrow_cast(dmlInputEdges.size()); + dmlGraphDesc.InputEdges = dmlInputEdges.data(); + dmlGraphDesc.OutputEdgeCount = gsl::narrow_cast(dmlOutputEdges.size()); + dmlGraphDesc.OutputEdges = dmlOutputEdges.data(); + dmlGraphDesc.IntermediateEdgeCount = gsl::narrow_cast(dmlIntermediateEdges.size()); + dmlGraphDesc.IntermediateEdges = dmlIntermediateEdges.data(); + } + + void CreateIDmlCompiledOperatorAndRegisterKernel( + onnxruntime::Graph& graph, + const onnxruntime::IndexedSubGraph& indexedSubGraph, + const onnxruntime::Node& fusedNode, + const std::unordered_map& partitionNodePropsMap, + const std::unordered_map>& initializerNameToInitializerMap, + const ExecutionProviderImpl* providerImpl, + onnxruntime::KernelRegistry* registryForPartitionKernels) + { + // convert partitionONNXGraph into DML EP GraphDesc + const uint32_t fusedNodeInputCount = gsl::narrow_cast(indexedSubGraph.GetMetaDef()->inputs.size()); + const uint32_t fusedNodeOutputCount = gsl::narrow_cast(indexedSubGraph.GetMetaDef()->outputs.size()); + + std::vector isInputsUploadedByDmlEP(fusedNodeInputCount); + for (uint32_t index = 0; index < fusedNodeInputCount; ++index) + { + auto iter = initializerNameToInitializerMap.find(indexedSubGraph.GetMetaDef()->inputs[index]); + isInputsUploadedByDmlEP[index] = iter != initializerNameToInitializerMap.end() ? true : false; + } + + ComPtr device; + ORT_THROW_IF_FAILED(providerImpl->GetDmlDevice(device.GetAddressOf())); + GraphDescBuilder::GraphDesc graphDesc = GraphDescBuilder::BuildGraphDesc( + isInputsUploadedByDmlEP.data(), + isInputsUploadedByDmlEP.size(), + initializerNameToInitializerMap, + graph, + indexedSubGraph, + partitionNodePropsMap, + device.Get(), + providerImpl); + + // convert DML EP GraphDesc into DML_GRAPH_DESC and create IDMLCompiledOperator + DML_GRAPH_DESC dmlGraphDesc = {}; + std::vector dmlOperatorGraphNodes(graphDesc.nodes.size()); + std::vector dmlGraphNodes(graphDesc.nodes.size()); + std::vector dmlInputEdges(graphDesc.inputEdges.size()); + std::vector dmlOutputEdges(graphDesc.outputEdges.size()); + std::vector dmlIntermediateEdges(graphDesc.intermediateEdges.size()); + ConvertGraphDesc( + graphDesc, + dmlGraphDesc, + fusedNodeInputCount, + fusedNodeOutputCount, + dmlOperatorGraphNodes, + dmlGraphNodes, + dmlInputEdges, + dmlOutputEdges, + dmlIntermediateEdges); + + DML_EXECUTION_FLAGS executionFlags = DML_EXECUTION_FLAG_NONE; + if (graphDesc.reuseCommandList) + { + executionFlags |= DML_EXECUTION_FLAG_DESCRIPTORS_VOLATILE; + } + + // Query DML execution provider to see if metacommands is enabled + if (!providerImpl->MetacommandsEnabled()) + { + executionFlags |= DML_EXECUTION_FLAG_DISABLE_META_COMMANDS; + } + + ComPtr device1; + ORT_THROW_IF_FAILED(device.As(&device1)); + ComPtr compiledExecutionPlanOperator; + ORT_THROW_IF_FAILED(device1->CompileGraph( + &dmlGraphDesc, + executionFlags, + IID_PPV_ARGS(&compiledExecutionPlanOperator))); + + // Populate input bindings for operator initialization + std::vector> initializeResourceRefs; // For lifetime control + std::vector initInputBindings(fusedNodeInputCount); + std::vector> initInputResources; // For lifetime control + std::vector> nonOwnedGraphInputsFromInitializers(fusedNodeInputCount); + + std::vector inputsUsed; + ProcessInputData( + providerImpl, + isInputsUploadedByDmlEP, + graphDesc.inputEdges, + indexedSubGraph.GetMetaDef()->inputs, + initializerNameToInitializerMap, + graph, + inputsUsed, + initInputBindings, + initInputResources, + nonOwnedGraphInputsFromInitializers, + initializeResourceRefs, + nullptr); + + // lamda captures for the kernel registration + Windows::AI::MachineLearning::Adapter::EdgeShapes outputShapes; + ORT_THROW_HR_IF(E_UNEXPECTED, !TryGetStaticOutputShapes(fusedNode, outputShapes)); + bool resuableCommandList = graphDesc.reuseCommandList; + auto fused_kernel_func = [compiledExecutionPlanOperator, + outputShapes, + resuableCommandList, + nonOwnedGraphInputsFromInitializers, + initializeResourceRefs, + initInputBindings, + isInputsUploadedByDmlEP, + inputsUsed] + (onnxruntime::FuncManager& func_mgr, const onnxruntime::OpKernelInfo& info, std::unique_ptr& out) mutable ->onnxruntime::Status + { + out.reset(CreateFusedGraphKernel(info, + compiledExecutionPlanOperator, + outputShapes, + resuableCommandList, + nonOwnedGraphInputsFromInitializers, + initializeResourceRefs, + initInputBindings, + isInputsUploadedByDmlEP, + inputsUsed)); + return Status::OK(); + }; + + // build the kernel definition on the fly, and register it to the fused_kernel_regisitry. + onnxruntime::KernelDefBuilder builder; + builder.SetName(indexedSubGraph.GetMetaDef()->name) + .SetDomain(indexedSubGraph.GetMetaDef()->domain) + .SinceVersion(indexedSubGraph.GetMetaDef()->since_version) + .Provider(onnxruntime::kDmlExecutionProvider); + ORT_THROW_IF_ERROR(registryForPartitionKernels->Register(builder, fused_kernel_func)); + } + + void FusePartitionAndRegisterKernel( + GraphPartition* partition, + uint32_t partitionIndex, + onnxruntime::Graph& graph, + std::unordered_map& graphNodePropertyMap, + onnxruntime::KernelRegistry* registryForPartitionKernels, + const std::string& partitionKernelPrefix, + const std::unordered_map>& initializerNameToInitializerMap, + const ExecutionProviderImpl* providerImpl) + { + assert(partition->IsDmlGraphPartition()); + + onnxruntime::IndexedSubGraph indexedSubGraph; + // Create a definition for the node. The name must be unique. + auto def = std::make_unique(); + def->name = DmlGraphFusionTransformer::DML_GRAPH_FUSION_NODE_NAME_PREFIX + partitionKernelPrefix + std::to_string(partitionIndex); + def->domain = DmlGraphFusionTransformer::DML_GRAPH_FUSION_NODE_DOMAIN; + def->since_version = 1; + def->inputs.insert(def->inputs.begin(), partition->GetInputs().begin(), partition->GetInputs().end()); + def->outputs.insert(def->outputs.begin(), partition->GetOutputs().begin(), partition->GetOutputs().end()); + + indexedSubGraph.SetMetaDef(std::move(def)); + indexedSubGraph.nodes = std::move(partition->GetNodeIndices()); + auto& fusedNode = graph.BeginFuseSubGraph(indexedSubGraph, indexedSubGraph.GetMetaDef()->name); + fusedNode.SetExecutionProviderType(onnxruntime::kDmlExecutionProvider); + + // Populate properties which will be passed to OpKernel for this graph via the function below + std::unordered_map partitionNodePropsMap; + for (auto nodeIndex : indexedSubGraph.nodes) + { + const onnxruntime::Node* node = graph.GetNode(nodeIndex); + +#ifdef PRINT_PARTITON_INFO + printf("Partition %u\t%s\n", partitionIndex, GraphDescBuilder::GetUniqueNodeName(*node).c_str()); +#endif + partitionNodePropsMap.insert(std::make_pair( + GraphDescBuilder::GetUniqueNodeName(*node), std::move(graphNodePropertyMap[node]))); + } + +#ifdef PRINT_PARTITON_INFO + printf("\n"); +#endif + CreateIDmlCompiledOperatorAndRegisterKernel( + graph, + indexedSubGraph, + fusedNode, + partitionNodePropsMap, + initializerNameToInitializerMap, + providerImpl, + registryForPartitionKernels); + graph.FinalizeFuseSubGraph(indexedSubGraph, fusedNode); + } +} +} diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.h new file mode 100644 index 0000000000..27b6e87b30 --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.h @@ -0,0 +1,92 @@ +#pragma once + +#include "precomp.h" +#include "GraphDescBuilder.h" +#include "ExecutionProvider.h" +#include "GraphPartitioner.h" +#include "FusedGraphKernel.h" +#include "MLOperatorAuthorImpl.h" + + +namespace Dml +{ +namespace DmlGraphFusionHelper +{ + template + static T AlignToPow2(T offset, T alignment) + { + static_assert(std::is_unsigned_v); + assert(alignment != 0); + assert((alignment & (alignment - 1)) == 0); + return (offset + alignment - 1) & ~(alignment - 1); + } + + Microsoft::WRL::ComPtr + CreateResource( + const ExecutionProviderImpl* provider, + const std::byte* tensorPtr, + size_t tensorByteSize); + + Microsoft::WRL::ComPtr + CreateCpuResource( + const ExecutionProviderImpl* provider, + const std::byte* tensorPtr, + size_t tensorByteSize); + + void UnwrapTensor( + Windows::AI::MachineLearning::Adapter::IWinmlExecutionProvider* winmlProvider, + const onnxruntime::Tensor* tensor, + ID3D12Resource** resource, + uint64_t* allocId); + + void ProcessInputData( + const ExecutionProviderImpl* providerImpl, + const std::vector& isInputsUploadedByDmlEP, + std::vector& inputEdges, + const gsl::span subGraphInputArgNames, + const std::unordered_map>& isInitializerTransferable, + onnxruntime::Graph& graph, + _Out_ std::vector& inputsUsed, + _Inout_ std::vector& initInputBindings, + _Inout_ std::vector>& initInputResources, + _Inout_ std::vector>& nonOwnedGraphInputsFromInitializers, + _Inout_ std::vector>& initializeResourceRefs, + _Inout_opt_ std::vector>* inputRawData); + + std::unordered_map> + GetInitializerToPartitionMap( + const onnxruntime::GraphViewer& graph, + gsl::span> partitions + ); + + void ConvertGraphDesc( + const Dml::GraphDescBuilder::GraphDesc& graphDesc, + _Out_ DML_GRAPH_DESC& dmlGraphDesc, + const uint32_t inputCount, + const uint32_t outputCount, + _Inout_ std::vector& dmlOperatorGraphNodes, + _Inout_ std::vector& dmlGraphNodes, + _Inout_ std::vector& dmlInputEdges, + _Inout_ std::vector& dmlOutputEdges, + _Inout_ std::vector& dmlIntermediateEdges); + + void CreateIDmlCompiledOperatorAndRegisterKernel( + onnxruntime::Graph& graph, + const onnxruntime::IndexedSubGraph& indexedSubGraph, + const onnxruntime::Node& fusedNode, + const std::unordered_map& partitionNodePropsMap, + const std::unordered_map>& isInitializerTransferable, + const ExecutionProviderImpl* providerImpl, + onnxruntime::KernelRegistry* registryForPartitionKernels); + + void FusePartitionAndRegisterKernel( + GraphPartition* partition, + uint32_t partitionIndex, + onnxruntime::Graph& graph, + std::unordered_map& graphNodePropertyMap, + onnxruntime::KernelRegistry* registryForPartitionKernels, + const std::string& partitionKernelPrefix, + const std::unordered_map>& isInitializerTransferable, + const ExecutionProviderImpl* providerImpl); +} +} \ No newline at end of file diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionTransformer.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionTransformer.cpp new file mode 100644 index 0000000000..1b3954de2f --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionTransformer.cpp @@ -0,0 +1,122 @@ +#pragma once + +#include "precomp.h" +#include "GraphDescBuilder.h" +#include "ExecutionProvider.h" +#include "DmlGraphFusionTransformer.h" +#include "GraphPartitioner.h" +#include "core/framework/kernel_type_str_resolver.h" +#include "core/framework/kernel_lookup.h" +#include "FusedGraphKernel.h" +#include "MLOperatorAuthorImpl.h" +#include "DmlGraphFusionHelper.h" + + +namespace Dml +{ + DmlGraphFusionTransformer::DmlGraphFusionTransformer( + const std::string& name, + const onnxruntime::IExecutionProvider* provider + ) + :onnxruntime::GraphTransformer(name), + m_providerImpl(static_cast(provider)->GetImpl()) + { + } + + onnxruntime::common::Status DmlGraphFusionTransformer::ApplyImpl( + onnxruntime::Graph& graph, + bool& modified, + int graph_level, + const onnxruntime::logging::Logger& logger) const + { + onnxruntime::ProviderType provider_type = onnxruntime::kDmlExecutionProvider; + const gsl::not_null registry = m_providerImpl->GetKernelRegistry().get(); + const auto kernel_type_str_resolver = onnxruntime::OpSchemaKernelTypeStrResolver{}; + const auto kernel_lookup = onnxruntime::KernelLookup{provider_type, + gsl::make_span(®istry, 1), + kernel_type_str_resolver}; + + // Initializers needed by any graph partition + std::unordered_set requiredInitializerMap; + std::unordered_map graphNodePropertyMap; + onnxruntime::GraphViewer graphViewer(graph); + std::vector> partitions = BuildPartitions( + graphViewer, + *m_providerImpl->GetInternalRegistrationInfoMap(), + kernel_lookup, + m_providerImpl->GetSupportedDeviceDataTypeMask(), + graphNodePropertyMap, + requiredInitializerMap); + + // Create a map between each initialized tensor and the partition(s) it is part of. + auto initializerPartitionMap = DmlGraphFusionHelper::GetInitializerToPartitionMap(graphViewer, partitions); + + for (uint32_t partitionIndex = 0; partitionIndex < partitions.size(); ++partitionIndex) + { + auto& partition = partitions[partitionIndex]; + + if (partition->GetRootMergedPartition() != partition.get() || + !partition->IsDmlPartition()) + { + continue; + } + + // This map will tell which initializer can be removed from onnxruntime::Graph (and from it's field + // onnx::GraphProto) while we upload the initializer to GPU. + // Why we want to remove the initializer from ORT? + // 1. To keep the peak memory usage as low as possible. That's why we are doing incremental upload to GPU. + // What is initializer? + // An initializer is a input tensor to an operator or the graph itself, which is contant and will never change. + // Why are we uploading the initialzer now? + // This prevents OnnxRuntime from allocating GPU resources and uploading those initializers, + // so the partiton's kernel can do so. In the process, it will pre-process weights while consuming a CPU + // backed resource, avoiding an extra set of GPU resources in memory. + std::unordered_map> isInitializerTransferable; + + + if (partition->IsDmlGraphPartition()) + { + // populate transferredInitializerMap + for (const auto& input : partition->GetInputs()) + { + const onnx::TensorProto* tensor = nullptr; + if (graph.GetInitializedTensor(input, tensor)) + { + // It's only safe to transfer tensors which are used by this partition alone. + auto iter = initializerPartitionMap.find(tensor); + assert(iter != initializerPartitionMap.end()); + if (iter->second.size() > 1) + { + if (requiredInitializerMap.find(input) != requiredInitializerMap.end()) + { + // The kernel relies on this input to be initialized, and it should be small enough to copy + // cheaply. FusedGraphKernel only handles constant CPU inputs through transferred initializers, + // rather than ORT, to avoid mismatches in policy or implementation causing failures. + isInitializerTransferable[input] = {tensor, false}; + } + + continue; + } + isInitializerTransferable[input] = {tensor, true}; + } + } + + std::string partitionKernelPrefix = std::to_string(m_providerImpl->GetPartitionKernelPrefixVal()) + "_"; + m_providerImpl->IncreasePartitionKernelPrefixVal(); + + DmlGraphFusionHelper::FusePartitionAndRegisterKernel( + partition.get(), + partitionIndex, + graph, + graphNodePropertyMap, + m_providerImpl->GetKernelRegistry().get(), + partitionKernelPrefix, + isInitializerTransferable, + m_providerImpl + ); + } + } + + return onnxruntime::common::Status::OK(); + } +} diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionTransformer.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionTransformer.h new file mode 100644 index 0000000000..b546f29f59 --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionTransformer.h @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + + +#include "core/optimizer/graph_transformer.h" +#include "core/framework/execution_providers.h" + +namespace Dml +{ + class ExecutionProviderImpl; + + class DmlGraphFusionTransformer : public onnxruntime::GraphTransformer + { + public: + DmlGraphFusionTransformer( + const std::string& name, + const onnxruntime::IExecutionProvider* provider + ); + + public: + inline const static char* const DML_GRAPH_FUSION_NODE_NAME_PREFIX = "DmlFusedNode_"; + inline const static char* const DML_GRAPH_FUSION_NODE_DOMAIN = "DmlFusedNodeDomain"; + + private: + onnxruntime::common::Status ApplyImpl(onnxruntime::Graph& graph, + bool& modified, + int graph_level, + const onnxruntime::logging::Logger& logger) const final; + private: + const ExecutionProviderImpl* m_providerImpl = nullptr; + }; +} diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index 8d6b50592f..64e3f73ec8 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -518,22 +518,165 @@ namespace Dml return Dml::GetSupportedDeviceDataTypeMask(m_dmlDevice.Get()); } + bool TryGetTensorDataType( + const onnxruntime::NodeArg& nodeArg, + _Out_ MLOperatorTensorDataType* onnxElementType + ) + { + *onnxElementType = MLOperatorTensorDataType::Undefined; + + const ::onnx::TypeProto* typeProto = nodeArg.TypeAsProto(); + if (typeProto != nullptr && typeProto->has_tensor_type()) + { + const ::onnx::TypeProto_Tensor& tensorTypeProto = typeProto->tensor_type(); + if (tensorTypeProto.has_elem_type()) + { + *onnxElementType = static_cast(tensorTypeProto.elem_type()); + return true; + } + } + + return false; + } + + bool DoesNodeContainSupportedDataTypes( + const onnxruntime::Node& node, + _In_opt_ const InternalRegistrationInfo* regInfo, + uint32_t supportedDeviceDataTypeMask // Each bit corresponds to each DML_TENSOR_DATA_TYPE. + ) + { + std::vector constantCpuInputs; + + if (regInfo != nullptr) + { + // Collect the list of CPU-bound input tensors, needed when checking 64-bit fallback + // or for other data types like int-8 which may be supported for CPU inputs but not + // GPU inputs. + auto inputDefinitions = node.InputDefs(); + for (uint32_t i : regInfo->requiredConstantCpuInputs) + { + if (i < inputDefinitions.size()) + { + constantCpuInputs.push_back(inputDefinitions[i]); + } + } + } + + // Assume data types are supported until proven otherwise. + bool nodeContainsSupportedDataTypes = true; + + // Callback to check each node's data type against registered operator support. + std::function nodeCallback = [&](const onnxruntime::NodeArg& nodeArg, bool isInput) -> void + { + // Get the tensor element data type for this node, comparing against what the device actually supports. + // Use the enumeration from the proto instead of nodeArg.Type() which returns a string. + + // Reject node if undefined data type or non-tensor, as DML cannot handle it. + MLOperatorTensorDataType onnxElementType; + if (!TryGetTensorDataType(nodeArg, &onnxElementType)) + { + // We shouldn't have arrived here because (1) no DML operators should have been + // registered which use non-tensor types (2) ONNX validation should have already + // been done, checking for the right kind of inputs and attributes. In theory, + // this branch could be reached with a bad custom operator or malformed file. If + // a legitimate case reaches here and DML needs to support a new input/output type + // besides tensors, then remove the assert. + assert(false); + nodeContainsSupportedDataTypes = false; + return; + } + + // Reject node for unknown DML data types. + DML_TENSOR_DATA_TYPE dmlElementType = GetDmlDataTypeFromMlDataTypeNoThrow(onnxElementType); + if (dmlElementType == DML_TENSOR_DATA_TYPE_UNKNOWN) + { + nodeContainsSupportedDataTypes = false; + return; + } + + // Succeed if the tensor is CPU-bound, as the CPU-side reading code is generic enough + // to handle multiple types regardless of GPU capability (typically these are just + // scalars or simple 1D arrays). + bool isConstantCpuInput = isInput && std::find(constantCpuInputs.begin(), constantCpuInputs.end(), &nodeArg) != constantCpuInputs.end(); + if (isConstantCpuInput) + { + // Leave nodeContainsSupportedDataTypes alone. + return; + } + + bool isDataTypeSupported = (1 << dmlElementType) & supportedDeviceDataTypeMask; + + // Reject node if the data type is unsupported by the device. + if (!isDataTypeSupported) + { + nodeContainsSupportedDataTypes = false; + return; + } + + // Otherwise the node supports the tensor data type. + }; + + // Check whether the node uses any data types which are unsupported by the device. + node.ForEachDef(nodeCallback); + + return nodeContainsSupportedDataTypes; + } + + bool ExecutionProviderImpl::IsNodeSupportedByDml( + const onnxruntime::Node& node, + const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup, + uint32_t supportedDeviceDataTypeMask // Each bit corresponds to each DML_TENSOR_DATA_TYPE. + ) const + { + const onnxruntime::KernelCreateInfo* createInfo = kernel_lookup.LookUpKernel(node); + if (!createInfo) + { + return false; + } + + auto regInfoIter = m_internalRegInfoMap->find(createInfo->kernel_def.get()); + std::shared_ptr internalRegInfo; + if (regInfoIter != m_internalRegInfoMap->end()) + { + internalRegInfo = regInfoIter->second; + if (internalRegInfo->supportQuery && !internalRegInfo->supportQuery(node)) + { + return false; + } + } + + // Check whether the node uses any data types which are unsupported by the device. + if (!DoesNodeContainSupportedDataTypes(node, internalRegInfo.get(), supportedDeviceDataTypeMask)) + { + return false; + } + + return true; + } + std::vector> ExecutionProviderImpl::GetCapability( const onnxruntime::GraphViewer& graph, const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup) const { - std::string partitionKernelPrefix = std::to_string(m_partitionKernelPrefixVal++) + "_"; - uint32_t deviceDataTypeMask = GetSupportedDeviceDataTypeMask(); + uint32_t deviceDataTypeMask = GetSupportedDeviceDataTypeMask(); // Each bit corresponds to each DML_TENSOR_DATA_TYPE. - return PartitionGraph( - graph, - *m_internalRegInfoMap, - kernel_lookup, - deviceDataTypeMask, - m_kernelRegistry.get(), - partitionKernelPrefix - ); + std::vector> result; + + // Get the list of node indices in toplogical order, so nodes are visited before + // downstream nodes consuming them. + const std::vector& toplogicalOrder = graph.GetNodesInTopologicalOrder(); + for (size_t nodeIndex : toplogicalOrder) + { + const onnxruntime::Node& node = *graph.GetNode(nodeIndex); + if (IsNodeSupportedByDml(node, kernel_lookup, deviceDataTypeMask)) + { + std::unique_ptr subGraph = std::make_unique(); + subGraph->nodes = {nodeIndex}; + result.push_back(std::make_unique(std::move(subGraph))); + } + } + return result; } bool IsGpuTensor(const onnxruntime::Tensor& tensor) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h index 3fba9c755a..68400eaeae 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h @@ -158,11 +158,27 @@ namespace Dml std::shared_ptr GetInternalRegistrationInfoMap() const; + void IncreasePartitionKernelPrefixVal() const + { + m_partitionKernelPrefixVal++; + } + + uint64_t GetPartitionKernelPrefixVal() const + { + return m_partitionKernelPrefixVal; + } + onnxruntime::common::Status OnSessionInitializationEnd(); private: void Initialize(ID3D12CommandQueue* queue, ExecutionProvider& executionProvider); + bool IsNodeSupportedByDml( + const onnxruntime::Node& node, + const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup, + uint32_t supportedDeviceDataTypeMask // Each bit corresponds to each DML_TENSOR_DATA_TYPE. + ) const; + ComPtr m_d3d12Device; ComPtr m_dmlDevice; bool m_isMcdmDevice = false; @@ -176,7 +192,6 @@ namespace Dml std::shared_ptr m_kernelRegistry; std::shared_ptr m_internalRegInfoMap; mutable uint64_t m_partitionKernelPrefixVal = 0; - bool m_closed = false; }; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp index f1ef0ed1df..64ff21a1dd 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp @@ -5,7 +5,7 @@ #include "MLOperatorAuthorImpl.h" #include "FusedGraphKernel.h" -#include "GraphKernelHelper.h" +#include "DmlGraphFusionHelper.h" using namespace Windows::AI::MachineLearning::Adapter; @@ -18,23 +18,21 @@ namespace Dml FusedGraphKernel( const onnxruntime::OpKernelInfo& kernelInfo, - const std::unordered_map &graphNodePropertyMap, - std::unordered_map& transferredInitializerMap, - const gsl::span fusedNodeInputArgOriginalNames, - const gsl::span fusedNodeOutputArgOriginalNames) : OpKernel(kernelInfo) - { - // Get the graph for the function which was created according to the computational - // capacity returned by the execution provider's graph partitioner - auto& node = kernelInfo.node(); - ORT_THROW_HR_IF(E_UNEXPECTED, node.NodeType() != onnxruntime::Node::Type::Fused); - auto func = node.GetFunctionBody(); - const onnxruntime::Graph& graph = func->Body(); - - // Get the shapes for outputs of the overall graph. These should be static, because - // the partitioner checked that each node has static shapes before fusing into a - // graph partition. - ORT_THROW_HR_IF(E_UNEXPECTED, !TryGetStaticOutputShapes(node, m_outputShapes)); - + ComPtr compiledExecutionPlanOperator, + Windows::AI::MachineLearning::Adapter::EdgeShapes& outputShapes, + bool reuseCommandList, + std::vector>& nonOwnedGraphInputsFromInitializers, + std::vector>& initializeResourceRefs, + std::vector initInputBindings, + std::vector& isInputsUploadedByDmlEP, + std::vector& inputsUsed) : + OpKernel(kernelInfo), + m_compiledExecutionPlanOperator(compiledExecutionPlanOperator), + m_inputsUsed(inputsUsed), + m_outputShapes(outputShapes), + m_isInputsUploadedByDmlEP(isInputsUploadedByDmlEP), + m_nonOwnedGraphInputsFromInitializers(nonOwnedGraphInputsFromInitializers) + { // Get the execution provider interfaces m_executionHandle = kernelInfo.GetExecutionProvider()->GetExecutionHandle(); if (m_executionHandle) @@ -48,105 +46,19 @@ namespace Dml } TranslateAndCompileGraph( - kernelInfo, - graph, - fusedNodeInputArgOriginalNames, - fusedNodeOutputArgOriginalNames, - graphNodePropertyMap, - transferredInitializerMap); + kernelInfo, + initializeResourceRefs, + initInputBindings, + reuseCommandList); } void TranslateAndCompileGraph( const onnxruntime::OpKernelInfo& kernelInfo, - const onnxruntime::Graph& graph, - const gsl::span fusedNodeInputArgOriginalNames, - const gsl::span fusedNodeOutputArgOriginalNames, - const std::unordered_map& graphNodePropertyMap, - std::unordered_map& transferredInitializerMap + std::vector>& initializeResourceRefs, + std::vector initInputBindings, + bool reuseCommandList ) { - ComPtr device; - ORT_THROW_IF_FAILED(m_provider->GetDmlDevice(device.GetAddressOf())); - - ComPtr device1; - ORT_THROW_IF_FAILED(device.As(&device1)); - - const uint32_t graphInputCount = kernelInfo.GetInputCount(); - - m_inputsConstant.resize(graphInputCount); - for (uint32_t i = 0; i < graphInputCount; ++i) - { - m_inputsConstant[i] = GraphKernelHelper::GetGraphInputConstness(i, kernelInfo, fusedNodeInputArgOriginalNames, transferredInitializerMap); - } - - GraphDescBuilder::GraphDesc graphDesc = GraphDescBuilder::BuildGraphDesc( - kernelInfo, - m_inputsConstant.data(), - m_inputsConstant.size(), - transferredInitializerMap, - graph, - fusedNodeInputArgOriginalNames, - fusedNodeOutputArgOriginalNames, - graphNodePropertyMap, - device.Get(), - m_executionHandle); - - // Populate input bindings for operator initialization - std::vector> initInputResources; // For lifetime control - std::vector initInputBindings(graphInputCount); - m_nonOwnedGraphInputsFromInitializers.resize(graphInputCount); - std::vector> initializeResourceRefs; - - GraphKernelHelper::ProcessInputData( - m_provider.Get(), - m_winmlProvider.Get(), - m_inputsConstant, - kernelInfo, - graphDesc, - fusedNodeInputArgOriginalNames, - m_inputsUsed, - initInputBindings, - initInputResources, - m_nonOwnedGraphInputsFromInitializers, - initializeResourceRefs, - nullptr, - transferredInitializerMap); - - DML_GRAPH_DESC dmlGraphDesc = {}; - std::vector dmlOperatorGraphNodes(graphDesc.nodes.size()); - std::vector dmlGraphNodes(graphDesc.nodes.size()); - - std::vector dmlInputEdges(graphDesc.inputEdges.size()); - std::vector dmlOutputEdges(graphDesc.outputEdges.size()); - std::vector dmlIntermediateEdges(graphDesc.intermediateEdges.size()); - - GraphKernelHelper::ConvertGraphDesc( - graphDesc, - dmlGraphDesc, - kernelInfo, - dmlOperatorGraphNodes, - dmlGraphNodes, - dmlInputEdges, - dmlOutputEdges, - dmlIntermediateEdges); - - DML_EXECUTION_FLAGS executionFlags = DML_EXECUTION_FLAG_NONE; - if (graphDesc.reuseCommandList) - { - executionFlags |= DML_EXECUTION_FLAG_DESCRIPTORS_VOLATILE; - } - - // Query DML execution provider to see if metacommands is enabled - if (!m_provider->MetacommandsEnabled()) - { - executionFlags |= DML_EXECUTION_FLAG_DISABLE_META_COMMANDS; - } - - ORT_THROW_IF_FAILED(device1->CompileGraph( - &dmlGraphDesc, - executionFlags, - IID_PPV_ARGS(&m_compiledExecutionPlanOperator))); - // Allocate a persistent resource and initialize the operator UINT64 persistentResourceSize = m_compiledExecutionPlanOperator->GetBindingProperties().PersistentResourceSize; if (persistentResourceSize > 0) @@ -175,7 +87,7 @@ namespace Dml [&](ComPtr& resource){ m_winmlProvider->QueueReference(WRAP_GRAPHICS_UNKNOWN(resource).Get()); } ); - if (graphDesc.reuseCommandList) + if (reuseCommandList) { BuildReusableCommandList(); } @@ -213,7 +125,7 @@ namespace Dml { inputPtrs[i] = m_nonOwnedGraphInputsFromInitializers[i].Get(); } - else if (!m_inputsConstant[i]) + else if (!m_isInputsUploadedByDmlEP[i]) { ORT_THROW_IF_FAILED(contextWrapper.GetInputTensor(i, inputTensors[i].GetAddressOf())); inputPtrs[i] = m_provider->DecodeResource(MLOperatorTensor(inputTensors[i].Get()).GetDataInterface().Get()); @@ -380,7 +292,7 @@ namespace Dml for (uint32_t i = 0; i < inputBindings.size(); ++i) { - if (!m_inputsConstant[i] && m_inputsUsed[i]) + if (!m_isInputsUploadedByDmlEP[i] && m_inputsUsed[i]) { if (m_nonOwnedGraphInputsFromInitializers[i]) { @@ -393,10 +305,10 @@ namespace Dml const onnxruntime::Tensor* tensor = kernelContext->Input(i); uint64_t allocId; - GraphKernelHelper::UnwrapTensor(m_winmlProvider.Get(), tensor, &inputBindings[i].Buffer, &allocId); + DmlGraphFusionHelper::UnwrapTensor(m_winmlProvider.Get(), tensor, &inputBindings[i].Buffer, &allocId); inputBindingsChanged = inputBindingsChanged || (!allocId || m_inputBindingAllocIds[i] != allocId); inputBindings[i].Buffer->Release(); // Avoid holding an additional reference - inputBindings[i].SizeInBytes = GraphKernelHelper::AlignToPow2(tensor->SizeInBytes(), 4); + inputBindings[i].SizeInBytes = DmlGraphFusionHelper::AlignToPow2(tensor->SizeInBytes(), 4); inputBindingDescs[i] = {DML_BINDING_TYPE_BUFFER, &inputBindings[i]}; m_inputBindingAllocIds[i] = allocId; } @@ -430,10 +342,10 @@ namespace Dml ); uint64_t allocId; - GraphKernelHelper::UnwrapTensor(m_winmlProvider.Get(), tensor, &outputBindings[i].Buffer, &allocId); + DmlGraphFusionHelper::UnwrapTensor(m_winmlProvider.Get(), tensor, &outputBindings[i].Buffer, &allocId); outputBindingsChanged = outputBindingsChanged || (!allocId || m_outputBindingAllocIds[i] != allocId); outputBindings[i].Buffer->Release(); // Avoid holding an additional reference - outputBindings[i].SizeInBytes = GraphKernelHelper::AlignToPow2(tensor->SizeInBytes(), 4); + outputBindings[i].SizeInBytes = DmlGraphFusionHelper::AlignToPow2(tensor->SizeInBytes(), 4); outputBindingDescs[i] = {DML_BINDING_TYPE_BUFFER, &outputBindings[i]}; m_outputBindingAllocIds[i] = allocId; } @@ -488,7 +400,7 @@ namespace Dml const void* m_executionHandle = nullptr; ComPtr m_winmlProvider; ComPtr m_provider; - Windows::AI::MachineLearning::Adapter::EdgeShapes m_outputShapes; + Windows::AI::MachineLearning::Adapter::EdgeShapes& m_outputShapes; // Re-usable command list, supporting descriptor heap, and DML binding table to update that heap. ComPtr m_graphicsCommandList; @@ -509,23 +421,32 @@ namespace Dml mutable ComPtr m_fence; mutable uint64_t m_completionValue = 0; - std::vector m_inputsConstant; + std::vector m_isInputsUploadedByDmlEP; std::vector> m_nonOwnedGraphInputsFromInitializers; }; onnxruntime::OpKernel* CreateFusedGraphKernel( - const onnxruntime::OpKernelInfo& info, - const std::unordered_map &graphNodePropertyMap, - std::unordered_map& transferredInitializerMap, - const gsl::span fusedNodeInputArgOriginalNames, - const gsl::span fusedNodeOutputArgOriginalNames + const onnxruntime::OpKernelInfo& info, + ComPtr compiledExecutionPlanOperator, + Windows::AI::MachineLearning::Adapter::EdgeShapes& outputShapes, + bool reuseCommandList, + std::vector>& nonOwnedGraphInputsFromInitializers, + std::vector>& initializeResourceRefs, + std::vector initInputBindings, + std::vector& isInputsUploadedByDmlEP, + std::vector& inputsUsed ) { return new FusedGraphKernel( info, - graphNodePropertyMap, - transferredInitializerMap, - fusedNodeInputArgOriginalNames, - fusedNodeOutputArgOriginalNames); + compiledExecutionPlanOperator, + outputShapes, + reuseCommandList, + nonOwnedGraphInputsFromInitializers, + initializeResourceRefs, + initInputBindings, + isInputsUploadedByDmlEP, + inputsUsed + ); } } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.h index 45a08550b9..00a858d54e 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.h @@ -3,14 +3,19 @@ #include "core/framework/op_kernel.h" #include "GraphDescBuilder.h" +#include "DmlGraphFusionTransformer.h" namespace Dml { onnxruntime::OpKernel* CreateFusedGraphKernel( - const onnxruntime::OpKernelInfo& info, - const std::unordered_map& graphNodePropertyMap, - std::unordered_map& transferredInitializerMap, - const gsl::span fusedNodeInputArgOriginalNames, - const gsl::span fusedNodeOutputArgOriginalNames + const onnxruntime::OpKernelInfo& info, + ComPtr compiledExecutionPlanOperator, + Windows::AI::MachineLearning::Adapter::EdgeShapes& outputShapes, + bool reuseCommandList, + std::vector>& nonOwnedGraphInputsFromInitializers, + std::vector>& initializeResourceRefs, + std::vector initInputBindings, + std::vector& isInputsUploadedByDmlEP, + std::vector& inputsUsed ); } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp index cadcf0c8d0..2727f81761 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp @@ -3,7 +3,6 @@ #include "precomp.h" #include "GraphDescBuilder.h" -#include "GraphKernelHelper.h" using namespace Windows::AI::MachineLearning::Adapter; @@ -33,17 +32,18 @@ namespace Dml::GraphDescBuilder #pragma warning(pop) GraphDesc BuildGraphDesc( - const onnxruntime::OpKernelInfo& kernelInfo, const uint8_t* isConstGpuGraphInput, const size_t isConstGpuGraphInputCount, - std::unordered_map& transferredInitializerMap, + const std::unordered_map>& isInitializerTransferable, const onnxruntime::Graph& graph, - const gsl::span fusedNodeInputArgOriginalNames, - const gsl::span fusedNodeOutputArgOriginalNames, + const onnxruntime::IndexedSubGraph& indexedSubGraph, const std::unordered_map& graphNodePropertyMap, IDMLDevice* device, const void* executionHandle) { + + const gsl::span subGraphInputArgNames = indexedSubGraph.GetMetaDef()->inputs; + const gsl::span subGraphOutputArgNames = indexedSubGraph.GetMetaDef()->outputs; struct NodeAndIndex { uint32_t nodeIndex; // The index of the node itself @@ -56,9 +56,9 @@ namespace Dml::GraphDescBuilder // Map from Lotus node argument names to input indices of the fused kernel node. std::unordered_map nameToDmlFusedNodeInputIndex; - for (size_t inputIndex = 0; inputIndex < fusedNodeInputArgOriginalNames.size(); ++inputIndex) + for (size_t inputIndex = 0; inputIndex < subGraphInputArgNames.size(); ++inputIndex) { - const onnxruntime::NodeArg* graphInput = graph.GetNodeArg(fusedNodeInputArgOriginalNames[inputIndex]); + const onnxruntime::NodeArg* graphInput = graph.GetNodeArg(subGraphInputArgNames[inputIndex]); if (!graphInput) { @@ -79,37 +79,34 @@ namespace Dml::GraphDescBuilder std::vector graphIntermediateEdges; std::vector graphOutputEdges; - // Get the topological sorting of Lotus nodes - // paulm: breaking change from LOTUS that removed GetNodesInTopologicalOrder from Graph - onnxruntime::GraphViewer viewer(graph); - const std::vector& orderedNodeIndices = viewer.GetNodesInTopologicalOrder(); - // Avoid using separate command lists for small graphs. This value can be reduced by tuning the // flushing behavior of DmlCommandRecorder. Its current behavior is to assume that graphs contain // enough GPU work to be worth flushing immediately. const uint32_t minNodeCountToReuseCommandList = 5; bool reuseCommandList = false; - if (orderedNodeIndices.size() >= minNodeCountToReuseCommandList) + if (indexedSubGraph.nodes.size() >= minNodeCountToReuseCommandList) { reuseCommandList = true; } - auto constantCpuGraphInputGetter = [&transferredInitializerMap](const std::string& argName) + auto constantCpuGraphInputGetter = [&isInitializerTransferable](const std::string& argName) { ComPtr tensorWrapper; - auto iter = transferredInitializerMap.find(argName); - if (iter != transferredInitializerMap.end()) + auto iter = isInitializerTransferable.find(argName); + if (iter != isInitializerTransferable.end()) { - tensorWrapper = wil::MakeOrThrow(&iter->second); + // Using const_cast here is simpler than making surrounding code const correct. + tensorWrapper = wil::MakeOrThrow(const_cast(iter->second.first)); } return tensorWrapper; }; // Iterate through each node and create a corresponding node in the new graph - for (size_t sortedNodeIndex : orderedNodeIndices) + // We can iterate the nodes in any order because the edge connectivity will take care of the topological order + for (size_t sortedNodeIndex : indexedSubGraph.nodes) { const onnxruntime::Node& node = *graph.GetNode(sortedNodeIndex); @@ -257,9 +254,9 @@ namespace Dml::GraphDescBuilder } // Add graph output nodes, which might be in a different order from the encapsulating node - for (size_t outputIndex = 0; outputIndex < fusedNodeOutputArgOriginalNames.size(); ++outputIndex) + for (size_t outputIndex = 0; outputIndex < subGraphOutputArgNames.size(); ++outputIndex) { - const onnxruntime::NodeArg* graphOutput = graph.GetNodeArg(fusedNodeOutputArgOriginalNames[outputIndex]); + const onnxruntime::NodeArg* graphOutput = graph.GetNodeArg(subGraphOutputArgNames[outputIndex]); ORT_THROW_HR_IF_NULL_MSG(E_POINTER, graphOutput, "FusedNode's nodeArgList does not contain one of the nodeArg"); const auto& outputNodeAndIndex = nameToNodeAndIndexMap.at(graphOutput->Name()); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h index e787332884..bf2ccae55e 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h @@ -40,13 +40,11 @@ namespace Dml }; GraphDesc BuildGraphDesc( - const onnxruntime::OpKernelInfo& kernelInfo, const uint8_t* isConstGpuGraphInput, const size_t isConstGpuGraphInputCount, - std::unordered_map& transferredInitializerMap, + const std::unordered_map>& isInitializerTransferable, const onnxruntime::Graph& graph, - const gsl::span fusedNodeInputArgOriginalNames, - const gsl::span fusedNodeOutputArgOriginalNames, + const onnxruntime::IndexedSubGraph& indexedSubGraph, const std::unordered_map& graphNodePropertyMap, IDMLDevice* device, const void* executionHandle); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphKernelHelper.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphKernelHelper.cpp deleted file mode 100644 index a55a55c952..0000000000 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphKernelHelper.cpp +++ /dev/null @@ -1,324 +0,0 @@ -#include "precomp.h" - -#include "GraphKernelHelper.h" - -namespace Dml -{ -namespace GraphKernelHelper -{ - Microsoft::WRL::ComPtr - CreateResource( - Dml::IExecutionProvider* provider, - const std::byte* tensorPtr, - size_t tensorByteSize) - { - Microsoft::WRL::ComPtr buffer; - - D3D12_HEAP_PROPERTIES heapProperties = { - D3D12_HEAP_TYPE_DEFAULT, D3D12_CPU_PAGE_PROPERTY_UNKNOWN, D3D12_MEMORY_POOL_UNKNOWN, 0, 0}; - - D3D12_RESOURCE_DESC resourceDesc = {D3D12_RESOURCE_DIMENSION_BUFFER, - 0, - static_cast((tensorByteSize + 3) & ~3), - 1, - 1, - 1, - DXGI_FORMAT_UNKNOWN, - {1, 0}, - D3D12_TEXTURE_LAYOUT_ROW_MAJOR, - D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS}; - - Microsoft::WRL::ComPtr d3dDevice; - ORT_THROW_IF_FAILED(provider->GetD3DDevice(d3dDevice.GetAddressOf())); - - ORT_THROW_IF_FAILED(d3dDevice->CreateCommittedResource( - &heapProperties, - D3D12_HEAP_FLAG_NONE, - &resourceDesc, - D3D12_RESOURCE_STATE_UNORDERED_ACCESS, - nullptr, - IID_GRAPHICS_PPV_ARGS(buffer.GetAddressOf()))); - - ORT_THROW_IF_FAILED(provider->UploadToResource(buffer.Get(), tensorPtr, tensorByteSize)); - - return buffer; - } - - Microsoft::WRL::ComPtr - CreateCpuResource( - Dml::IExecutionProvider* provider, - const std::byte* tensorPtr, - size_t tensorByteSize) - { - Microsoft::WRL::ComPtr buffer; - - D3D12_HEAP_PROPERTIES heapProperties = { - D3D12_HEAP_TYPE_CUSTOM, D3D12_CPU_PAGE_PROPERTY_WRITE_COMBINE, D3D12_MEMORY_POOL_L0, 0, 0}; - - D3D12_RESOURCE_DESC resourceDesc = {D3D12_RESOURCE_DIMENSION_BUFFER, - 0, - static_cast((tensorByteSize + 3) & ~3), - 1, - 1, - 1, - DXGI_FORMAT_UNKNOWN, - {1, 0}, - D3D12_TEXTURE_LAYOUT_ROW_MAJOR, - D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS}; - - Microsoft::WRL::ComPtr d3dDevice; - ORT_THROW_IF_FAILED(provider->GetD3DDevice(d3dDevice.GetAddressOf())); - - ORT_THROW_IF_FAILED(d3dDevice->CreateCommittedResource( - &heapProperties, - D3D12_HEAP_FLAG_NONE, - &resourceDesc, - D3D12_RESOURCE_STATE_UNORDERED_ACCESS, - nullptr, - IID_GRAPHICS_PPV_ARGS(buffer.GetAddressOf()))); - - // Map the buffer and copy the data - void* bufferData = nullptr; - D3D12_RANGE range = {0, tensorByteSize}; - ORT_THROW_IF_FAILED(buffer->Map(0, &range, &bufferData)); - memcpy(bufferData, tensorPtr, tensorByteSize); - buffer->Unmap(0, &range); - - return buffer; - } - - void UnwrapTensor( - IWinmlExecutionProvider* winmlProvider, - const onnxruntime::Tensor* tensor, - ID3D12Resource** resource, - uint64_t* allocId) - { - IUnknown* allocationUnk = static_cast(const_cast(tensor->DataRaw())); - Microsoft::WRL::ComPtr resourceUnk; - winmlProvider->GetABIDataInterface(false, allocationUnk, &resourceUnk); - - *allocId = winmlProvider->TryGetPooledAllocationId(allocationUnk, 0); - - ORT_THROW_IF_FAILED(resourceUnk->QueryInterface(resource)); - } - - bool GetGraphInputConstness( - uint32_t index, - const onnxruntime::OpKernelInfo& kernelInfo, - const gsl::span fusedNodeInputArgOriginalNames, - const std::unordered_map& transferredInitializerMap) - { - // Transferred initializers are uploaded to GPU memory - auto iter = transferredInitializerMap.find(fusedNodeInputArgOriginalNames[index]); - if (iter != transferredInitializerMap.end()) - { - return true; - } - - // If an initializer wasn't transferred, the constant input may be available from ORT - const onnxruntime::Tensor* inputTensor = nullptr; - if (!kernelInfo.TryGetConstantInput(index, &inputTensor) || inputTensor == nullptr) - { - return false; - } - - // Check that the constant ORT input is in GPU memory - if (!strcmp(inputTensor->Location().name, onnxruntime::CPU) || - inputTensor->Location().mem_type == ::OrtMemType::OrtMemTypeCPUOutput || - inputTensor->Location().mem_type == ::OrtMemType::OrtMemTypeCPUInput) - { - return false; - } - - return true; - }; - - void ProcessInputData( - Dml::IExecutionProvider* provider, - IWinmlExecutionProvider* winmlProvider, - const std::vector& inputsConstant, - const onnxruntime::OpKernelInfo& kernelInfo, - const Dml::GraphDescBuilder::GraphDesc& graphDesc, - const gsl::span fusedNodeInputArgOriginalNames, - _Out_ std::vector& inputsUsed, - _Inout_ std::vector& initInputBindings, - _Inout_ std::vector>& initInputResources, - _Inout_ std::vector>& nonOwnedGraphInputsFromInitializers, - _Inout_ std::vector>& initializeResourceRefs, - _Inout_opt_ std::vector>* inputRawData, - _Inout_ std::unordered_map& transferredInitializerMap) - { - const uint32_t graphInputCount = kernelInfo.GetInputCount(); - // Determine the last input which uses an initializer, so initializers can be freed incrementally - // while processing each input in order. - std::map initializerToLastInputIndexMap; - for (uint32_t i = 0; i < graphInputCount; i++) - { - auto iter = transferredInitializerMap.find(fusedNodeInputArgOriginalNames[i]); - if (iter != transferredInitializerMap.end()) { - initializerToLastInputIndexMap[&iter->second] = i; - } - } - - // Walk through each graph edge and mark used inputs - inputsUsed.assign(graphInputCount, false); - for (const DML_INPUT_GRAPH_EDGE_DESC& edge : graphDesc.inputEdges) { - inputsUsed[edge.GraphInputIndex] = true; - } - - for (uint32_t i = 0; i < initInputBindings.size(); i++) - { - // If the input isn't actually used by the graph, nothing ever needs to be bound (either for - // initialization or execution). So just throw away the transferred initializer and skip this input. - if (!inputsUsed[i]) - { - transferredInitializerMap.erase(fusedNodeInputArgOriginalNames[i]); - - if (inputRawData) - { - inputRawData->push_back(std::vector()); - } - - continue; - } - - // Look for the initializer among those transferred from the graph during partitioning - auto iter = transferredInitializerMap.find(fusedNodeInputArgOriginalNames[i]); - if (iter != transferredInitializerMap.end()) - { - std::byte* tensorPtr = nullptr; - size_t tensorByteSize = 0; - std::unique_ptr unpackedTensor; - - auto& initializer = iter->second; - - // The tensor may be stored as raw data or in typed fields. - if (initializer.has_raw_data()) - { - tensorPtr = (std::byte*)(initializer.raw_data().c_str()); - tensorByteSize = initializer.raw_data().size(); - } - else - { - std::tie(unpackedTensor, tensorByteSize) = UnpackTensor(initializer); - tensorPtr = unpackedTensor.get(); - } - - // Tensor sizes in DML must be a multiple of 4 bytes large. - tensorByteSize = AlignToPow2(tensorByteSize, 4); - - if (inputRawData) - { - inputRawData->push_back(std::vector(tensorPtr, tensorPtr + tensorByteSize)); - } - - if (!inputsConstant[i]) - { - // Store the resource to use during execution - ComPtr defaultBuffer = CreateResource(provider, tensorPtr, tensorByteSize); - nonOwnedGraphInputsFromInitializers[i] = defaultBuffer; - initializeResourceRefs.push_back(std::move(defaultBuffer)); - } - else - { - ComPtr initializeInputBuffer; - - // D3D_FEATURE_LEVEL_1_0_CORE doesn't support Custom heaps - if (provider->IsMcdmDevice()) - { - initializeInputBuffer = CreateResource(provider, tensorPtr, tensorByteSize); - } - else - { - initializeInputBuffer = CreateCpuResource(provider, tensorPtr, tensorByteSize); - } - - // Set the binding for operator initialization to the buffer - initInputBindings[i].Buffer = initializeInputBuffer.Get(); - initInputBindings[i].SizeInBytes = tensorByteSize; - initializeResourceRefs.push_back(std::move(initializeInputBuffer)); - } - - // Free the initializer if this is the last usage of it. - if (initializerToLastInputIndexMap[&initializer] == i) - { - transferredInitializerMap.erase(iter); - } - } - else if (inputsConstant[i]) - { - const onnxruntime::Tensor* inputTensor = nullptr; - ORT_THROW_HR_IF(E_UNEXPECTED, !kernelInfo.TryGetConstantInput(i, &inputTensor)); - - const std::byte* tensorData = reinterpret_cast(inputTensor->DataRaw()); - - if (inputRawData) - { - inputRawData->push_back( - std::vector(tensorData, tensorData + inputTensor->SizeInBytes())); - } - - uint64_t allocId; - UnwrapTensor(winmlProvider, inputTensor, &initInputBindings[i].Buffer, &allocId); - initInputBindings[i].SizeInBytes = initInputBindings[i].Buffer->GetDesc().Width; - - initInputBindings[i].Buffer->Release(); // Avoid holding an additional reference - initInputResources.push_back(initInputBindings[i].Buffer); - } - else if (inputRawData) - { - inputRawData->push_back(std::vector()); - } - } - - // All initializers should have been consumed and freed above - assert(transferredInitializerMap.empty()); - } - - void ConvertGraphDesc( - const Dml::GraphDescBuilder::GraphDesc& graphDesc, - _Out_ DML_GRAPH_DESC& dmlGraphDesc, - const onnxruntime::OpKernelInfo& kernelInfo, - _Inout_ std::vector& dmlOperatorGraphNodes, - _Inout_ std::vector& dmlGraphNodes, - _Inout_ std::vector& dmlInputEdges, - _Inout_ std::vector& dmlOutputEdges, - _Inout_ std::vector& dmlIntermediateEdges) - { - const uint32_t graphInputCount = kernelInfo.GetInputCount(); - - for (size_t i = 0; i < graphDesc.nodes.size(); ++i) - { - dmlOperatorGraphNodes[i] = DML_OPERATOR_GRAPH_NODE_DESC{graphDesc.nodes[i].op.Get()}; - dmlGraphNodes[i] = DML_GRAPH_NODE_DESC{DML_GRAPH_NODE_TYPE_OPERATOR, &dmlOperatorGraphNodes[i]}; - } - - for (size_t i = 0; i < graphDesc.inputEdges.size(); ++i) - { - dmlInputEdges[i] = DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_INPUT, &graphDesc.inputEdges[i]}; - } - - for (size_t i = 0; i < graphDesc.outputEdges.size(); ++i) - { - dmlOutputEdges[i] = DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_OUTPUT, &graphDesc.outputEdges[i]}; - } - - for (size_t i = 0; i < graphDesc.intermediateEdges.size(); ++i) - { - dmlIntermediateEdges[i] = - DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_INTERMEDIATE, &graphDesc.intermediateEdges[i]}; - } - - dmlGraphDesc.InputCount = graphInputCount; - dmlGraphDesc.OutputCount = kernelInfo.GetOutputCount(); - dmlGraphDesc.NodeCount = gsl::narrow_cast(dmlGraphNodes.size()); - dmlGraphDesc.Nodes = dmlGraphNodes.data(); - dmlGraphDesc.InputEdgeCount = gsl::narrow_cast(dmlInputEdges.size()); - dmlGraphDesc.InputEdges = dmlInputEdges.data(); - dmlGraphDesc.OutputEdgeCount = gsl::narrow_cast(dmlOutputEdges.size()); - dmlGraphDesc.OutputEdges = dmlOutputEdges.data(); - dmlGraphDesc.IntermediateEdgeCount = gsl::narrow_cast(dmlIntermediateEdges.size()); - dmlGraphDesc.IntermediateEdges = dmlIntermediateEdges.data(); - } -} // namespace GraphKernelHelper -} // namespace Dml \ No newline at end of file diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphKernelHelper.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphKernelHelper.h deleted file mode 100644 index d6b6db8e87..0000000000 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphKernelHelper.h +++ /dev/null @@ -1,73 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "GraphDescBuilder.h" - -namespace Dml -{ -namespace GraphKernelHelper -{ - using namespace Windows::AI::MachineLearning::Adapter; - - template - static T AlignToPow2(T offset, T alignment) - { - static_assert(std::is_unsigned_v); - assert(alignment != 0); - assert((alignment & (alignment - 1)) == 0); - return (offset + alignment - 1) & ~(alignment - 1); - } - - Microsoft::WRL::ComPtr - CreateResource( - Dml::IExecutionProvider* provider, - const std::byte* tensorPtr, - size_t tensorByteSize); - - Microsoft::WRL::ComPtr - CreateCpuResource( - Dml::IExecutionProvider* provider, - const std::byte* tensorPtr, - size_t tensorByteSize); - - void UnwrapTensor( - IWinmlExecutionProvider* winmlProvider, - const onnxruntime::Tensor* tensor, - ID3D12Resource** resource, - uint64_t* allocId); - - bool GetGraphInputConstness( - uint32_t index, - const onnxruntime::OpKernelInfo& kernelInfo, - const gsl::span fusedNodeInputArgOriginalNames, - const std::unordered_map& transferredInitializerMap); - - void ProcessInputData( - Dml::IExecutionProvider* provider, - IWinmlExecutionProvider* winmlProvider, - const std::vector& inputsConstant, - const onnxruntime::OpKernelInfo& kernelInfo, - const Dml::GraphDescBuilder::GraphDesc& graphDesc, - const gsl::span fusedNodeInputArgOriginalNames, - _Out_ std::vector& inputsUsed, - _Inout_ std::vector& initInputBindings, - _Inout_ std::vector>& initInputResources, - _Inout_ std::vector>& nonOwnedGraphInputsFromInitializers, - _Inout_ std::vector>& initializeResourceRefs, - _Inout_opt_ std::vector>* inputRawData, - _Inout_ std::unordered_map& transferredInitializerMap); - - void ConvertGraphDesc( - const Dml::GraphDescBuilder::GraphDesc& graphDesc, - _Out_ DML_GRAPH_DESC& dmlGraphDesc, - const onnxruntime::OpKernelInfo& kernelInfo, - _Inout_ std::vector& dmlOperatorGraphNodes, - _Inout_ std::vector& dmlGraphNodes, - _Inout_ std::vector& dmlInputEdges, - _Inout_ std::vector& dmlOutputEdges, - _Inout_ std::vector& dmlIntermediateEdges); - -} // namespace GraphKernelHelper -} // namespace Dml \ No newline at end of file diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.cpp index 362e110fbd..b9c5f8849a 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.cpp @@ -141,144 +141,6 @@ namespace Dml } }; - bool TryGetTensorDataType( - const onnxruntime::NodeArg& nodeArg, - _Out_ MLOperatorTensorDataType* onnxElementType - ) - { - *onnxElementType = MLOperatorTensorDataType::Undefined; - - const ::onnx::TypeProto* typeProto = nodeArg.TypeAsProto(); - if (typeProto != nullptr && typeProto->has_tensor_type()) - { - const ::onnx::TypeProto_Tensor& tensorTypeProto = typeProto->tensor_type(); - if (tensorTypeProto.has_elem_type()) - { - *onnxElementType = static_cast(tensorTypeProto.elem_type()); - return true; - } - } - - return false; - } - - bool DoesNodeContainSupportedDataTypes( - const onnxruntime::Node& node, - _In_opt_ const InternalRegistrationInfo* regInfo, - uint32_t supportedDeviceDataTypeMask // Each bit corresponds to each DML_TENSOR_DATA_TYPE. - ) - { - std::vector constantCpuInputs; - - if (regInfo != nullptr) - { - // Collect the list of CPU-bound input tensors, needed when checking 64-bit fallback - // or for other data types like int-8 which may be supported for CPU inputs but not - // GPU inputs. - auto inputDefinitions = node.InputDefs(); - for (uint32_t i : regInfo->requiredConstantCpuInputs) - { - if (i < inputDefinitions.size()) - { - constantCpuInputs.push_back(inputDefinitions[i]); - } - } - } - - // Assume data types are supported until proven otherwise. - bool nodeContainsSupportedDataTypes = true; - - // Callback to check each node's data type against registered operator support. - std::function nodeCallback = [&](const onnxruntime::NodeArg& nodeArg, bool isInput) -> void - { - // Get the tensor element data type for this node, comparing against what the device actually supports. - // Use the enumeration from the proto instead of nodeArg.Type() which returns a string. - - // Reject node if undefined data type or non-tensor, as DML cannot handle it. - MLOperatorTensorDataType onnxElementType; - if (!TryGetTensorDataType(nodeArg, &onnxElementType)) - { - // We shouldn't have arrived here because (1) no DML operators should have been - // registered which use non-tensor types (2) ONNX validation should have already - // been done, checking for the right kind of inputs and attributes. In theory, - // this branch could be reached with a bad custom operator or malformed file. If - // a legitimate case reaches here and DML needs to support a new input/output type - // besides tensors, then remove the assert. - assert(false); - nodeContainsSupportedDataTypes = false; - return; - } - - // Reject node for unknown DML data types. - DML_TENSOR_DATA_TYPE dmlElementType = GetDmlDataTypeFromMlDataTypeNoThrow(onnxElementType); - if (dmlElementType == DML_TENSOR_DATA_TYPE_UNKNOWN) - { - nodeContainsSupportedDataTypes = false; - return; - } - - // Succeed if the tensor is CPU-bound, as the CPU-side reading code is generic enough - // to handle multiple types regardless of GPU capability (typically these are just - // scalars or simple 1D arrays). - bool isConstantCpuInput = isInput && std::find(constantCpuInputs.begin(), constantCpuInputs.end(), &nodeArg) != constantCpuInputs.end(); - if (isConstantCpuInput) - { - // Leave nodeContainsSupportedDataTypes alone. - return; - } - - bool isDataTypeSupported = (1 << dmlElementType) & supportedDeviceDataTypeMask; - - // Reject node if the data type is unsupported by the device. - if (!isDataTypeSupported) - { - nodeContainsSupportedDataTypes = false; - return; - } - - // Otherwise the node supports the tensor data type. - }; - - // Check whether the node uses any data types which are unsupported by the device. - node.ForEachDef(nodeCallback); - - return nodeContainsSupportedDataTypes; - } - - bool IsNodeSupportedByDml( - const onnxruntime::Node& node, - const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup, - uint32_t supportedDeviceDataTypeMask, // Each bit corresponds to each DML_TENSOR_DATA_TYPE. - const InternalRegistrationInfoMap& internalRegInfoMap - ) - { - const onnxruntime::KernelCreateInfo* createInfo = kernel_lookup.LookUpKernel(node); - if (!createInfo) - { - return false; - } - - auto regInfoIter = internalRegInfoMap.find(createInfo->kernel_def.get()); - std::shared_ptr internalRegInfo; - if (regInfoIter != internalRegInfoMap.end()) - { - internalRegInfo = regInfoIter->second; - if (internalRegInfo->supportQuery && !internalRegInfo->supportQuery(node)) - { - return false; - } - } - - // Check whether the node uses any data types which are unsupported by the device. - if (!DoesNodeContainSupportedDataTypes(node, internalRegInfo.get(), supportedDeviceDataTypeMask)) - { - return false; - } - - return true; - } - - // Gets properties of the registration for a node void GetRegistrationProperties( const onnxruntime::GraphViewer& graph, @@ -289,68 +151,57 @@ namespace Dml _In_opt_ const std::unordered_map* nodeNameToPartitionMap, _Inout_ std::unordered_map& dmlNodePropertyMap, _Inout_ std::unordered_set& requiredInitializerMap, - _Out_ bool* isDmlNode, _Out_ bool* isDmlGraphNode ) { - *isDmlNode = false; *isDmlGraphNode = false; - // Find the highest priority DML registry supporting this node, and get its highest-priority - // registration. Determine if that registration supports usage as a graph node. + // Get the kernel creation info for the registration, and check if it carries the property + // set during registration of kernels that support DML graph node usage. + auto graphNodeProperty = dmlNodePropertyMap.insert(std::make_pair(&node, GraphNodeProperties())); - if (IsNodeSupportedByDml(node, kernel_lookup, supportedDeviceDataTypeMask, - internalRegInfoMap)) + // Ensure that shape information is known statically for the inputs and outputs of the node, + // which is required for MLGraph compilation. + const onnxruntime::KernelCreateInfo* createInfo = kernel_lookup.LookUpKernel(node); + assert(createInfo != nullptr); // since GetRegistrationProperties is called only when node is a DML node + + auto regInfoIter = internalRegInfoMap.find(createInfo->kernel_def.get()); + if (regInfoIter != internalRegInfoMap.end()) { - *isDmlNode = true; + auto internalRegInfo = regInfoIter->second; - // Get the kernel creation info for the registration, and check if it carries the property - // set during registration of kernels that support DML graph node usage. - auto graphNodeProperty = dmlNodePropertyMap.insert(std::make_pair(&node, GraphNodeProperties())); - - // Ensure that shape information is known statically for the inputs and outputs of the node, - // which is required for MLGraph compilation. - const onnxruntime::KernelCreateInfo* createInfo = kernel_lookup.LookUpKernel(node); - assert(createInfo != nullptr); // since IsNodeSupportedByDml() returned true - - auto regInfoIter = internalRegInfoMap.find(createInfo->kernel_def.get()); - if (regInfoIter != internalRegInfoMap.end()) + if (internalRegInfo && internalRegInfo->graphNodeFactoryRegistration) { - auto internalRegInfo = regInfoIter->second; - - if (internalRegInfo && internalRegInfo->graphNodeFactoryRegistration) + bool requiredCpuInputsConstant = true; + for (uint32_t inputIndex : internalRegInfo->requiredConstantCpuInputs) { - bool requiredCpuInputsConstant = true; - for (uint32_t inputIndex : internalRegInfo->requiredConstantCpuInputs) + if (inputIndex >= node.InputDefs().size() || !node.InputDefs()[inputIndex]->Exists()) { - if (inputIndex >= node.InputDefs().size() || !node.InputDefs()[inputIndex]->Exists()) - { - continue; - } - - const onnx::TensorProto* tensor = nullptr; - const std::string& inputName = node.InputDefs()[inputIndex]->Name(); - - if (!graph.GetInitializedTensor(inputName, tensor)) - { - requiredCpuInputsConstant = false; - break; - } - - requiredInitializerMap.insert(inputName); + continue; } - std::optional requiredInputCount = internalRegInfo->graphNodeFactoryRegistration->requiredInputCount; - if (requiredCpuInputsConstant && - TryGetStaticInputShapes( node, graphNodeProperty.first->second.inputShapes) && - !ContainsEmptyDimensions(graphNodeProperty.first->second.inputShapes, internalRegInfo->requiredConstantCpuInputs) && - TryGetStaticOutputShapes(node, graphNodeProperty.first->second.outputShapes) && - !ContainsEmptyDimensions(graphNodeProperty.first->second.outputShapes, internalRegInfo->requiredConstantCpuInputs) && - (requiredInputCount == std::nullopt || *requiredInputCount == node.InputDefs().size())) + const onnx::TensorProto* tensor = nullptr; + const std::string& inputName = node.InputDefs()[inputIndex]->Name(); + + if (!graph.GetInitializedTensor(inputName, tensor)) { - *isDmlGraphNode = true; - graphNodeProperty.first->second.internalRegInfo = internalRegInfo; + requiredCpuInputsConstant = false; + break; } + + requiredInitializerMap.insert(inputName); + } + + std::optional requiredInputCount = internalRegInfo->graphNodeFactoryRegistration->requiredInputCount; + if (requiredCpuInputsConstant && + TryGetStaticInputShapes( node, graphNodeProperty.first->second.inputShapes) && + !ContainsEmptyDimensions(graphNodeProperty.first->second.inputShapes, internalRegInfo->requiredConstantCpuInputs) && + TryGetStaticOutputShapes(node, graphNodeProperty.first->second.outputShapes) && + !ContainsEmptyDimensions(graphNodeProperty.first->second.outputShapes, internalRegInfo->requiredConstantCpuInputs) && + (requiredInputCount == std::nullopt || *requiredInputCount == node.InputDefs().size())) + { + *isDmlGraphNode = true; + graphNodeProperty.first->second.internalRegInfo = internalRegInfo; } } } @@ -490,78 +341,6 @@ namespace Dml return partition; } - std::unique_ptr ComputationCapacityFromPartition( - GraphPartition* partition, - uint32_t partitionIndex, - const onnxruntime::GraphViewer& graph, - std::unordered_map&& graphNodePropertyMap, - onnxruntime::KernelRegistry* registryForPartitionKernels, - const std::string& partitionKernelPrefix, - std::shared_ptr> transferredInitializerMap) - { - std::unique_ptr subGraph = std::make_unique(); - - if (partition->IsDmlGraphPartition()) - { - assert(partition->IsDmlGraphPartition()); - - // Create a definition for the node. The name must be unique. - auto def = std::make_unique(); - def->name = std::string("DmlFusedNode_") + partitionKernelPrefix + std::to_string(partitionIndex); - def->domain = "DmlFusedNodeDomain"; - def->since_version = 1; - def->inputs.insert(def->inputs.begin(), partition->GetInputs().begin(), partition->GetInputs().end()); - def->outputs.insert(def->outputs.begin(), partition->GetOutputs().begin(), partition->GetOutputs().end()); - - // Populate properties which will be passed to OpKernel for this graph via the function below - std::unordered_map partitionNodePropsMap; - for (auto nodeIndex : partition->GetNodeIndices()) - { - const onnxruntime::Node* node = graph.GetNode(nodeIndex); - -#ifdef PRINT_PARTITON_INFO - printf("Partition %u\t%s\n", partitionIndex, GraphDescBuilder::GetUniqueNodeName(*node).c_str()); -#endif - partitionNodePropsMap.insert(std::make_pair( - GraphDescBuilder::GetUniqueNodeName(*node), std::move(graphNodePropertyMap[node]))); - } - -#ifdef PRINT_PARTITON_INFO - printf("\n"); -#endif - - // These nodeArgNames will be used while creating DML Graph inside FusedGraphKernel.cpp - // Ordering of input/output nodeArgs in below vector will be same as Node::Definitions::input_defs because - // ORT is populating these args as it is while creating the FusedNode at Graph::CreateFusedSubGraphNode() - // Why we need these names? - // After Partitioning and before reaching to FusedGraphKernel, ORT may modify the input/output nodeArg names - // present in FusedNode (Node::Definitions::input_defs) as part of some transformers like memcopy, or L1/L2/L3 transformers. - std::vector fusedNodeInputArgOriginalNames = def->inputs; - std::vector fusedNodeOutputArgOriginalNames = def->outputs; - auto fused_kernel_func = [partitionNodePropsMap, transferredInitializerMap, fusedNodeInputArgOriginalNames, fusedNodeOutputArgOriginalNames](onnxruntime::FuncManager& func_mgr, const onnxruntime::OpKernelInfo& info, std::unique_ptr& out) mutable ->onnxruntime::Status - { - out.reset(CreateFusedGraphKernel(info, partitionNodePropsMap, *transferredInitializerMap, fusedNodeInputArgOriginalNames, fusedNodeOutputArgOriginalNames)); - return Status::OK(); - }; - - // build the kernel definition on the fly, and register it to the fused_kernel_regisitry. - onnxruntime::KernelDefBuilder builder; - - builder.SetName(def->name) - .SetDomain(def->domain) - .SinceVersion(def->since_version) - .Provider(onnxruntime::kDmlExecutionProvider); - - ORT_THROW_IF_ERROR(registryForPartitionKernels->Register(builder, fused_kernel_func)); - - subGraph->SetMetaDef(std::move(def)); - } - - subGraph->nodes = std::move(partition->GetNodeIndices()); - - return std::make_unique(std::move(subGraph)); - } - // Whether any operator in the model contains a subgraph. This is true // if the graph being partitioned is itself within a subgraph, or contains // an operator with a subgraph. @@ -647,25 +426,27 @@ namespace Dml const onnxruntime::Node& node = *graph.GetNode(nodeIndex); // Whether the node is implemented through DML. - bool isDmlNode = false; + bool isDmlNode = node.GetExecutionProviderType() == onnxruntime::kDmlExecutionProvider; // Whether the node is implemented through DML and as a graph node, meaning it // can generate DML operations through a private interface for use as an MLGraph node. bool isDmlGraphNode = false; // Get the registration properties above and populate nodeNameToPartitionMap. - GetRegistrationProperties( - graph, - node, - kernel_lookup, - supportedDeviceDataTypeMask, - internalRegInfoMap, - &nodeNameToPartitionMap, - graphNodePropertyMap, - requiredInitializerMap, - /*out*/ &isDmlNode, - /*out*/ &isDmlGraphNode - ); + if (isDmlNode) + { + GetRegistrationProperties( + graph, + node, + kernel_lookup, + supportedDeviceDataTypeMask, + internalRegInfoMap, + &nodeNameToPartitionMap, + graphNodePropertyMap, + requiredInitializerMap, + /*out*/ &isDmlGraphNode + ); + } // Add a unique partition if graph node usage is not supported. // @@ -739,131 +520,4 @@ namespace Dml return partitions; } - - std::unordered_map> - GetInitializerToPartitionMap( - const onnxruntime::GraphViewer& graph, - gsl::span> partitions - ) - { - std::unordered_map> initializerPartitionMap; - - for (uint32_t partitionIndex = 0; partitionIndex < gsl::narrow_cast(partitions.size()); ++partitionIndex) - { - auto& partition = partitions[partitionIndex]; - - // Skip partitions which have been merged into other partitions - if (partition->GetRootMergedPartition() != partition.get()) - { - continue; - } - - std::unordered_map transferredInitializerMap; - - for (const std::string& input : partition->GetInputs()) - { - const onnx::TensorProto* tensor = nullptr; - if (graph.GetInitializedTensor(input, tensor)) - { - initializerPartitionMap[tensor].push_back(partitionIndex); - } - } - } - - return initializerPartitionMap; - } - - std::vector> - PartitionGraph( - const onnxruntime::GraphViewer& graph, - const InternalRegistrationInfoMap& internalRegInfoMap, - const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup, - uint32_t supportedDeviceDataTypeMask, // Each bit corresponds to each DML_TENSOR_DATA_TYPE. - onnxruntime::KernelRegistry* registryForPartitionKernels, - const std::string& partitionKernelPrefix - ) - { - std::vector> result; - - // Initializers needed by any graph partition - std::unordered_set requiredInitializerMap; - - std::unordered_map graphNodePropertyMap; - std::vector> partitions = BuildPartitions( - graph, - internalRegInfoMap, - kernel_lookup, - supportedDeviceDataTypeMask, - graphNodePropertyMap, - requiredInitializerMap); - - // Create a map between each initialized tensor and the partition(s) it is part of. - auto initializerPartitionMap = GetInitializerToPartitionMap(graph, partitions); - - for (uint32_t partitionIndex = 0; partitionIndex < partitions.size(); ++partitionIndex) - { - auto& partition = partitions[partitionIndex]; - - if (partition->GetRootMergedPartition() != partition.get() || - !partition->IsDmlPartition()) - { - continue; - } - - // Create a map which will store by name each initializer which should be transferred to the - // partition. This prevents OnnxRuntime from allocating GPU resources and uploading those initializers, - // so the partiton's kernel can do so. In the process, it will pre-process weights while consuming a CPU - // backed resource, avoiding an extra set of GPU resources in memory. - // A shared pointer is used so the functor and contained initializer captures can be cheaply copied within ORT. - auto transferredInitializerMap = std::make_shared>(); - - for (const auto& input : partition->GetInputs()) - { - if (partition->IsDmlGraphPartition()) - { - const onnx::TensorProto* tensor = nullptr; - if (graph.GetInitializedTensor(input, tensor)) - { - // It's only safe to transfer tensors which are used by this partition alone. - auto iter = initializerPartitionMap.find(tensor); - assert(iter != initializerPartitionMap.end()); - if (iter->second.size() > 1) - { - if (requiredInitializerMap.find(input) != requiredInitializerMap.end()) - { - // The kernel relies on this input to be initialized, and it should be small enough to copy - // cheaply. FusedGraphKernel only handles constant CPU inputs through transferred initializers, - // rather than ORT, to avoid mismatches in policy or implementation causing failures. - (*transferredInitializerMap)[input] = const_cast(*tensor); - } - - continue; - } - - // Transfer the initializer - auto& graphTensor = const_cast(*tensor); - - onnx::TensorProto partitionTensor; - graphTensor.Swap(&partitionTensor); - (*transferredInitializerMap)[input] = std::move(partitionTensor); - - const_cast(graph.GetAllInitializedTensors()).erase(graph.GetAllInitializedTensors().find(input)); - } - } - } - - result.push_back(ComputationCapacityFromPartition( - partition.get(), - partitionIndex, - graph, - std::move(graphNodePropertyMap), - registryForPartitionKernels, - partitionKernelPrefix, - transferredInitializerMap - )); - } - - return result; - } - } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.h index b82999caff..c5b15c7b1c 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphPartitioner.h @@ -49,22 +49,4 @@ namespace Dml std::unordered_map& graphNodePropertyMap, std::unordered_set& requiredInitializerMap, std::function onNodeUnsupportedInGraph = nullptr); - - std::vector> - PartitionGraph( - const onnxruntime::GraphViewer& graph, - const Windows::AI::MachineLearning::Adapter::InternalRegistrationInfoMap& internalRegInfoMap, - const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup, - uint32_t supportedDeviceDataTypeMask, // Each bit corresponds to each DML_TENSOR_DATA_TYPE. - onnxruntime::KernelRegistry* registryForPartitionKernels, - const std::string& partitionKernelPrefix - ); - - bool IsNodeSupportedByDml( - const onnxruntime::Node& node, - const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup, - uint32_t supportedDeviceDataTypeMask, // Each bit corresponds to each DML_TENSOR_DATA_TYPE. - const Windows::AI::MachineLearning::Adapter::InternalRegistrationInfoMap& internalRegInfoMap - ); - } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.cpp index e1076ab235..09922310b5 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.cpp @@ -17,15 +17,6 @@ namespace Dml { - GraphTransformer::GraphTransformer( - const std::string& name, - const onnxruntime::IExecutionProvider* provider - ) - : onnxruntime::GraphTransformer(name), - m_providerImpl(static_cast(provider)->GetImpl()) - { - } - onnxruntime::common::Status GraphTransformer::ApplyImpl( onnxruntime::Graph& graph, bool& modified, @@ -39,10 +30,6 @@ namespace Dml PerformOperatorFusion(&graph, &transformModifiedGraph); modified |= transformModifiedGraph; - transformModifiedGraph = false; - PerformQuantizedOperatorDecomposition(&graph, &transformModifiedGraph); - modified |= transformModifiedGraph; - if (modified) { ORT_RETURN_IF_ERROR(graph.Resolve()); @@ -84,24 +71,10 @@ namespace Dml // graph while iterating over it std::vector nodesToAdd; - onnxruntime::ProviderType provider_type = onnxruntime::kDmlExecutionProvider; - const gsl::not_null registry = m_providerImpl->GetKernelRegistry().get(); - const auto kernel_type_str_resolver = onnxruntime::OpSchemaKernelTypeStrResolver{}; - const auto kernel_lookup = onnxruntime::KernelLookup{provider_type, - gsl::make_span(®istry, 1), - kernel_type_str_resolver}; - for (auto& node : graph->Nodes()) { - // We need to predict whether the nodes will be assigned to the DML transformer by Lotus, - // which occurs in IExecutionProvider::GetCapability. - - if (!IsNodeSupportedByDml( - node, - kernel_lookup, - m_providerImpl->GetSupportedDeviceDataTypeMask(), - *m_providerImpl->GetInternalRegistrationInfoMap().get() - )) + // Ignore the nodes which were not assigned to the DML by ORT during IExecutionProvider::GetCapability() + if (!onnxruntime::graph_utils::IsSupportedProvider(node, {onnxruntime::kDmlExecutionProvider})) { // Can't fuse nodes that don't belong to this execution provider continue; @@ -119,9 +92,8 @@ namespace Dml const auto& outputNode = *node.OutputNodesBegin(); - // We need to predict whether the nodes will be assigned to the DML transformer by Lotus, - // which occurs in IExecutionProvider::GetCapability. - if (!kernel_lookup.LookUpKernel(outputNode)) + // Can't fuse if outputNode was not assigned to the DML by ORT during IExecutionProvider::GetCapability() + if (!onnxruntime::graph_utils::IsSupportedProvider(outputNode, GetCompatibleExecutionProviders())) { // Can't fuse nodes that don't belong to this execution provider continue; @@ -207,7 +179,8 @@ namespace Dml nodeToAdd.outputs, &nodeToAdd.attributes, nodeToAdd.domain); - + + node.SetExecutionProviderType(onnxruntime::kDmlExecutionProvider); // Add a dynamic attribute to the fuseable operator to specify activation node.AddAttribute(AttrName::FusedActivation, nodeToAdd.activationOpType); node.AddAttribute(AttrName::FusedActivationDomain, nodeToAdd.activationOpDomain); @@ -224,103 +197,4 @@ namespace Dml } } - // Converts certain QLinear operations unsupported by the DML API into a sequence of DeQuantizeLinear, 32-bit operator, QuantizeLinear - void GraphTransformer::PerformQuantizedOperatorDecomposition(onnxruntime::Graph* graph, bool* modified) const - { - struct NodeToAdd - { - std::string name; - std::string description; - std::string opType; - std::string domain; - onnxruntime::NodeAttributes attributes; - std::vector inputs; - std::vector outputs; - }; - - // Defer adding and removing nodes in the graph until after we're done iterating over it, because we can't mutate the - // graph while iterating over it - std::vector nodesToAdd; - std::vector nodesToRemove; - - for (auto& node : graph->Nodes()) - { - // For now, only QLinearSigmoid is handled - if (node.Domain() == onnxruntime::kMSDomain && - node.OpType() == "QLinearSigmoid") - { - // Intermediate node arg type proto with floating point format - onnx::TypeProto floatTensorProto; - floatTensorProto.mutable_tensor_type()->set_elem_type(onnx::TensorProto_DataType_FLOAT); - - // Add intermediate graph edges for the input and output of the FP32 sigmoid operator - auto* sigmoidInputArg = &graph->GetOrCreateNodeArg("decomposed_QLinearSigmoid_input_" + GetUniqueNodeName(&node), &floatTensorProto); - auto* sigmoidOutputArg = &graph->GetOrCreateNodeArg("decomposed_QLinearSigmoid_output_" + GetUniqueNodeName(&node), &floatTensorProto); - - { - NodeToAdd dequantizeNode; - dequantizeNode.name = "decomposed_QLinearSigmoid_DequantizeLinear_" + GetUniqueNodeName(&node); - dequantizeNode.description = ""; - dequantizeNode.opType = "DequantizeLinear"; - dequantizeNode.domain = ""; - - dequantizeNode.inputs.push_back(graph->GetNodeArg(node.InputDefs()[0]->Name())); - dequantizeNode.inputs.push_back(graph->GetNodeArg(node.InputDefs()[1]->Name())); - dequantizeNode.inputs.push_back(graph->GetNodeArg(node.InputDefs()[2]->Name())); - dequantizeNode.outputs.push_back(sigmoidInputArg); - - nodesToAdd.push_back(std::move(dequantizeNode)); - } - - { - NodeToAdd sigmoidNode; - sigmoidNode.name = "decomposed_QLinearSigmoid_Sigmoid_" + GetUniqueNodeName(&node); - sigmoidNode.description = ""; - sigmoidNode.opType = "Sigmoid"; - sigmoidNode.domain = ""; - sigmoidNode.inputs.push_back(sigmoidInputArg); - sigmoidNode.outputs.push_back(sigmoidOutputArg); - nodesToAdd.push_back(std::move(sigmoidNode)); - } - - { - NodeToAdd quantizeNode; - quantizeNode.name = "decomposed_QLinearSigmoid_QuantizeLinear_" + GetUniqueNodeName(&node); - quantizeNode.description = ""; - quantizeNode.opType = "QuantizeLinear"; - quantizeNode.domain = ""; - - quantizeNode.inputs.push_back(sigmoidOutputArg); - quantizeNode.inputs.push_back(graph->GetNodeArg(node.InputDefs()[3]->Name())); - quantizeNode.inputs.push_back(graph->GetNodeArg(node.InputDefs()[4]->Name())); - quantizeNode.outputs.push_back(graph->GetNodeArg(node.OutputDefs()[0]->Name())); - - nodesToAdd.push_back(std::move(quantizeNode)); - } - - nodesToRemove.push_back(node.Index()); - *modified = true; - } - } - - for (auto& nodeToAdd : nodesToAdd) - { - graph->AddNode( - nodeToAdd.name, - nodeToAdd.opType, - nodeToAdd.description, - nodeToAdd.inputs, - nodeToAdd.outputs, - &nodeToAdd.attributes, - nodeToAdd.domain); - } - - for (const auto& nodeIndex : nodesToRemove) - { - onnxruntime::Node* node = graph->GetNode(nodeIndex); - onnxruntime::graph_utils::RemoveNodeOutputEdges(*graph, *node); - graph->RemoveNode(node->Index()); - } - } - } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h index 4ce7064e25..dfcb3bd835 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h @@ -11,7 +11,6 @@ namespace Dml { - class ExecutionProviderImpl; // Applies transforms to a Lotus graph. The graph transformer is responsible for setting the execution provider // on the graph nodes which DML supports. @@ -19,19 +18,16 @@ namespace Dml { public: GraphTransformer( - const std::string& name, - const onnxruntime::IExecutionProvider* provider - ); + const std::string& name + ) : onnxruntime::GraphTransformer(name) + { + } private: onnxruntime::common::Status ApplyImpl(onnxruntime::Graph& graph, bool& modified, int graph_level, const onnxruntime::logging::Logger& logger) const final; private: void PerformOperatorFusion(onnxruntime::Graph* graph, bool* modified) const; - void PerformQuantizedOperatorDecomposition(onnxruntime::Graph* graph, bool* modified) const; - - std::shared_ptr m_registry; - const ExecutionProviderImpl* m_providerImpl = nullptr; }; } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp index bd7aa3f0d6..3ae29629ef 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp @@ -55,7 +55,7 @@ namespace Dml std::vector inputEdges; for (uint32_t inputIndex = 0; inputIndex < m_kernelInputIndices.size(); inputIndex++) { - if (m_kernelInputIndices[inputIndex].has_value()) + if (m_kernelInputIndices[inputIndex].has_value()) { DML_INPUT_GRAPH_EDGE_DESC inputEdge = {}; inputEdge.GraphInputIndex = *m_kernelInputIndices[inputIndex]; @@ -67,11 +67,11 @@ namespace Dml operatorGraphDesc.inputEdgeCount = gsl::narrow_cast(inputEdges.size()); operatorGraphDesc.inputEdges = inputEdges.data(); - + std::vector outputEdges; for (uint32_t outputIndex = 0; outputIndex < m_kernelOutputIndices.size(); outputIndex++) { - if (m_kernelOutputIndices[outputIndex].has_value()) + if (m_kernelOutputIndices[outputIndex].has_value()) { DML_OUTPUT_GRAPH_EDGE_DESC outputEdge = {}; outputEdge.FromNodeIndex = 0; @@ -111,6 +111,105 @@ namespace Dml } } + void DmlOperator::SetDmlOperatorGraphDesc( + const MLOperatorGraphDesc&& operatorGraphDesc, + const MLOperatorKernelCreationContext& kernelInfo + ) + { + // Initialize should only be called once. + assert(m_compiledOperator == nullptr); + + // DML doesn't support empty tensors. If an operator is still executable with empty tensors, the empty tensors + // should be removed or massaged depending on the definition. + for (const TensorDesc& desc : m_inputTensorDescs) + { + if (OperatorHelper::ContainsEmptyDimensions(desc.GetSizes())) + { + return; + } + } + + for (const TensorDesc& desc : m_outputTensorDescs) + { + if (OperatorHelper::ContainsEmptyDimensions(desc.GetSizes())) + { + return; + } + } + + // m_kernelInputIndices should be identity + for (uint32_t idx = 0; idx < m_kernelInputIndices.size(); idx++) + { + if (m_kernelInputIndices[idx] == std::nullopt || !kernelInfo.IsInputValid(*m_kernelInputIndices[idx])) + { + continue; + } + assert(m_kernelInputIndices[idx] == idx); + } + + // m_kernelOutputIndices should be identity + for (uint32_t idx = 0; idx < m_kernelOutputIndices.size(); idx++) + { + if (m_kernelOutputIndices[idx] == std::nullopt || !kernelInfo.IsOutputValid(*m_kernelOutputIndices[idx])) + { + continue; + } + assert(m_kernelOutputIndices[idx] == idx); + } + + ComPtr contextPrivate; + ORT_THROW_IF_FAILED(kernelInfo.GetInterface()->QueryInterface(contextPrivate.GetAddressOf())); + if (contextPrivate->IsDmlGraphNode()) + { + ORT_THROW_IF_FAILED(contextPrivate->SetDmlOperator(&operatorGraphDesc)); + } + else + { + DML_GRAPH_DESC graphDesc = {}; + std::vector dmlGraphNodes(operatorGraphDesc.nodeCount); + std::vector> dmlOperators(operatorGraphDesc.nodeCount); + std::vector dmlOperatorGraphNodes(operatorGraphDesc.nodeCount); + std::vector dmlInputEdges(operatorGraphDesc.inputEdgeCount); + std::vector dmlOutputEdges(operatorGraphDesc.outputEdgeCount); + std::vector dmlIntermediateEdges(operatorGraphDesc.intermediateEdgeCount); + + // DML Graph validator will check the validity of the graph. No need to check here. + ConvertToDmlGraphDesc(operatorGraphDesc, + graphDesc, + dmlOperators, + dmlOperatorGraphNodes, + dmlGraphNodes, + dmlInputEdges, + dmlOutputEdges, + dmlIntermediateEdges); + + // compile the graph and create IDMLCompiledOperator + Microsoft::WRL::ComPtr dmlDevice1; + DMLX_THROW_IF_FAILED(m_dmlDevice->QueryInterface(IID_PPV_ARGS(&dmlDevice1))); + DML_EXECUTION_FLAGS executionFlags = GetExecutionFlags(); + ORT_THROW_IF_FAILED(dmlDevice1->CompileGraph(&graphDesc, executionFlags, IID_PPV_ARGS(&m_compiledOperator))); + + UINT64 persistentResourceSize = m_compiledOperator->GetBindingProperties().PersistentResourceSize; + if (persistentResourceSize > 0) + { + ORT_THROW_IF_FAILED(m_executionProvider->AllocatePooledResource( + static_cast(persistentResourceSize), + AllocatorRoundingMode::Enabled, + m_persistentResource.GetAddressOf(), + m_persistentResourcePoolingUnk.GetAddressOf())); + + m_persistentResourceBinding = DML_BUFFER_BINDING{ m_persistentResource.Get(), 0, persistentResourceSize }; + } + + std::vector initializationInputBindings(m_kernelInputIndices.size()); + + ORT_THROW_IF_FAILED(m_executionProvider->InitializeOperator( + m_compiledOperator.Get(), + m_persistentResourceBinding ? &*m_persistentResourceBinding : nullptr, + gsl::make_span(initializationInputBindings))); + } + } + void DmlOperator::SetDmlOperatorDesc( const DML_OPERATOR_DESC& operatorDesc, const MLOperatorKernelContext& kernelInfo @@ -596,4 +695,52 @@ namespace Dml ); } + void DmlOperator::ConvertToDmlGraphDesc(const MLOperatorGraphDesc& operatorGraphDesc, + _Out_ DML_GRAPH_DESC& graphDesc, + _Inout_ std::vector>& dmlOperators, + _Inout_ std::vector& dmlOperatorGraphNodes, + _Inout_ std::vector& dmlGraphNodes, + _Inout_ std::vector& dmlInputEdges, + _Inout_ std::vector& dmlOutputEdges, + _Inout_ std::vector& dmlIntermediateEdges) + { + graphDesc.InputCount = gsl::narrow_cast(m_kernelInputIndices.size()); + graphDesc.OutputCount = gsl::narrow_cast(m_kernelOutputIndices.size()); + + // set the graph nodes + graphDesc.NodeCount = operatorGraphDesc.nodeCount; + for (size_t i = 0; i < graphDesc.NodeCount; ++i) + { + // Create the operator. + ORT_THROW_IF_FAILED(m_dmlDevice->CreateOperator(operatorGraphDesc.nodesAsOpDesc[i], IID_PPV_ARGS(&dmlOperators[i]))); + dmlOperatorGraphNodes[i] = DML_OPERATOR_GRAPH_NODE_DESC{dmlOperators[i].Get()}; + dmlGraphNodes[i] = DML_GRAPH_NODE_DESC{DML_GRAPH_NODE_TYPE_OPERATOR, &dmlOperatorGraphNodes[i]}; + } + graphDesc.Nodes = dmlGraphNodes.data(); + + // set the input edges + graphDesc.InputEdgeCount = operatorGraphDesc.inputEdgeCount; + for (size_t i = 0; i < operatorGraphDesc.inputEdgeCount; ++i) + { + dmlInputEdges[i] = DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_INPUT, &operatorGraphDesc.inputEdges[i]}; + } + graphDesc.InputEdges = dmlInputEdges.data(); + + // set the output edges + graphDesc.OutputEdgeCount = operatorGraphDesc.outputEdgeCount; + for (size_t i = 0; i < operatorGraphDesc.outputEdgeCount; ++i) + { + dmlOutputEdges[i] = DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_OUTPUT, &operatorGraphDesc.outputEdges[i]}; + } + graphDesc.OutputEdges = dmlOutputEdges.data(); + + // set the intermediate edges + graphDesc.IntermediateEdgeCount = operatorGraphDesc.intermediateEdgeCount; + for (size_t i = 0; i < operatorGraphDesc.intermediateEdgeCount; ++i) + { + dmlIntermediateEdges[i] = DML_GRAPH_EDGE_DESC{DML_GRAPH_EDGE_TYPE_INTERMEDIATE, &operatorGraphDesc.intermediateEdges[i]}; + } + graphDesc.IntermediateEdges = dmlIntermediateEdges.data(); + } + } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h index 8efa8d2f82..d398233e4c 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h @@ -63,6 +63,15 @@ namespace Dml const MLOperatorKernelCreationContext& kernelInfo ); + // This method only works with DML_GRAPH. + // To make it work without DML_GRAPH, we need to add new functionality + // in DMLX i.e. DMLX should also give access to DML_OPERATOR_DESC + // rather than IDMLOperator. + void SetDmlOperatorGraphDesc( + const MLOperatorGraphDesc&& operatorGraphDesc, + const MLOperatorKernelCreationContext& kernelInfo + ); + void SetDmlOperatorDesc( const DML_OPERATOR_DESC& operatorDesc, const MLOperatorKernelContext& kernelInfo @@ -122,6 +131,16 @@ namespace Dml // kernel. Entries for unused DML inputs are nullopt. std::vector> m_kernelInputIndices; std::vector> m_kernelOutputIndices; + + void ConvertToDmlGraphDesc(const MLOperatorGraphDesc& operatorGraphDesc, + _Out_ DML_GRAPH_DESC& graphDesc, + _Inout_ std::vector>& dmlOperators, + _Inout_ std::vector& dmlOperatorGraphNodes, + _Inout_ std::vector& dmlGraphNodes, + _Inout_ std::vector& dmlInputEdges, + _Inout_ std::vector& dmlOutputEdges, + _Inout_ std::vector& dmlIntermediateEdges); + }; } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorQLinearSigmoid.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorQLinearSigmoid.cpp new file mode 100644 index 0000000000..9c67e1e26f --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorQLinearSigmoid.cpp @@ -0,0 +1,173 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "precomp.h" + +namespace Dml +{ +// QLinearSigmoid = Dequantize + Sigmoid + Quantize +// This kernel is the first usage of graph based implementation +class DmlOperatorQLinearSigmoid : public DmlOperator +{ + // This order matches the ONNX schema. + enum OnnxInputIndex + { + X, // Input + X_scale, + X_zero_point, + Y_scale, + Y_zero_point, + Count, + }; + +public: + DmlOperatorQLinearSigmoid(const MLOperatorKernelCreationContext& kernelCreationContext) + : DmlOperator(kernelCreationContext) + { + DmlOperator::Initialize(kernelCreationContext); + + std::vector outputShape = kernelCreationContext.GetTensorShapeDescription().GetOutputTensorShape(0); + const uint32_t outputShapeDimCount = gsl::narrow_cast(outputShape.size()); + + uint32_t axis = 0; + // Explicitly reshape each of the inputs after the first input (scale and zero point tensors). + for (uint32_t index = OnnxInputIndex::X_scale, inputCount = gsl::narrow_cast(OnnxInputIndex::Count); index < inputCount; ++index) + { + auto edgeDesc = kernelCreationContext.GetInputEdgeDescription(index); + assert(edgeDesc.edgeType == MLOperatorEdgeType::Tensor); + + // Fix up the the tensor shape by filling with trailing ones. So input[2,3] with axis=0 and scale[2] + // becomes scale[2,1], so that broadcasting works correctly. + std::vector inputTensorShape = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(index); + + // If the input tensor is a 1D vector, then extra massaging is needed to project their + // 1D vectors back to the full shape for broadcasting along the given axis. + // The 1D vector should have a length equal to the output tensor's dimension on that axis. + if (inputTensorShape.size() == 1 && inputTensorShape != outputShape) + { + ML_CHECK_VALID_ARGUMENT(axis < outputShapeDimCount); + uint32_t broadcastAxisLength = outputShape[axis]; + ML_CHECK_VALID_ARGUMENT(inputTensorShape[0] == broadcastAxisLength); + inputTensorShape.insert(inputTensorShape.begin(), axis, 1); + inputTensorShape.insert(inputTensorShape.end(), outputShapeDimCount - 1 - axis, 1); + } + // For any other shape (scalar/ND), leave it alone, and the TensorDesc constructor + // will apply broadcasting with standard elementwise alignment. + + m_inputTensorDescs[index] = TensorDesc( + edgeDesc.tensorDataType, + gsl::make_span(outputShape), + gsl::make_span(inputTensorShape), + TensorAxis::DoNotCoerce, + TensorAxis::W, + TensorAxis::RightAligned, + NchwDimensionCount, // minDimensionCount + 0 // guaranteedBaseOffsetAlignment + ); + } + + // 1. output edge between Dequantize and Sigmoid node + // 2. input edge between Sigmoid and Quantize node + TensorDesc intermediateOutputTensorDesc = TensorDesc( + MLOperatorTensorDataType::Float, + gsl::make_span(outputShape), + gsl::make_span(outputShape), + TensorAxis::DoNotCoerce, + TensorAxis::W, + TensorAxis::RightAligned, + NchwDimensionCount, // minDimensionCount + 0 // guaranteedBaseOffsetAlignment + ); + + std::vector inputDescs = GetDmlInputDescs(); + std::vector outputDescs = GetDmlOutputDescs(); + + DML_ELEMENT_WISE_DEQUANTIZE_LINEAR_OPERATOR_DESC dequantizeOperatorDesc = {}; + dequantizeOperatorDesc.InputTensor = &inputDescs[OnnxInputIndex::X]; + dequantizeOperatorDesc.ScaleTensor = &inputDescs[OnnxInputIndex::X_scale]; + dequantizeOperatorDesc.ZeroPointTensor = &inputDescs[OnnxInputIndex::X_zero_point]; + dequantizeOperatorDesc.OutputTensor = &intermediateOutputTensorDesc.GetDmlDesc(); + + const DML_OPERATOR_DESC opDesc1{DML_OPERATOR_ELEMENT_WISE_DEQUANTIZE_LINEAR, &dequantizeOperatorDesc}; + + DML_ACTIVATION_SIGMOID_OPERATOR_DESC sigmoidOperatorDesc = {}; + sigmoidOperatorDesc.InputTensor = dequantizeOperatorDesc.OutputTensor; + sigmoidOperatorDesc.OutputTensor = dequantizeOperatorDesc.OutputTensor; + const DML_OPERATOR_DESC opDesc2{DML_OPERATOR_ACTIVATION_SIGMOID, &sigmoidOperatorDesc}; + + DML_ELEMENT_WISE_QUANTIZE_LINEAR_OPERATOR_DESC quantizeOperatorDesc = {}; + quantizeOperatorDesc.InputTensor = sigmoidOperatorDesc.OutputTensor; + quantizeOperatorDesc.ScaleTensor = &inputDescs[OnnxInputIndex::Y_scale]; + quantizeOperatorDesc.ZeroPointTensor = &inputDescs[OnnxInputIndex::Y_zero_point]; + quantizeOperatorDesc.OutputTensor = &outputDescs[0]; + const DML_OPERATOR_DESC opDesc3{DML_OPERATOR_ELEMENT_WISE_QUANTIZE_LINEAR, &quantizeOperatorDesc}; + + MLOperatorGraphDesc operatorGraphDesc = {}; + operatorGraphDesc.nodeCount = 3; + std::vector opDescs{&opDesc1, &opDesc2, &opDesc3}; + operatorGraphDesc.nodesAsOpDesc = opDescs.data(); + + // set input edges + std::pair nodeToNodeInputIndex[5] {{0, 0}, {0, 1}, {0, 2}, {2, 1}, {2, 2}}; + std::vector inputEdges(OnnxInputIndex::Count); + for (uint32_t inputIndex = 0; inputIndex < OnnxInputIndex::Count; inputIndex++) + { + DML_INPUT_GRAPH_EDGE_DESC inputEdge = {}; + inputEdge.GraphInputIndex = inputIndex; // OnnxInputIndex and DmlInputIndex are identity for QLinearSigmoid + inputEdge.ToNodeIndex = nodeToNodeInputIndex[inputIndex].first; + inputEdge.ToNodeInputIndex = nodeToNodeInputIndex[inputIndex].second; + inputEdges[inputIndex] = inputEdge; + } + operatorGraphDesc.inputEdgeCount = gsl::narrow_cast(inputEdges.size()); + operatorGraphDesc.inputEdges = inputEdges.data(); + + // set intermediate edges + std::vector intermediateEdges; + + DML_INTERMEDIATE_GRAPH_EDGE_DESC dequantizeToSigmoidEdge = {}; + dequantizeToSigmoidEdge.FromNodeIndex = 0; + dequantizeToSigmoidEdge.FromNodeOutputIndex = 0; + dequantizeToSigmoidEdge.ToNodeIndex = 1; + dequantizeToSigmoidEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(dequantizeToSigmoidEdge); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC sigmoidToQuantizeEdge = {}; + sigmoidToQuantizeEdge.FromNodeIndex = 1; + sigmoidToQuantizeEdge.FromNodeOutputIndex = 0; + sigmoidToQuantizeEdge.ToNodeIndex = 2; + sigmoidToQuantizeEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(sigmoidToQuantizeEdge); + + operatorGraphDesc.intermediateEdgeCount = gsl::narrow_cast(intermediateEdges.size()); + operatorGraphDesc.intermediateEdges = intermediateEdges.data(); + + // set the output edges + std::vector outputEdges; + DML_OUTPUT_GRAPH_EDGE_DESC outputEdge = {}; + outputEdge.FromNodeIndex = 2; + outputEdge.FromNodeOutputIndex = 0; + outputEdge.GraphOutputIndex = 0; + outputEdges.push_back(outputEdge); + operatorGraphDesc.outputEdgeCount = gsl::narrow_cast(outputEdges.size()); + operatorGraphDesc.outputEdges = outputEdges.data(); + + SetDmlOperatorGraphDesc(std::move(operatorGraphDesc), kernelCreationContext); + } +}; + +void CALLBACK QueryQLinearSigmoid(IMLOperatorSupportQueryContextPrivate* context, /*out*/ bool* isSupported) +{ + *isSupported = false; + // Right now the contract is if optional input tensors (like x_zero_point, y_zero_point) are + // not present, then fallback to CPU because DML Quantize_Linear and Dequantize_Linear does not support + // optionality of the zero_point tensor. BUG: https://microsoft.visualstudio.com/OS/_queries/edit/41599005 + if (context->GetInputCount() < 5) + { + return; + } + + *isSupported = true; +} + +DML_OP_DEFINE_CREATION_FUNCTION(QLinearSigmoid, DmlOperatorQLinearSigmoid); +} // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp index 7ddbd39398..04c4a6eacc 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -220,6 +220,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(DmlFusedAdd); DML_OP_EXTERN_CREATION_FUNCTION(DmlFusedSum); DML_OP_EXTERN_CREATION_FUNCTION(QuantizeLinear); DML_OP_EXTERN_CREATION_FUNCTION(DequantizeLinear); +DML_OP_EXTERN_CREATION_FUNCTION(QLinearSigmoid); DML_OP_EXTERN_CREATION_FUNCTION(Sign); DML_OP_EXTERN_CREATION_FUNCTION(IsNaN); DML_OP_EXTERN_CREATION_FUNCTION(Sinh); @@ -270,6 +271,7 @@ DML_OP_EXTERN_QUERY_FUNCTION(RecurrentNeuralNetwork); DML_OP_EXTERN_QUERY_FUNCTION(BatchNormalization); DML_OP_EXTERN_QUERY_FUNCTION(Pad); DML_OP_EXTERN_QUERY_FUNCTION(LayerNormalization); +DML_OP_EXTERN_QUERY_FUNCTION(QLinearSigmoid); constexpr static std::array typeNameListDefault = {"T"}; constexpr static std::array typeNameListTwo = { "T1", "T2" }; @@ -330,6 +332,7 @@ constexpr static std::array supportedTypeListIntege constexpr static std::array supportedTypeListRoiAlign = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Int32|SupportedTensorDataTypes::Int64 }; constexpr static std::array supportedTypeListArgMinMax = {SupportedTensorDataTypes::Float16to32|SupportedTensorDataTypes::Ints8to64}; constexpr static std::array supportedTypeListLayerNormalization = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Float32}; +constexpr static std::array supportedTypeListQLinearSigmoid = {SupportedTensorDataTypes::UInt8 | SupportedTensorDataTypes::Int8}; constexpr static std::array supportedTypeListQLinearMatMul = { SupportedTensorDataTypes::Int8|SupportedTensorDataTypes::UInt8, @@ -699,6 +702,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation // Contrib operators {REG_INFO_MS( 1, Gelu, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO_MS( 1, FusedMatMul, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, + {REG_INFO_MS( 1, QLinearSigmoid, typeNameListDefault, supportedTypeListQLinearSigmoid, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryQLinearSigmoid)}, {REG_INFO( 10, IsInf, typeNameListTwo, supportedTypeListIsInf, DmlGraphSupport::Supported)}, {REG_INFO( 10, Mod, typeNameListDefault, supportedTypeListNumericDefault, DmlGraphSupport::Supported)}, diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index f14ef30e72..990bac4b19 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -1461,6 +1461,7 @@ using ShapeInferenceHelper_Atan = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_Affine = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_QuantizeLinear = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_DequantizeLinear = GetOutputShapeAsInputShapeHelper; +using ShapeInferenceHelper_QLinearSigmoid = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_Sign = GetBroadcastedOutputShapeHelper; using ShapeInferenceHelper_IsNaN = GetBroadcastedOutputShapeHelper; using ShapeInferenceHelper_Erf = GetBroadcastedOutputShapeHelper; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h index 0904312542..3e2096d41d 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -388,6 +388,7 @@ namespace OperatorHelper static const int sc_sinceVer_QLinearAdd = 1; static const int sc_sinceVer_Gelu = 1; static const int sc_sinceVer_FusedMatMul = 1; + static const int sc_sinceVer_QLinearSigmoid = 1; } // namespace MsftOperatorSet1 } // namespace OperatorHelper diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 322fbac0be..89315c0e7c 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -49,6 +49,7 @@ #include "core/providers/cpu/controlflow/utils.h" #include "core/providers/cpu/cpu_execution_provider.h" #ifdef USE_DML // TODO: This is necessary for the workaround in TransformGraph +#include "core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionTransformer.h" #include "core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h" #endif #include "core/session/environment.h" @@ -899,23 +900,6 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, ORT_RETURN_IF_ERROR_SESSIONID_( graph_transformer_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *session_logger_)); -#ifdef USE_DML - // TODO: this is a temporary workaround to apply the DML EP's custom graph transformer prior to partitioning. This - // transformer applies DML-specific fusions that go beyond what ORT offers by default. Ideally the DML EP should - // apply these transforms during partitioning, but the full mutable Graph object isn't exposed to - // IExecutionProvider::GetCapability, which is necessary for the DML EP's transforms. - // - // To prevent this from interfering with other EPs, we only apply this transform if the DML EP is the only one that's - // registered (aside from the CPU EP, which is always registered by default.) - if (execution_providers_.Get(kDmlExecutionProvider) && execution_providers_.NumProviders() <= 2) { - Dml::GraphTransformer dml_transformer(onnxruntime::kDmlExecutionProvider, - execution_providers_.Get(kDmlExecutionProvider)); - - bool modified = false; - ORT_RETURN_IF_ERROR_SESSIONID_(dml_transformer.Apply(graph, modified, *session_logger_)); - } -#endif - // if saving model to ORT format we only assign nodes a custom EP can handle and don't compile them. // we do this to preserve the original nodes in the model but prevent optimizers from changing them. // at runtime, the ORT format model will re-do the partitioning/compilation of these nodes, which may change @@ -1376,6 +1360,24 @@ common::Status InferenceSession::Initialize() { minimal_build_optimization_handling, record_runtime_optimization_produced_op_schema)); +#ifdef USE_DML + if (execution_providers_.Get(kDmlExecutionProvider)) { + std::unique_ptr dmlGraphFusionTransformer = std::make_unique("DmlGraphFusionTransformer", + execution_providers_.Get(kDmlExecutionProvider)); + if (dmlGraphFusionTransformer == nullptr) { + return Status(common::ONNXRUNTIME, common::FAIL, "DmlGraphFusionTransformer is nullptr"); + } + ORT_RETURN_IF_ERROR_SESSIONID_(graph_transformation_mgr_.Register(std::move(dmlGraphFusionTransformer), onnxruntime::TransformerLevel::Level3)); + + // This transformer applies DML-specific fusions that go beyond what ORT offers by default + std::unique_ptr dmlOperatorFusionTransformer = std::make_unique("DmlOperatorFusionTransformer"); + if (dmlOperatorFusionTransformer == nullptr) { + return Status(common::ONNXRUNTIME, common::FAIL, "DmlOperatorFusionTransformer is nullptr"); + } + ORT_RETURN_IF_ERROR_SESSIONID_(graph_transformation_mgr_.Register(std::move(dmlOperatorFusionTransformer), onnxruntime::TransformerLevel::Level2)); + } +#endif + // apply any transformations to the main graph and any subgraphs ORT_RETURN_IF_ERROR_SESSIONID_(TransformGraph(graph, graph_transformation_mgr_, execution_providers_, kernel_registry_manager_,