mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
[DML EP] Re-architect | Partitioning as Transformer (#13131)
### Description
Re-architect DML EP to allow ORT L2/L3 transformers. This change
includes:
- During ORT graph partitioning, DML EP will only set the
dmlExecutionProvider to all eligible nodes.
- Moved DML specific operator transformer as L2 transformer
- Introduced a new DMLGraphFusionTransformer, applicable only for DML
EP, which is responsible to
- partition the graph
- fuse each partition into a IDMLCompiledOperator
- register the kernel for each partition
### Motivation and Context
- Why is this change required? What problem does it solve?
It enables ORT L2/L3 transformers for DML EP, which will increase the
perf of Transformer-based models.
- If it fixes an open issue, please link to the issue here. N/A
Co-authored-by: Sumit Agarwal <sumitagarwal@microsoft.com>
This commit is contained in:
sumitsays
GitHub
parent
38906625a3
commit
e01a8519e0
23 changed files with 1422 additions and 1141 deletions
|
|
@ -0,0 +1,499 @@
|
|||
#pragma once
|
||||
|
||||
#include "DmlGraphFusionHelper.h"
|
||||
|
||||
|
||||
namespace Dml
|
||||
{
|
||||
namespace DmlGraphFusionHelper
|
||||
{
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource>
|
||||
CreateResource(
|
||||
const ExecutionProviderImpl* provider,
|
||||
const std::byte* tensorPtr,
|
||||
size_t tensorByteSize)
|
||||
{
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource> 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<uint64_t>((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<ID3D12Device> 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<ID3D12Resource>
|
||||
CreateCpuResource(
|
||||
const ExecutionProviderImpl* provider,
|
||||
const std::byte* tensorPtr,
|
||||
size_t tensorByteSize)
|
||||
{
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource> 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<uint64_t>((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<ID3D12Device> 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<IUnknown*>(const_cast<void*>(tensor->DataRaw()));
|
||||
Microsoft::WRL::ComPtr<IUnknown> 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<uint8_t>& isInputsUploadedByDmlEP,
|
||||
std::vector<DML_INPUT_GRAPH_EDGE_DESC>& inputEdges,
|
||||
const gsl::span<const std::string> subGraphInputArgNames,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& initializerNameToInitializerMap,
|
||||
onnxruntime::Graph& graph,
|
||||
_Out_ std::vector<bool>& inputsUsed,
|
||||
_Inout_ std::vector<DML_BUFFER_BINDING>& initInputBindings,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initInputResources,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& nonOwnedGraphInputsFromInitializers,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
_Inout_opt_ std::vector<std::vector<std::byte>>* inputRawData)
|
||||
{
|
||||
|
||||
const uint32_t fusedNodeInputCount = gsl::narrow_cast<uint32_t>(subGraphInputArgNames.size());
|
||||
|
||||
// Determine the last input which uses an initializer, so initializers can be freed incrementally
|
||||
// while processing each input in order.
|
||||
std::map<const onnx::TensorProto*, uint32_t> 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<std::byte>());
|
||||
}
|
||||
|
||||
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<std::byte[]> 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<size_t>(tensorByteSize, 4);
|
||||
|
||||
if (inputRawData)
|
||||
{
|
||||
inputRawData->push_back(std::vector<std::byte>(tensorPtr, tensorPtr + tensorByteSize));
|
||||
}
|
||||
|
||||
if (!isInputsUploadedByDmlEP[i])
|
||||
{
|
||||
// Store the resource to use during execution
|
||||
ComPtr<ID3D12Resource> defaultBuffer = CreateResource(providerImpl, tensorPtr, tensorByteSize);
|
||||
nonOwnedGraphInputsFromInitializers[i] = defaultBuffer;
|
||||
initializeResourceRefs.push_back(std::move(defaultBuffer));
|
||||
}
|
||||
else
|
||||
{
|
||||
ComPtr<ID3D12Resource> 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::byte>());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::unordered_map<const onnx::TensorProto*, std::vector<uint32_t>>
|
||||
GetInitializerToPartitionMap(
|
||||
const onnxruntime::GraphViewer& graph,
|
||||
gsl::span<std::unique_ptr<GraphPartition>> partitions
|
||||
)
|
||||
{
|
||||
std::unordered_map<const onnx::TensorProto*, std::vector<uint32_t>> initializerPartitionMap;
|
||||
for (uint32_t partitionIndex = 0; partitionIndex < gsl::narrow_cast<uint32_t>(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<DML_OPERATOR_GRAPH_NODE_DESC>& dmlOperatorGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_NODE_DESC>& dmlGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlInputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlOutputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& 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<uint32_t>(dmlGraphNodes.size());
|
||||
dmlGraphDesc.Nodes = dmlGraphNodes.data();
|
||||
dmlGraphDesc.InputEdgeCount = gsl::narrow_cast<uint32_t>(dmlInputEdges.size());
|
||||
dmlGraphDesc.InputEdges = dmlInputEdges.data();
|
||||
dmlGraphDesc.OutputEdgeCount = gsl::narrow_cast<uint32_t>(dmlOutputEdges.size());
|
||||
dmlGraphDesc.OutputEdges = dmlOutputEdges.data();
|
||||
dmlGraphDesc.IntermediateEdgeCount = gsl::narrow_cast<uint32_t>(dmlIntermediateEdges.size());
|
||||
dmlGraphDesc.IntermediateEdges = dmlIntermediateEdges.data();
|
||||
}
|
||||
|
||||
void CreateIDmlCompiledOperatorAndRegisterKernel(
|
||||
onnxruntime::Graph& graph,
|
||||
const onnxruntime::IndexedSubGraph& indexedSubGraph,
|
||||
const onnxruntime::Node& fusedNode,
|
||||
const std::unordered_map<std::string, GraphNodeProperties>& partitionNodePropsMap,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& initializerNameToInitializerMap,
|
||||
const ExecutionProviderImpl* providerImpl,
|
||||
onnxruntime::KernelRegistry* registryForPartitionKernels)
|
||||
{
|
||||
// convert partitionONNXGraph into DML EP GraphDesc
|
||||
const uint32_t fusedNodeInputCount = gsl::narrow_cast<uint32_t>(indexedSubGraph.GetMetaDef()->inputs.size());
|
||||
const uint32_t fusedNodeOutputCount = gsl::narrow_cast<uint32_t>(indexedSubGraph.GetMetaDef()->outputs.size());
|
||||
|
||||
std::vector<uint8_t> 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<IDMLDevice> 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<DML_OPERATOR_GRAPH_NODE_DESC> dmlOperatorGraphNodes(graphDesc.nodes.size());
|
||||
std::vector<DML_GRAPH_NODE_DESC> dmlGraphNodes(graphDesc.nodes.size());
|
||||
std::vector<DML_GRAPH_EDGE_DESC> dmlInputEdges(graphDesc.inputEdges.size());
|
||||
std::vector<DML_GRAPH_EDGE_DESC> dmlOutputEdges(graphDesc.outputEdges.size());
|
||||
std::vector<DML_GRAPH_EDGE_DESC> 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<IDMLDevice1> device1;
|
||||
ORT_THROW_IF_FAILED(device.As(&device1));
|
||||
ComPtr<IDMLCompiledOperator> compiledExecutionPlanOperator;
|
||||
ORT_THROW_IF_FAILED(device1->CompileGraph(
|
||||
&dmlGraphDesc,
|
||||
executionFlags,
|
||||
IID_PPV_ARGS(&compiledExecutionPlanOperator)));
|
||||
|
||||
// Populate input bindings for operator initialization
|
||||
std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>> initializeResourceRefs; // For lifetime control
|
||||
std::vector<DML_BUFFER_BINDING> initInputBindings(fusedNodeInputCount);
|
||||
std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>> initInputResources; // For lifetime control
|
||||
std::vector<ComPtr<ID3D12Resource>> nonOwnedGraphInputsFromInitializers(fusedNodeInputCount);
|
||||
|
||||
std::vector<bool> 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<onnxruntime::OpKernel>& 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<const onnxruntime::Node*, GraphNodeProperties>& graphNodePropertyMap,
|
||||
onnxruntime::KernelRegistry* registryForPartitionKernels,
|
||||
const std::string& partitionKernelPrefix,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& 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<onnxruntime::IndexedSubGraph::MetaDef>();
|
||||
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<std::string, GraphNodeProperties> 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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 <typename T>
|
||||
static T AlignToPow2(T offset, T alignment)
|
||||
{
|
||||
static_assert(std::is_unsigned_v<T>);
|
||||
assert(alignment != 0);
|
||||
assert((alignment & (alignment - 1)) == 0);
|
||||
return (offset + alignment - 1) & ~(alignment - 1);
|
||||
}
|
||||
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource>
|
||||
CreateResource(
|
||||
const ExecutionProviderImpl* provider,
|
||||
const std::byte* tensorPtr,
|
||||
size_t tensorByteSize);
|
||||
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource>
|
||||
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<uint8_t>& isInputsUploadedByDmlEP,
|
||||
std::vector<DML_INPUT_GRAPH_EDGE_DESC>& inputEdges,
|
||||
const gsl::span<const std::string> subGraphInputArgNames,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& isInitializerTransferable,
|
||||
onnxruntime::Graph& graph,
|
||||
_Out_ std::vector<bool>& inputsUsed,
|
||||
_Inout_ std::vector<DML_BUFFER_BINDING>& initInputBindings,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initInputResources,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& nonOwnedGraphInputsFromInitializers,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
_Inout_opt_ std::vector<std::vector<std::byte>>* inputRawData);
|
||||
|
||||
std::unordered_map<const onnx::TensorProto*, std::vector<uint32_t>>
|
||||
GetInitializerToPartitionMap(
|
||||
const onnxruntime::GraphViewer& graph,
|
||||
gsl::span<std::unique_ptr<GraphPartition>> partitions
|
||||
);
|
||||
|
||||
void ConvertGraphDesc(
|
||||
const Dml::GraphDescBuilder::GraphDesc& graphDesc,
|
||||
_Out_ DML_GRAPH_DESC& dmlGraphDesc,
|
||||
const uint32_t inputCount,
|
||||
const uint32_t outputCount,
|
||||
_Inout_ std::vector<DML_OPERATOR_GRAPH_NODE_DESC>& dmlOperatorGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_NODE_DESC>& dmlGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlInputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlOutputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlIntermediateEdges);
|
||||
|
||||
void CreateIDmlCompiledOperatorAndRegisterKernel(
|
||||
onnxruntime::Graph& graph,
|
||||
const onnxruntime::IndexedSubGraph& indexedSubGraph,
|
||||
const onnxruntime::Node& fusedNode,
|
||||
const std::unordered_map<std::string, GraphNodeProperties>& partitionNodePropsMap,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& isInitializerTransferable,
|
||||
const ExecutionProviderImpl* providerImpl,
|
||||
onnxruntime::KernelRegistry* registryForPartitionKernels);
|
||||
|
||||
void FusePartitionAndRegisterKernel(
|
||||
GraphPartition* partition,
|
||||
uint32_t partitionIndex,
|
||||
onnxruntime::Graph& graph,
|
||||
std::unordered_map<const onnxruntime::Node*, GraphNodeProperties>& graphNodePropertyMap,
|
||||
onnxruntime::KernelRegistry* registryForPartitionKernels,
|
||||
const std::string& partitionKernelPrefix,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& isInitializerTransferable,
|
||||
const ExecutionProviderImpl* providerImpl);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<const ExecutionProvider*>(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<const onnxruntime::KernelRegistry*> 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<std::string> requiredInitializerMap;
|
||||
std::unordered_map<const onnxruntime::Node*, GraphNodeProperties> graphNodePropertyMap;
|
||||
onnxruntime::GraphViewer graphViewer(graph);
|
||||
std::vector<std::unique_ptr<GraphPartition>> 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<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>> 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();
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
};
|
||||
}
|
||||
|
|
@ -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<MLOperatorTensorDataType>(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<onnxruntime::NodeArg const*> 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<void(const onnxruntime::NodeArg& nodeArg, bool isInput)> 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<InternalRegistrationInfo> 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<std::unique_ptr<onnxruntime::ComputeCapability>>
|
||||
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<std::unique_ptr<onnxruntime::ComputeCapability>> result;
|
||||
|
||||
// Get the list of node indices in toplogical order, so nodes are visited before
|
||||
// downstream nodes consuming them.
|
||||
const std::vector<onnxruntime::NodeIndex>& toplogicalOrder = graph.GetNodesInTopologicalOrder();
|
||||
for (size_t nodeIndex : toplogicalOrder)
|
||||
{
|
||||
const onnxruntime::Node& node = *graph.GetNode(nodeIndex);
|
||||
if (IsNodeSupportedByDml(node, kernel_lookup, deviceDataTypeMask))
|
||||
{
|
||||
std::unique_ptr<onnxruntime::IndexedSubGraph> subGraph = std::make_unique<onnxruntime::IndexedSubGraph>();
|
||||
subGraph->nodes = {nodeIndex};
|
||||
result.push_back(std::make_unique<onnxruntime::ComputeCapability>(std::move(subGraph)));
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
bool IsGpuTensor(const onnxruntime::Tensor& tensor)
|
||||
|
|
|
|||
|
|
@ -158,11 +158,27 @@ namespace Dml
|
|||
std::shared_ptr<const Windows::AI::MachineLearning::Adapter::InternalRegistrationInfoMap>
|
||||
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<ID3D12Device> m_d3d12Device;
|
||||
ComPtr<IDMLDevice> m_dmlDevice;
|
||||
bool m_isMcdmDevice = false;
|
||||
|
|
@ -176,7 +192,6 @@ namespace Dml
|
|||
std::shared_ptr<onnxruntime::KernelRegistry> m_kernelRegistry;
|
||||
std::shared_ptr<const Windows::AI::MachineLearning::Adapter::InternalRegistrationInfoMap> m_internalRegInfoMap;
|
||||
mutable uint64_t m_partitionKernelPrefixVal = 0;
|
||||
|
||||
bool m_closed = false;
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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<std::string, GraphNodeProperties> &graphNodePropertyMap,
|
||||
std::unordered_map<std::string, onnx::TensorProto>& transferredInitializerMap,
|
||||
const gsl::span<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const gsl::span<const std::string> 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<IDMLCompiledOperator> compiledExecutionPlanOperator,
|
||||
Windows::AI::MachineLearning::Adapter::EdgeShapes& outputShapes,
|
||||
bool reuseCommandList,
|
||||
std::vector<ComPtr<ID3D12Resource>>& nonOwnedGraphInputsFromInitializers,
|
||||
std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
std::vector<DML_BUFFER_BINDING> initInputBindings,
|
||||
std::vector<uint8_t>& isInputsUploadedByDmlEP,
|
||||
std::vector<bool>& 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<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const gsl::span<const std::string> fusedNodeOutputArgOriginalNames,
|
||||
const std::unordered_map<std::string, GraphNodeProperties>& graphNodePropertyMap,
|
||||
std::unordered_map<std::string, onnx::TensorProto>& transferredInitializerMap
|
||||
std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
std::vector<DML_BUFFER_BINDING> initInputBindings,
|
||||
bool reuseCommandList
|
||||
)
|
||||
{
|
||||
ComPtr<IDMLDevice> device;
|
||||
ORT_THROW_IF_FAILED(m_provider->GetDmlDevice(device.GetAddressOf()));
|
||||
|
||||
ComPtr<IDMLDevice1> 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<Microsoft::WRL::ComPtr<ID3D12Resource>> initInputResources; // For lifetime control
|
||||
std::vector<DML_BUFFER_BINDING> initInputBindings(graphInputCount);
|
||||
m_nonOwnedGraphInputsFromInitializers.resize(graphInputCount);
|
||||
std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>> 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<DML_OPERATOR_GRAPH_NODE_DESC> dmlOperatorGraphNodes(graphDesc.nodes.size());
|
||||
std::vector<DML_GRAPH_NODE_DESC> dmlGraphNodes(graphDesc.nodes.size());
|
||||
|
||||
std::vector<DML_GRAPH_EDGE_DESC> dmlInputEdges(graphDesc.inputEdges.size());
|
||||
std::vector<DML_GRAPH_EDGE_DESC> dmlOutputEdges(graphDesc.outputEdges.size());
|
||||
std::vector<DML_GRAPH_EDGE_DESC> 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<ID3D12Resource>& 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<onnxruntime::Tensor>(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<size_t>(tensor->SizeInBytes(), 4);
|
||||
inputBindings[i].SizeInBytes = DmlGraphFusionHelper::AlignToPow2<size_t>(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<size_t>(tensor->SizeInBytes(), 4);
|
||||
outputBindings[i].SizeInBytes = DmlGraphFusionHelper::AlignToPow2<size_t>(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<IWinmlExecutionProvider> m_winmlProvider;
|
||||
ComPtr<Dml::IExecutionProvider> 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<ID3D12GraphicsCommandList> m_graphicsCommandList;
|
||||
|
|
@ -509,23 +421,32 @@ namespace Dml
|
|||
mutable ComPtr<ID3D12Fence> m_fence;
|
||||
mutable uint64_t m_completionValue = 0;
|
||||
|
||||
std::vector<uint8_t> m_inputsConstant;
|
||||
std::vector<uint8_t> m_isInputsUploadedByDmlEP;
|
||||
std::vector<ComPtr<ID3D12Resource>> m_nonOwnedGraphInputsFromInitializers;
|
||||
};
|
||||
|
||||
onnxruntime::OpKernel* CreateFusedGraphKernel(
|
||||
const onnxruntime::OpKernelInfo& info,
|
||||
const std::unordered_map<std::string, GraphNodeProperties> &graphNodePropertyMap,
|
||||
std::unordered_map<std::string, onnx::TensorProto>& transferredInitializerMap,
|
||||
const gsl::span<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const gsl::span<const std::string> fusedNodeOutputArgOriginalNames
|
||||
const onnxruntime::OpKernelInfo& info,
|
||||
ComPtr<IDMLCompiledOperator> compiledExecutionPlanOperator,
|
||||
Windows::AI::MachineLearning::Adapter::EdgeShapes& outputShapes,
|
||||
bool reuseCommandList,
|
||||
std::vector<ComPtr<ID3D12Resource>>& nonOwnedGraphInputsFromInitializers,
|
||||
std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
std::vector<DML_BUFFER_BINDING> initInputBindings,
|
||||
std::vector<uint8_t>& isInputsUploadedByDmlEP,
|
||||
std::vector<bool>& inputsUsed
|
||||
)
|
||||
{
|
||||
return new FusedGraphKernel(
|
||||
info,
|
||||
graphNodePropertyMap,
|
||||
transferredInitializerMap,
|
||||
fusedNodeInputArgOriginalNames,
|
||||
fusedNodeOutputArgOriginalNames);
|
||||
compiledExecutionPlanOperator,
|
||||
outputShapes,
|
||||
reuseCommandList,
|
||||
nonOwnedGraphInputsFromInitializers,
|
||||
initializeResourceRefs,
|
||||
initInputBindings,
|
||||
isInputsUploadedByDmlEP,
|
||||
inputsUsed
|
||||
);
|
||||
}
|
||||
} // namespace Dml
|
||||
|
|
|
|||
|
|
@ -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<std::string, GraphNodeProperties>& graphNodePropertyMap,
|
||||
std::unordered_map<std::string, onnx::TensorProto>& transferredInitializerMap,
|
||||
const gsl::span<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const gsl::span<const std::string> fusedNodeOutputArgOriginalNames
|
||||
const onnxruntime::OpKernelInfo& info,
|
||||
ComPtr<IDMLCompiledOperator> compiledExecutionPlanOperator,
|
||||
Windows::AI::MachineLearning::Adapter::EdgeShapes& outputShapes,
|
||||
bool reuseCommandList,
|
||||
std::vector<ComPtr<ID3D12Resource>>& nonOwnedGraphInputsFromInitializers,
|
||||
std::vector<Microsoft::WRL::ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
std::vector<DML_BUFFER_BINDING> initInputBindings,
|
||||
std::vector<uint8_t>& isInputsUploadedByDmlEP,
|
||||
std::vector<bool>& inputsUsed
|
||||
);
|
||||
} // namespace Dml
|
||||
|
|
|
|||
|
|
@ -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<std::string, onnx::TensorProto>& transferredInitializerMap,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& isInitializerTransferable,
|
||||
const onnxruntime::Graph& graph,
|
||||
const gsl::span<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const gsl::span<const std::string> fusedNodeOutputArgOriginalNames,
|
||||
const onnxruntime::IndexedSubGraph& indexedSubGraph,
|
||||
const std::unordered_map<std::string, GraphNodeProperties>& graphNodePropertyMap,
|
||||
IDMLDevice* device,
|
||||
const void* executionHandle)
|
||||
{
|
||||
|
||||
const gsl::span<const std::string> subGraphInputArgNames = indexedSubGraph.GetMetaDef()->inputs;
|
||||
const gsl::span<const std::string> 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<std::string, uint32_t> 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<DML_INTERMEDIATE_GRAPH_EDGE_DESC> graphIntermediateEdges;
|
||||
std::vector<DML_OUTPUT_GRAPH_EDGE_DESC> 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<onnxruntime::NodeIndex>& 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<OnnxTensorWrapper> tensorWrapper;
|
||||
|
||||
auto iter = transferredInitializerMap.find(argName);
|
||||
if (iter != transferredInitializerMap.end())
|
||||
auto iter = isInitializerTransferable.find(argName);
|
||||
if (iter != isInitializerTransferable.end())
|
||||
{
|
||||
tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(&iter->second);
|
||||
// Using const_cast here is simpler than making surrounding code const correct.
|
||||
tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(const_cast<ONNX_NAMESPACE::TensorProto*>(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());
|
||||
|
|
|
|||
|
|
@ -40,13 +40,11 @@ namespace Dml
|
|||
};
|
||||
|
||||
GraphDesc BuildGraphDesc(
|
||||
const onnxruntime::OpKernelInfo& kernelInfo,
|
||||
const uint8_t* isConstGpuGraphInput,
|
||||
const size_t isConstGpuGraphInputCount,
|
||||
std::unordered_map<std::string, onnx::TensorProto>& transferredInitializerMap,
|
||||
const std::unordered_map<std::string, std::pair<const ONNX_NAMESPACE::TensorProto*, bool>>& isInitializerTransferable,
|
||||
const onnxruntime::Graph& graph,
|
||||
const gsl::span<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const gsl::span<const std::string> fusedNodeOutputArgOriginalNames,
|
||||
const onnxruntime::IndexedSubGraph& indexedSubGraph,
|
||||
const std::unordered_map<std::string, GraphNodeProperties>& graphNodePropertyMap,
|
||||
IDMLDevice* device,
|
||||
const void* executionHandle);
|
||||
|
|
|
|||
|
|
@ -1,324 +0,0 @@
|
|||
#include "precomp.h"
|
||||
|
||||
#include "GraphKernelHelper.h"
|
||||
|
||||
namespace Dml
|
||||
{
|
||||
namespace GraphKernelHelper
|
||||
{
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource>
|
||||
CreateResource(
|
||||
Dml::IExecutionProvider* provider,
|
||||
const std::byte* tensorPtr,
|
||||
size_t tensorByteSize)
|
||||
{
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource> 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<uint64_t>((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<ID3D12Device> 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<ID3D12Resource>
|
||||
CreateCpuResource(
|
||||
Dml::IExecutionProvider* provider,
|
||||
const std::byte* tensorPtr,
|
||||
size_t tensorByteSize)
|
||||
{
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource> 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<uint64_t>((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<ID3D12Device> 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<IUnknown*>(const_cast<void*>(tensor->DataRaw()));
|
||||
Microsoft::WRL::ComPtr<IUnknown> 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<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const std::unordered_map<std::string, onnx::TensorProto>& 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<uint8_t>& inputsConstant,
|
||||
const onnxruntime::OpKernelInfo& kernelInfo,
|
||||
const Dml::GraphDescBuilder::GraphDesc& graphDesc,
|
||||
const gsl::span<const std::string> fusedNodeInputArgOriginalNames,
|
||||
_Out_ std::vector<bool>& inputsUsed,
|
||||
_Inout_ std::vector<DML_BUFFER_BINDING>& initInputBindings,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initInputResources,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& nonOwnedGraphInputsFromInitializers,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
_Inout_opt_ std::vector<std::vector<std::byte>>* inputRawData,
|
||||
_Inout_ std::unordered_map<std::string, onnx::TensorProto>& 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<const onnx::TensorProto*, uint32_t> 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<std::byte>());
|
||||
}
|
||||
|
||||
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<std::byte[]> 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<size_t>(tensorByteSize, 4);
|
||||
|
||||
if (inputRawData)
|
||||
{
|
||||
inputRawData->push_back(std::vector<std::byte>(tensorPtr, tensorPtr + tensorByteSize));
|
||||
}
|
||||
|
||||
if (!inputsConstant[i])
|
||||
{
|
||||
// Store the resource to use during execution
|
||||
ComPtr<ID3D12Resource> defaultBuffer = CreateResource(provider, tensorPtr, tensorByteSize);
|
||||
nonOwnedGraphInputsFromInitializers[i] = defaultBuffer;
|
||||
initializeResourceRefs.push_back(std::move(defaultBuffer));
|
||||
}
|
||||
else
|
||||
{
|
||||
ComPtr<ID3D12Resource> 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<const std::byte*>(inputTensor->DataRaw());
|
||||
|
||||
if (inputRawData)
|
||||
{
|
||||
inputRawData->push_back(
|
||||
std::vector<std::byte>(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<std::byte>());
|
||||
}
|
||||
}
|
||||
|
||||
// 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<DML_OPERATOR_GRAPH_NODE_DESC>& dmlOperatorGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_NODE_DESC>& dmlGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlInputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlOutputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& 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<uint32_t>(dmlGraphNodes.size());
|
||||
dmlGraphDesc.Nodes = dmlGraphNodes.data();
|
||||
dmlGraphDesc.InputEdgeCount = gsl::narrow_cast<uint32_t>(dmlInputEdges.size());
|
||||
dmlGraphDesc.InputEdges = dmlInputEdges.data();
|
||||
dmlGraphDesc.OutputEdgeCount = gsl::narrow_cast<uint32_t>(dmlOutputEdges.size());
|
||||
dmlGraphDesc.OutputEdges = dmlOutputEdges.data();
|
||||
dmlGraphDesc.IntermediateEdgeCount = gsl::narrow_cast<uint32_t>(dmlIntermediateEdges.size());
|
||||
dmlGraphDesc.IntermediateEdges = dmlIntermediateEdges.data();
|
||||
}
|
||||
} // namespace GraphKernelHelper
|
||||
} // namespace Dml
|
||||
|
|
@ -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 <typename T>
|
||||
static T AlignToPow2(T offset, T alignment)
|
||||
{
|
||||
static_assert(std::is_unsigned_v<T>);
|
||||
assert(alignment != 0);
|
||||
assert((alignment & (alignment - 1)) == 0);
|
||||
return (offset + alignment - 1) & ~(alignment - 1);
|
||||
}
|
||||
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource>
|
||||
CreateResource(
|
||||
Dml::IExecutionProvider* provider,
|
||||
const std::byte* tensorPtr,
|
||||
size_t tensorByteSize);
|
||||
|
||||
Microsoft::WRL::ComPtr<ID3D12Resource>
|
||||
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<const std::string> fusedNodeInputArgOriginalNames,
|
||||
const std::unordered_map<std::string, onnx::TensorProto>& transferredInitializerMap);
|
||||
|
||||
void ProcessInputData(
|
||||
Dml::IExecutionProvider* provider,
|
||||
IWinmlExecutionProvider* winmlProvider,
|
||||
const std::vector<uint8_t>& inputsConstant,
|
||||
const onnxruntime::OpKernelInfo& kernelInfo,
|
||||
const Dml::GraphDescBuilder::GraphDesc& graphDesc,
|
||||
const gsl::span<const std::string> fusedNodeInputArgOriginalNames,
|
||||
_Out_ std::vector<bool>& inputsUsed,
|
||||
_Inout_ std::vector<DML_BUFFER_BINDING>& initInputBindings,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initInputResources,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& nonOwnedGraphInputsFromInitializers,
|
||||
_Inout_ std::vector<ComPtr<ID3D12Resource>>& initializeResourceRefs,
|
||||
_Inout_opt_ std::vector<std::vector<std::byte>>* inputRawData,
|
||||
_Inout_ std::unordered_map<std::string, onnx::TensorProto>& transferredInitializerMap);
|
||||
|
||||
void ConvertGraphDesc(
|
||||
const Dml::GraphDescBuilder::GraphDesc& graphDesc,
|
||||
_Out_ DML_GRAPH_DESC& dmlGraphDesc,
|
||||
const onnxruntime::OpKernelInfo& kernelInfo,
|
||||
_Inout_ std::vector<DML_OPERATOR_GRAPH_NODE_DESC>& dmlOperatorGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_NODE_DESC>& dmlGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlInputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlOutputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlIntermediateEdges);
|
||||
|
||||
} // namespace GraphKernelHelper
|
||||
} // namespace Dml
|
||||
|
|
@ -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<MLOperatorTensorDataType>(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<onnxruntime::NodeArg const*> 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<void(const onnxruntime::NodeArg& nodeArg, bool isInput)> 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<InternalRegistrationInfo> 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<std::string, GraphPartition*>* nodeNameToPartitionMap,
|
||||
_Inout_ std::unordered_map<const onnxruntime::Node*, GraphNodeProperties>& dmlNodePropertyMap,
|
||||
_Inout_ std::unordered_set<std::string>& 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<uint32_t> 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<uint32_t> 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<onnxruntime::ComputeCapability> ComputationCapacityFromPartition(
|
||||
GraphPartition* partition,
|
||||
uint32_t partitionIndex,
|
||||
const onnxruntime::GraphViewer& graph,
|
||||
std::unordered_map<const onnxruntime::Node*, GraphNodeProperties>&& graphNodePropertyMap,
|
||||
onnxruntime::KernelRegistry* registryForPartitionKernels,
|
||||
const std::string& partitionKernelPrefix,
|
||||
std::shared_ptr<std::unordered_map<std::string, onnx::TensorProto>> transferredInitializerMap)
|
||||
{
|
||||
std::unique_ptr<onnxruntime::IndexedSubGraph> subGraph = std::make_unique<onnxruntime::IndexedSubGraph>();
|
||||
|
||||
if (partition->IsDmlGraphPartition())
|
||||
{
|
||||
assert(partition->IsDmlGraphPartition());
|
||||
|
||||
// Create a definition for the node. The name must be unique.
|
||||
auto def = std::make_unique<onnxruntime::IndexedSubGraph::MetaDef>();
|
||||
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<std::string, GraphNodeProperties> 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<std::string> fusedNodeInputArgOriginalNames = def->inputs;
|
||||
std::vector<std::string> fusedNodeOutputArgOriginalNames = def->outputs;
|
||||
auto fused_kernel_func = [partitionNodePropsMap, transferredInitializerMap, fusedNodeInputArgOriginalNames, fusedNodeOutputArgOriginalNames](onnxruntime::FuncManager& func_mgr, const onnxruntime::OpKernelInfo& info, std::unique_ptr<onnxruntime::OpKernel>& 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<onnxruntime::ComputeCapability>(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<const onnx::TensorProto*, std::vector<uint32_t>>
|
||||
GetInitializerToPartitionMap(
|
||||
const onnxruntime::GraphViewer& graph,
|
||||
gsl::span<std::unique_ptr<GraphPartition>> partitions
|
||||
)
|
||||
{
|
||||
std::unordered_map<const onnx::TensorProto*, std::vector<uint32_t>> initializerPartitionMap;
|
||||
|
||||
for (uint32_t partitionIndex = 0; partitionIndex < gsl::narrow_cast<uint32_t>(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<std::string, onnx::TensorProto> 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<std::unique_ptr<onnxruntime::ComputeCapability>>
|
||||
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<std::unique_ptr<onnxruntime::ComputeCapability>> result;
|
||||
|
||||
// Initializers needed by any graph partition
|
||||
std::unordered_set<std::string> requiredInitializerMap;
|
||||
|
||||
std::unordered_map<const onnxruntime::Node*, GraphNodeProperties> graphNodePropertyMap;
|
||||
std::vector<std::unique_ptr<GraphPartition>> 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<std::unordered_map<std::string, onnx::TensorProto>>();
|
||||
|
||||
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<onnx::TensorProto&>(*tensor);
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
// Transfer the initializer
|
||||
auto& graphTensor = const_cast<onnx::TensorProto&>(*tensor);
|
||||
|
||||
onnx::TensorProto partitionTensor;
|
||||
graphTensor.Swap(&partitionTensor);
|
||||
(*transferredInitializerMap)[input] = std::move(partitionTensor);
|
||||
|
||||
const_cast<onnxruntime::InitializedTensorSet&>(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
|
||||
|
|
|
|||
|
|
@ -49,22 +49,4 @@ namespace Dml
|
|||
std::unordered_map<const onnxruntime::Node*, GraphNodeProperties>& graphNodePropertyMap,
|
||||
std::unordered_set<std::string>& requiredInitializerMap,
|
||||
std::function<void(const onnxruntime::Node&)> onNodeUnsupportedInGraph = nullptr);
|
||||
|
||||
std::vector<std::unique_ptr<onnxruntime::ComputeCapability>>
|
||||
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
|
||||
|
|
|
|||
|
|
@ -17,15 +17,6 @@
|
|||
|
||||
namespace Dml
|
||||
{
|
||||
GraphTransformer::GraphTransformer(
|
||||
const std::string& name,
|
||||
const onnxruntime::IExecutionProvider* provider
|
||||
)
|
||||
: onnxruntime::GraphTransformer(name),
|
||||
m_providerImpl(static_cast<const ExecutionProvider* >(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<NodeToAdd> nodesToAdd;
|
||||
|
||||
onnxruntime::ProviderType provider_type = onnxruntime::kDmlExecutionProvider;
|
||||
const gsl::not_null<const onnxruntime::KernelRegistry*> 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<onnxruntime::NodeArg*> inputs;
|
||||
std::vector<onnxruntime::NodeArg*> 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<NodeToAdd> nodesToAdd;
|
||||
std::vector<onnxruntime::NodeIndex> 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
|
||||
|
|
|
|||
|
|
@ -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<onnxruntime::KernelRegistry> m_registry;
|
||||
const ExecutionProviderImpl* m_providerImpl = nullptr;
|
||||
};
|
||||
|
||||
} // namespace Dml
|
||||
|
|
|
|||
|
|
@ -55,7 +55,7 @@ namespace Dml
|
|||
std::vector<DML_INPUT_GRAPH_EDGE_DESC> 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<uint32_t>(inputEdges.size());
|
||||
operatorGraphDesc.inputEdges = inputEdges.data();
|
||||
|
||||
|
||||
|
||||
std::vector<DML_OUTPUT_GRAPH_EDGE_DESC> 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<IMLOperatorKernelCreationContextPrivate> 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<DML_GRAPH_NODE_DESC> dmlGraphNodes(operatorGraphDesc.nodeCount);
|
||||
std::vector<ComPtr<IDMLOperator>> dmlOperators(operatorGraphDesc.nodeCount);
|
||||
std::vector<DML_OPERATOR_GRAPH_NODE_DESC> dmlOperatorGraphNodes(operatorGraphDesc.nodeCount);
|
||||
std::vector<DML_GRAPH_EDGE_DESC> dmlInputEdges(operatorGraphDesc.inputEdgeCount);
|
||||
std::vector<DML_GRAPH_EDGE_DESC> dmlOutputEdges(operatorGraphDesc.outputEdgeCount);
|
||||
std::vector<DML_GRAPH_EDGE_DESC> 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<IDMLDevice1> 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<size_t>(persistentResourceSize),
|
||||
AllocatorRoundingMode::Enabled,
|
||||
m_persistentResource.GetAddressOf(),
|
||||
m_persistentResourcePoolingUnk.GetAddressOf()));
|
||||
|
||||
m_persistentResourceBinding = DML_BUFFER_BINDING{ m_persistentResource.Get(), 0, persistentResourceSize };
|
||||
}
|
||||
|
||||
std::vector<DML_BUFFER_BINDING> 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<ComPtr<IDMLOperator>>& dmlOperators,
|
||||
_Inout_ std::vector<DML_OPERATOR_GRAPH_NODE_DESC>& dmlOperatorGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_NODE_DESC>& dmlGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlInputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlOutputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlIntermediateEdges)
|
||||
{
|
||||
graphDesc.InputCount = gsl::narrow_cast<uint32_t>(m_kernelInputIndices.size());
|
||||
graphDesc.OutputCount = gsl::narrow_cast<uint32_t>(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
|
||||
|
|
|
|||
|
|
@ -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<std::optional<uint32_t>> m_kernelInputIndices;
|
||||
std::vector<std::optional<uint32_t>> m_kernelOutputIndices;
|
||||
|
||||
void ConvertToDmlGraphDesc(const MLOperatorGraphDesc& operatorGraphDesc,
|
||||
_Out_ DML_GRAPH_DESC& graphDesc,
|
||||
_Inout_ std::vector<ComPtr<IDMLOperator>>& dmlOperators,
|
||||
_Inout_ std::vector<DML_OPERATOR_GRAPH_NODE_DESC>& dmlOperatorGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_NODE_DESC>& dmlGraphNodes,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlInputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlOutputEdges,
|
||||
_Inout_ std::vector<DML_GRAPH_EDGE_DESC>& dmlIntermediateEdges);
|
||||
|
||||
};
|
||||
|
||||
} // namespace Dml
|
||||
|
|
|
|||
|
|
@ -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<uint32_t> outputShape = kernelCreationContext.GetTensorShapeDescription().GetOutputTensorShape(0);
|
||||
const uint32_t outputShapeDimCount = gsl::narrow_cast<uint32_t>(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<uint32_t>(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<uint32_t> 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<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
|
||||
std::vector<DML_TENSOR_DESC> 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<const DML_OPERATOR_DESC*> opDescs{&opDesc1, &opDesc2, &opDesc3};
|
||||
operatorGraphDesc.nodesAsOpDesc = opDescs.data();
|
||||
|
||||
// set input edges
|
||||
std::pair<uint32_t, uint32_t> nodeToNodeInputIndex[5] {{0, 0}, {0, 1}, {0, 2}, {2, 1}, {2, 2}};
|
||||
std::vector<DML_INPUT_GRAPH_EDGE_DESC> 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<uint32_t>(inputEdges.size());
|
||||
operatorGraphDesc.inputEdges = inputEdges.data();
|
||||
|
||||
// set intermediate edges
|
||||
std::vector<DML_INTERMEDIATE_GRAPH_EDGE_DESC> 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<uint32_t>(intermediateEdges.size());
|
||||
operatorGraphDesc.intermediateEdges = intermediateEdges.data();
|
||||
|
||||
// set the output edges
|
||||
std::vector<DML_OUTPUT_GRAPH_EDGE_DESC> 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<uint32_t>(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
|
||||
|
|
@ -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<const char*, 1> typeNameListDefault = {"T"};
|
||||
constexpr static std::array<const char*, 2> typeNameListTwo = { "T1", "T2" };
|
||||
|
|
@ -330,6 +332,7 @@ constexpr static std::array<SupportedTensorDataTypes, 1> supportedTypeListIntege
|
|||
constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListRoiAlign = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Int32|SupportedTensorDataTypes::Int64 };
|
||||
constexpr static std::array<SupportedTensorDataTypes, 1> supportedTypeListArgMinMax = {SupportedTensorDataTypes::Float16to32|SupportedTensorDataTypes::Ints8to64};
|
||||
constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListLayerNormalization = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Float32};
|
||||
constexpr static std::array<SupportedTensorDataTypes, 1> supportedTypeListQLinearSigmoid = {SupportedTensorDataTypes::UInt8 | SupportedTensorDataTypes::Int8};
|
||||
|
||||
constexpr static std::array<SupportedTensorDataTypes, 3> 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)},
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<onnxruntime::GraphTransformer> dmlGraphFusionTransformer = std::make_unique<Dml::DmlGraphFusionTransformer>("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<onnxruntime::GraphTransformer> dmlOperatorFusionTransformer = std::make_unique<Dml::GraphTransformer>("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_,
|
||||
|
|
|
|||
Loading…
Reference in a new issue