mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
MKL-DNN Subgraphs (#1116)
* subgraph with memcpy fix * Linux compile errors fix * Linux compile errors fix * subgraph with memcpy fix * Linux compile errors fix * Linux compile errors fix * memcpy (PR1020) fix implemented * check graph viewer GetNode for nullptr at other plances * documents * Review changes (UseSubgraph simplified) * static_cast<int> removed * static_cast<int> removed 2 * fall back to CPU implementation in GetCapability() * check shape for null. fall back to CPU implementation in GetCapability() * backend data errors fixed * PR review changes * disable Opset10 tests * removed tests from main.cc of test runner. added a check at GetCapability() * backend data and Model-Zoo related fixes * subgraph with memcpy fix * Linux compile errors fix * Linux compile errors fix * subgraph with memcpy fix * Linux compile errors fix * memcpy (PR1020) fix implemented * documents * Review changes (UseSubgraph simplified) * static_cast<int> removed * fall back to CPU implementation in GetCapability() * check shape for null. fall back to CPU implementation in GetCapability() * backend data errors fixed * PR review changes * disable Opset10 tests * removed tests from main.cc of test runner. added a check at GetCapability() * backend data and Model-Zoo related fixes * patch to run tests and models separatly
This commit is contained in:
parent
3c3186c761
commit
24d6b0f5c4
18 changed files with 3239 additions and 2 deletions
35
docs/execution_providers/MKL-DNN-ExecutionProvider.md
Normal file
35
docs/execution_providers/MKL-DNN-ExecutionProvider.md
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
# MKL-DNN Execution Provider
|
||||
|
||||
Intel® Math Kernel Library for Deep Neural Networks (Intel® MKL-DNN) is an open-source performance library for deep-learning applications. The library accelerates deep-learning applications and frameworks on Intel® architecture and Intel® Processor Graphics Architecture. Intel MKL-DNN contains vectorized and threaded building blocks that you can use to implement deep neural networks (DNN) with C and C++ interfaces. For more visit MKL-DNN documentation at (https://intel.github.io/mkl-dnn/)
|
||||
|
||||
Intel and Microsoft have developed MKL-DNN Execution Provider (EP) for ONNX Runtime to accelerate performance of ONNX Runtime using Intel® Math Kernel Library for Deep Neural Networks (Intel® MKL-DNN) optimized primitives
|
||||
|
||||
## MKL-DNN/MKLML
|
||||
To build ONNX Runtime with MKL-DNN support, build it with `./build.sh --use_mkldnn`
|
||||
|
||||
To build ONNX Runtime using MKL-DNN built with dependency on MKL small libraries, build it with `./build.sh --use_mkldnn --use_mklml`
|
||||
|
||||
## Supported OS
|
||||
* Ubuntu 16.04
|
||||
* Windows 10
|
||||
* Mac OS X
|
||||
|
||||
## Supported backend
|
||||
* CPU
|
||||
* More to be added soon!
|
||||
|
||||
## Using the nGraph execution provider
|
||||
### C/C++
|
||||
The MKLDNNExecutionProvider execution provider needs to be registered with ONNX Runtime to enable in the inference session.
|
||||
```
|
||||
InferenceSession session_object{so};
|
||||
session_object.RegisterExecutionProvider(std::make_unique<::onnxruntime:: MKLDNNExecutionProvider >());
|
||||
status = session_object.Load(model_file_name);
|
||||
```
|
||||
The C API details are [here](https://github.com/Microsoft/onnxruntime/blob/master/docs/C_API.md#c-api).
|
||||
|
||||
## Python
|
||||
When using the python wheel from the ONNX Runtime built with MKL-DNN execution provider, it will be automatically prioritized over the CPU execution provider. Python APIs details are [here](https://github.com/Microsoft/onnxruntime/blob/master/docs/python/api_summary.rst#api-summary).
|
||||
|
||||
## Using onnxruntime_perf_test and onnx_test_runner
|
||||
You can test the performance of your ONNX Model with the MKL-DNN execution provider. Use the flag -e mkldnn in [onnxruntime_perf_test](https://github.com/Microsoft/onnxruntime/tree/master/onnxruntime/test/perftest#onnxruntime-performance-test) and [onnx_test_runner](https://github.com/Microsoft/onnxruntime/tree/master/onnxruntime/test/onnx/README.txt)..
|
||||
65
docs/execution_providers/MKL-DNN-Subgraphs.md
Normal file
65
docs/execution_providers/MKL-DNN-Subgraphs.md
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
# Subgraph Optimization
|
||||
|
||||
MKL-DNN uses blocked layout (example: nhwc with channels blocked by 16 – nChw16c) to take advantage of vector operations using AVX512. To get best performance, we avoid reorders (example. Nchw16c to nchw) and propagate blocked layout to next primitive.
|
||||
|
||||
Subgraph optimization achieves this in the following steps.
|
||||
1. Parses ONNX Runtime graph and creates an Internal Representation of subgraph..
|
||||
2. Subgraph Operator (MklDnnFunKernel) iterates through MKL-DNN nodes and creates a vector MKL-DNN Kernels
|
||||
3. Compute Function of MklDnnFunKernel iterates and binds data to MKL-DNN primitives in the vector and submits vector for execution.
|
||||
|
||||
|
||||
## Subgraph (IR) Internal Representation
|
||||
MklDnnExecutionProvicer::GetCapability() parses ONNX model graph and creates IR (Internal Representation) of subgraphs of MKL-DNN operators.
|
||||
Each subgraph contains a vector MklDnnNodes, inputs, outputs and attributes for all its MklDnnNodes. There can be attributes of same name. So, we prefix attribute names with Node name and its index.
|
||||
Unique id for subgraph is set as an attribute.
|
||||
|
||||
MklDnnNode has an index to its inputs and outputs and pointer to its parent nodes. MklDnnNode directly reads blocked memory from its parent to avoid data reordering.
|
||||
|
||||
<p align="left"><img src="images/mkl-dnn_node.png" /></p>
|
||||
|
||||
|
||||
## Subgraph Classes
|
||||
Primitive like MklDnnConv, MklDnnPool, etc are derived from MklDnnKernel base class.
|
||||
|
||||
The following UML diagram captures Subgraph classes.
|
||||
|
||||
<p align="left"><img src="images/mkl-dnn_subgraph.png" /></p>
|
||||
|
||||
|
||||
## Subgraph Execution
|
||||
|
||||
MklDnnExecutionProvicer::Compute() function creates MklDnnFuncKernel and call it’s Compute Function.
|
||||
|
||||
|
||||
MklDnnFuncKernel::Compute function creates SubgraphPrimitve pool and add the object to a map.
|
||||
|
||||
SubgraphPrimitve constructor calls the following member functions
|
||||
```
|
||||
SubgraphPrimitve::CreatePrimitives()
|
||||
for (auto& mklnode : mklnodes) {
|
||||
if (mklnode.name == "Conv") {
|
||||
kernel.reset(new MklDnnConv());
|
||||
kernels.push_back(kernel);
|
||||
} else if (mklnode.name == "BatchNormalization-Relu") {
|
||||
kernel.reset(new MklDnnBatchNorm());
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mklnode.name == "MaxPool") {
|
||||
kernel.reset(new MklDnnPool());
|
||||
context_.kernels.push_back(kernel);
|
||||
}
|
||||
.
|
||||
.
|
||||
.
|
||||
```
|
||||
In CreatePrimitives method, we iterate MklDnnNodes and creates MklDnnKernel objects and add MKL-DNN primitive to a vector. It also reads attributes. This is done only once, at first iteration.
|
||||
|
||||
```
|
||||
SubgraphPrimitve::Compute()
|
||||
for (auto& kernel : kernels) {
|
||||
kernel->Bind(input_tensors, output_tensors);
|
||||
}
|
||||
stream->submit(net);
|
||||
```
|
||||
|
||||
In SubgraphPrimitve::Compute() method, we iterate thru MklDnn Kernels and bind input data. Then we submit the vector of Primitives to MKL-DNN stream.
|
||||
|
||||
BIN
docs/execution_providers/images/mkl-dnn_node.png
Normal file
BIN
docs/execution_providers/images/mkl-dnn_node.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 50 KiB |
BIN
docs/execution_providers/images/mkl-dnn_subgraph.png
Normal file
BIN
docs/execution_providers/images/mkl-dnn_subgraph.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 42 KiB |
|
|
@ -1,11 +1,17 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#ifdef _MSC_VER
|
||||
#pragma warning(disable : 4996)
|
||||
#endif
|
||||
|
||||
#include "mkldnn_execution_provider.h"
|
||||
#include "core/framework/allocator.h"
|
||||
#include "core/framework/memcpy.h"
|
||||
#include "core/framework/kernel_registry.h"
|
||||
#include "mkldnn_fwd.h"
|
||||
#include "core/framework/compute_capability.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_func_kernel.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
|
|
@ -121,4 +127,337 @@ std::shared_ptr<KernelRegistry> MKLDNNExecutionProvider::GetKernelRegistry() con
|
|||
static std::shared_ptr<KernelRegistry> kernel_registry = onnxruntime::mkl_dnn::GetMklDnnKernelRegistry();
|
||||
return kernel_registry;
|
||||
}
|
||||
|
||||
bool MKLDNNExecutionProvider::UseSubgraph(const onnxruntime::GraphViewer& graph_viewer,
|
||||
const std::vector<const KernelRegistry*>& kernel_registries,
|
||||
std::vector<std::unique_ptr<ComputeCapability>>& result) const {
|
||||
// switch between mkldnn-vanilla and mkldnn-subgraph implementation using
|
||||
// MKLDNN_SUBGRAPH environment variable
|
||||
bool use_subgraph = true;
|
||||
|
||||
bool FP16_graph = false;
|
||||
if (graph_viewer.MaxNodeIndex() > 0) {
|
||||
int index = 0;
|
||||
auto node = graph_viewer.GetNode(index);
|
||||
while (node == NULL) {
|
||||
index++;
|
||||
node = graph_viewer.GetNode(index);
|
||||
}
|
||||
if (node->InputDefs()[0]->Type() != nullptr)
|
||||
FP16_graph = node->InputDefs()[0]->Type()->find("16") != std::string::npos;
|
||||
}
|
||||
|
||||
if (FP16_graph) {
|
||||
// FP16 not supported yet.
|
||||
use_subgraph = false;
|
||||
result = IExecutionProvider::GetCapability(graph_viewer, kernel_registries);
|
||||
} else {
|
||||
const char* env = getenv("ORT_MKLDNN_SUBGRAPH");
|
||||
if (env != nullptr) {
|
||||
if (atoi(env) == 0) {
|
||||
use_subgraph = false;
|
||||
result = IExecutionProvider::GetCapability(graph_viewer, kernel_registries);
|
||||
}
|
||||
}
|
||||
}
|
||||
return use_subgraph;
|
||||
}
|
||||
|
||||
void MKLDNNExecutionProvider::CreateOrUpdateMklDnnNode(const Node* node,
|
||||
std::shared_ptr<mkl_dnn::Subgraph>& subgraph_ptr,
|
||||
mkl_dnn::Subgraph::SubgraphVariables& sub_var,
|
||||
bool fused,
|
||||
std::map<std::string, size_t>& output_to_source_node_map,
|
||||
NodeAttributes& subgraph_attributes) const {
|
||||
const auto& node_inputs = node->InputDefs();
|
||||
sub_var.outputs.push_back(node->OutputDefs()[0]->Name());
|
||||
|
||||
if (!fused) {
|
||||
mkl_dnn::MklDnnNode mkldnn_node;
|
||||
mkldnn_node.name = node->OpType();
|
||||
mkldnn_node.num_inputs = static_cast<int>(node->InputDefs().size());
|
||||
mkldnn_node.input_start_index = static_cast<int>(sub_var.inputs.size()) - 1;
|
||||
mkldnn_node.node_index = static_cast<int>(subgraph_ptr->mkldnn_nodes.size()) + 1;
|
||||
const auto& node_outputs = node->OutputDefs();
|
||||
mkldnn_node.output_name = node_outputs[0]->Name();
|
||||
if (node->OpType() == "Conv") {
|
||||
mkldnn_node.weight_name = node->InputDefs()[1]->Name();
|
||||
}
|
||||
for (size_t i = 0; i < node_inputs.size(); i++) {
|
||||
auto iter = output_to_source_node_map.find(node_inputs[i]->Name());
|
||||
if (iter != output_to_source_node_map.end())
|
||||
mkldnn_node.parent_nodes.push_back(iter->second);
|
||||
}
|
||||
subgraph_ptr->mkldnn_nodes.push_back(mkldnn_node);
|
||||
output_to_source_node_map.insert(std::make_pair(node_outputs[0]->Name(), subgraph_ptr->mkldnn_nodes.size() - 1));
|
||||
} else {
|
||||
const auto& node_outputs = node->OutputDefs();
|
||||
output_to_source_node_map.erase(subgraph_ptr->mkldnn_nodes.back().output_name);
|
||||
subgraph_ptr->mkldnn_nodes.back().output_name = node_outputs[0]->Name();
|
||||
output_to_source_node_map.insert(std::make_pair(node_outputs[0]->Name(), subgraph_ptr->mkldnn_nodes.size() - 1));
|
||||
}
|
||||
|
||||
// Add inputs which are not in the outputs vector.
|
||||
for (size_t i = 0; i < node_inputs.size(); i++) {
|
||||
auto itr = std::find(sub_var.outputs.begin(), sub_var.outputs.end(), node_inputs[i]->Name());
|
||||
if (itr == sub_var.outputs.end()) {
|
||||
sub_var.inputs.push_back(node_inputs[i]->Name());
|
||||
} else {
|
||||
// Vector of node outputs, which is input to other node
|
||||
// if node output is not input to any other node, then it's the end node
|
||||
// which we will find later
|
||||
sub_var.outputs_as_input_other_node.push_back(node_inputs[i]->Name());
|
||||
}
|
||||
}
|
||||
|
||||
NodeAttributes attributes = node->GetAttributes();
|
||||
if (attributes.size() > 0) {
|
||||
size_t index = subgraph_ptr->mkldnn_nodes.size();
|
||||
|
||||
for (auto att_it = attributes.begin(); att_it != attributes.end(); ++att_it) {
|
||||
std::string key = node->OpType() + "-" + std::to_string(index) + "-" + att_it->first;
|
||||
std::pair<std::string, ONNX_NAMESPACE::AttributeProto> att(key, att_it->second);
|
||||
subgraph_attributes[key] = att_it->second;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::unique_ptr<ComputeCapability>> MKLDNNExecutionProvider::GetCapability(
|
||||
const onnxruntime::GraphViewer& graph_viewer,
|
||||
const std::vector<const KernelRegistry*>& kernel_registries) const {
|
||||
ORT_UNUSED_PARAMETER(kernel_registries);
|
||||
std::vector<std::unique_ptr<ComputeCapability>> result;
|
||||
|
||||
// temporary switch to toggle between mkldnn-vanilla and mkldnn-subgraph implementation using
|
||||
// ORT_MKLDNN_SUBGRAPH environment variable
|
||||
if (UseSubgraph(graph_viewer, kernel_registries, result) == false) {
|
||||
return result;
|
||||
}
|
||||
|
||||
LOGS_DEFAULT(INFO) << "Using MKL-DNN Subgraph";
|
||||
// use sub-graph implementation
|
||||
mkl_dnn::Subgraph::SubgraphVariables sub_var;
|
||||
std::shared_ptr<mkl_dnn::Subgraph> subgraph_ptr;
|
||||
|
||||
// We need graph name make PrimitivePool keys unique.
|
||||
// There are several identical graphs in Model zoo and only differ in
|
||||
// few attribute values. GetGraphName return graph-name + first-node-output name
|
||||
std::string graph_name = GetGraphName(graph_viewer);
|
||||
subgraph_ptr.reset(new mkl_dnn::Subgraph(graph_name));
|
||||
|
||||
// output name to node index map. Using it to find sub-graph end nodes
|
||||
// if output of a node is not an input to any node in a sub-graph is end node
|
||||
std::map<std::string, size_t> output_to_source_node_map;
|
||||
NodeAttributes subgraph_attributes;
|
||||
int node_index = 0;
|
||||
|
||||
while (node_index < graph_viewer.MaxNodeIndex()) {
|
||||
auto node = graph_viewer.GetNode(node_index);
|
||||
if (node == nullptr) {
|
||||
node_index++;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (IsDimensionSupported(node) == false) {
|
||||
node_index++;
|
||||
continue;
|
||||
}
|
||||
|
||||
auto op_it = mkldnn_ops_.find(node->OpType());
|
||||
if (op_it != mkldnn_ops_.end()) {
|
||||
sub_var.subgraph_node_indexes.push_back(node->Index());
|
||||
|
||||
// can we fuse (at mkldnn level) nodes?
|
||||
bool fused = false;
|
||||
if (sub_var.subgraph_node_indexes.size() > 1 && node->OpType() == "Relu") {
|
||||
if (subgraph_ptr->mkldnn_nodes.back().name == "BatchNormalization" || subgraph_ptr->mkldnn_nodes.back().name == "Conv") {
|
||||
subgraph_ptr->mkldnn_nodes.back().name += "-Relu";
|
||||
fused = true;
|
||||
}
|
||||
}
|
||||
|
||||
// Create MklDnn node:
|
||||
// Update inputs, outputs and parent nodes
|
||||
// Collect attributes and modify the key to make it unique
|
||||
CreateOrUpdateMklDnnNode(node, subgraph_ptr, sub_var, fused, output_to_source_node_map, subgraph_attributes);
|
||||
|
||||
auto temp_index = node_index + 1;
|
||||
if (temp_index < graph_viewer.MaxNodeIndex()) {
|
||||
if (!sub_var.subgraph_node_indexes.empty()) {
|
||||
// if next node is mkldnn node and if it's input is not output of current node
|
||||
// if next node input is output of any of the nodes in sub-graph continue
|
||||
// else
|
||||
// break and create sub-graph
|
||||
auto next_node = graph_viewer.GetNode(temp_index);
|
||||
while (next_node == nullptr) {
|
||||
temp_index++;
|
||||
next_node = graph_viewer.GetNode(temp_index);
|
||||
}
|
||||
auto sub_it = mkldnn_ops_.find(next_node->OpType());
|
||||
if (sub_it != mkldnn_ops_.end()) {
|
||||
const auto& next_node_inputs = next_node->InputDefs();
|
||||
bool input_from_subgraph = true;
|
||||
size_t inputs_count = 1;
|
||||
if (next_node->OpType() == "Sum")
|
||||
inputs_count = next_node_inputs.size();
|
||||
for (size_t i = 0; i < inputs_count; i++) {
|
||||
auto in = next_node_inputs[i];
|
||||
auto itr = std::find(sub_var.outputs.begin(), sub_var.outputs.end(), in->Name());
|
||||
if (itr == sub_var.outputs.end()) {
|
||||
input_from_subgraph = false;
|
||||
}
|
||||
}
|
||||
if (input_from_subgraph == false) {
|
||||
CreateMetaDef(graph_viewer, subgraph_attributes, subgraph_ptr, sub_var, result);
|
||||
subgraph_attributes.clear();
|
||||
subgraph_ptr.reset(new mkl_dnn::Subgraph(graph_name));
|
||||
output_to_source_node_map.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!sub_var.subgraph_node_indexes.empty()) {
|
||||
if (node->GetOutputEdgesCount() > 1) {
|
||||
// If current node has branches
|
||||
// iterate and see if all nodes are mkldnn ops OR
|
||||
// it ends in node with same number of input edges (mkldnn node or cpu node)
|
||||
// create sub-graph
|
||||
bool create_subgraph = false;
|
||||
bool break_loop = false;
|
||||
while (!break_loop) {
|
||||
if (temp_index > graph_viewer.MaxNodeIndex())
|
||||
break_loop = true;
|
||||
|
||||
auto next_node = graph_viewer.GetNode(temp_index);
|
||||
while (next_node == nullptr) {
|
||||
temp_index++;
|
||||
next_node = graph_viewer.GetNode(temp_index);
|
||||
}
|
||||
if (next_node->GetInputEdgesCount() == node->GetOutputEdgesCount()) {
|
||||
// if all nodes in the branch loop are mkldnn nodes
|
||||
// then continue with adding nodes to sub-graph
|
||||
break_loop = true;
|
||||
}
|
||||
// inner nodes. if inner nodes are not mkldnn nodes
|
||||
// create subgraph (inception v2)
|
||||
auto sub_it = mkldnn_ops_.find(next_node->OpType());
|
||||
if (sub_it == mkldnn_ops_.end()) {
|
||||
// break and create a sub-graph
|
||||
break_loop = true;
|
||||
create_subgraph = true;
|
||||
}
|
||||
temp_index++;
|
||||
}
|
||||
if (create_subgraph) {
|
||||
CreateMetaDef(graph_viewer, subgraph_attributes, subgraph_ptr, sub_var, result);
|
||||
subgraph_ptr.reset(new mkl_dnn::Subgraph(graph_name));
|
||||
subgraph_attributes.clear();
|
||||
output_to_source_node_map.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (!sub_var.subgraph_node_indexes.empty()) {
|
||||
CreateMetaDef(graph_viewer, subgraph_attributes, subgraph_ptr, sub_var, result);
|
||||
subgraph_ptr.reset(new mkl_dnn::Subgraph(graph_name));
|
||||
subgraph_attributes.clear();
|
||||
output_to_source_node_map.clear();
|
||||
}
|
||||
}
|
||||
node_index++;
|
||||
} // graph_viewer node iterator ends
|
||||
if (!sub_var.subgraph_node_indexes.empty()) {
|
||||
CreateMetaDef(graph_viewer, subgraph_attributes, subgraph_ptr, sub_var, result);
|
||||
subgraph_ptr.reset(new mkl_dnn::Subgraph(graph_name));
|
||||
subgraph_attributes.clear();
|
||||
output_to_source_node_map.clear();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
void MKLDNNExecutionProvider::CreateMetaDef(const onnxruntime::GraphViewer& graph_viewer,
|
||||
const NodeAttributes& subgraph_attributes,
|
||||
std::shared_ptr<mkl_dnn::Subgraph>& subgraph_ptr,
|
||||
mkl_dnn::Subgraph::SubgraphVariables& sub_var,
|
||||
std::vector<std::unique_ptr<ComputeCapability>>& result) const {
|
||||
std::string graph_fused_nodes;
|
||||
std::string node_list;
|
||||
std::string subgraph_id = std::to_string(sub_var.subgraph_index);
|
||||
sub_var.subgraph_index++;
|
||||
|
||||
// This is a list of initializers that subgraph considers as constants.
|
||||
// Example weights, reshape shape etc.
|
||||
std::unordered_set<std::string> input_initializers;
|
||||
|
||||
// Create ng_required_initializers attribute of NGraphCustomOp
|
||||
ONNX_NAMESPACE::AttributeProto initializers;
|
||||
initializers.set_name("initializers");
|
||||
initializers.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_TENSORS);
|
||||
|
||||
for (const auto& init : sub_var.inputs) {
|
||||
if (graph_viewer.GetAllInitializedTensors().count(init)) {
|
||||
auto tensor = initializers.add_tensors();
|
||||
*tensor = *(graph_viewer.GetAllInitializedTensors().at(init));
|
||||
}
|
||||
}
|
||||
|
||||
auto meta_def = std::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>();
|
||||
meta_def->attributes["initializers"] = initializers;
|
||||
meta_def->name = "MkldnnCustomOp" + std::to_string(sub_var.subgraph_index);
|
||||
meta_def->domain = kMSDomain;
|
||||
meta_def->since_version = 1;
|
||||
meta_def->status = ONNX_NAMESPACE::EXPERIMENTAL;
|
||||
meta_def->inputs = sub_var.inputs;
|
||||
meta_def->attributes.insert(subgraph_attributes.begin(), subgraph_attributes.end());
|
||||
|
||||
// Find the end nodes
|
||||
for (auto& mklnode : subgraph_ptr->mkldnn_nodes) {
|
||||
auto itr = std::find(sub_var.outputs_as_input_other_node.begin(),
|
||||
sub_var.outputs_as_input_other_node.end(), mklnode.output_name);
|
||||
if (itr == sub_var.outputs_as_input_other_node.end()) {
|
||||
meta_def->outputs.push_back(mklnode.output_name);
|
||||
mklnode.output_index = static_cast<int>(meta_def->outputs.size()) - 1;
|
||||
}
|
||||
}
|
||||
|
||||
ONNX_NAMESPACE::AttributeProto ap;
|
||||
ap.set_s(subgraph_id);
|
||||
ap.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_STRING);
|
||||
meta_def->attributes["subgraph_id"] = ap;
|
||||
std::unique_ptr<IndexedSubGraph> sub_graph = std::make_unique<IndexedSubGraph>();
|
||||
sub_graph->nodes = sub_var.subgraph_node_indexes;
|
||||
sub_graph->SetMetaDef(meta_def);
|
||||
result.push_back(std::make_unique<ComputeCapability>(std::move(sub_graph)));
|
||||
mkl_subgraphs_.insert(std::make_pair(subgraph_id, subgraph_ptr));
|
||||
|
||||
// Reset subgraph and meta_Def
|
||||
sub_var.Reset();
|
||||
}
|
||||
|
||||
Status MKLDNNExecutionProvider::Compile(const std::vector<onnxruntime::Node*>& fused_nodes,
|
||||
std::vector<NodeComputeInfo>& node_compute_funcs) {
|
||||
for (const auto* fused_node : fused_nodes) {
|
||||
auto attributes = fused_node->GetAttributes();
|
||||
NodeComputeInfo compute_info;
|
||||
|
||||
compute_info.create_state_func = [=](ComputeContext* context, FunctionState* state) {
|
||||
auto* p = new onnxruntime::mkl_dnn::MkldnnFuncKernel<float>(context, attributes, this);
|
||||
*state = p;
|
||||
return 0;
|
||||
};
|
||||
|
||||
compute_info.release_state_func = [](FunctionState state) {
|
||||
if (state)
|
||||
delete static_cast<onnxruntime::mkl_dnn::MkldnnFuncKernel<float>*>(state);
|
||||
};
|
||||
|
||||
compute_info.compute_func = [](FunctionState state, const OrtCustomOpApi* api, OrtKernelContext* context) {
|
||||
onnxruntime::mkl_dnn::MkldnnFuncKernel<float>* custom_op = reinterpret_cast<mkl_dnn::MkldnnFuncKernel<float>*>(state);
|
||||
return custom_op->Compute(api, context);
|
||||
};
|
||||
|
||||
node_compute_funcs.push_back(compute_info);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@
|
|||
#include "core/graph/constants.h"
|
||||
#include "core/framework/allocatormgr.h"
|
||||
#include "core/framework/execution_provider.h"
|
||||
#include "core/providers/mkldnn/subgraph/subgraph.h"
|
||||
|
||||
namespace mkldnn {
|
||||
struct memory;
|
||||
|
|
@ -46,7 +47,7 @@ class MKLDNNExecutionProvider : public IExecutionProvider {
|
|||
}
|
||||
|
||||
void SetWeightsMemoryBuffer(const std::string& weight_key,
|
||||
const std::shared_ptr<mkldnn::memory>& filter_dst_mem) {
|
||||
const std::shared_ptr<mkldnn::memory>& filter_dst_mem) {
|
||||
weights_mem_map_.insert(std::make_pair(weight_key, filter_dst_mem));
|
||||
}
|
||||
|
||||
|
|
@ -59,6 +60,13 @@ class MKLDNNExecutionProvider : public IExecutionProvider {
|
|||
reordered_buffers_.push_back(std::move(buffer));
|
||||
}
|
||||
|
||||
std::vector<std::unique_ptr<ComputeCapability>>
|
||||
GetCapability(const onnxruntime::GraphViewer& graph,
|
||||
const std::vector<const KernelRegistry*>& /*kernel_registries*/) const override;
|
||||
|
||||
common::Status Compile(const std::vector<onnxruntime::Node*>& fused_nodes,
|
||||
std::vector<NodeComputeInfo>& node_compute_funcs) override;
|
||||
|
||||
private:
|
||||
// mkldnn weights(filer data) memory blocks from first iteration
|
||||
// saved by weights name
|
||||
|
|
@ -66,6 +74,88 @@ class MKLDNNExecutionProvider : public IExecutionProvider {
|
|||
// Save reordered memory buffers in list so that memory is not freed.
|
||||
std::vector<IAllocatorUniquePtr<void>> reordered_buffers_;
|
||||
OrtMutex mutex_;
|
||||
|
||||
// SUBGRAPH
|
||||
private:
|
||||
static int GetOnnxOpSet(const GraphViewer& graph_viewer) {
|
||||
const auto& dm_to_ver = graph_viewer.DomainToVersionMap();
|
||||
return dm_to_ver.at(kOnnxDomain);
|
||||
}
|
||||
|
||||
std::string GetGraphName(const onnxruntime::GraphViewer& graph_viewer) const {
|
||||
std::string graph_name;
|
||||
|
||||
int opset = GetOnnxOpSet(graph_viewer);
|
||||
|
||||
int index = 0;
|
||||
if (graph_viewer.MaxNodeIndex() > 0) {
|
||||
auto first_node = graph_viewer.GetNode(index);
|
||||
while (first_node == nullptr) {
|
||||
index++;
|
||||
first_node = graph_viewer.GetNode(index);
|
||||
}
|
||||
auto first_node_outputs = first_node->OutputDefs();
|
||||
graph_name = graph_viewer.Name() + "_opset-" + std::to_string(opset) + "_" + first_node_outputs[0]->Name();
|
||||
}
|
||||
return graph_name;
|
||||
}
|
||||
|
||||
bool UseSubgraph(const onnxruntime::GraphViewer& graph_viewer,
|
||||
const std::vector<const KernelRegistry*>& kernel_registries,
|
||||
std::vector<std::unique_ptr<ComputeCapability>>& result) const;
|
||||
|
||||
// Some dimensions are not supported by MKL-DNN
|
||||
// example: Pool with NumDimensions <= 3 is not supported
|
||||
// Fall back to CPU implementation
|
||||
bool IsDimensionSupported(const Node* node) const {
|
||||
bool supported = true;
|
||||
if (node->OpType() == "BatchNormalization") {
|
||||
auto node_inputs = node->InputDefs();
|
||||
if (node_inputs[0]->Shape() != nullptr && node_inputs[0]->Shape()->dim_size() == 3) {
|
||||
supported = false;
|
||||
}
|
||||
}
|
||||
if (node->OpType().find("Pool") != std::string::npos) {
|
||||
auto node_inputs = node->InputDefs();
|
||||
if (node_inputs[0]->Shape() != nullptr && node_inputs[0]->Shape()->dim_size() <= 3) {
|
||||
supported = false;
|
||||
}
|
||||
|
||||
if (node->Op()->SinceVersion() == 10)
|
||||
supported = false;
|
||||
|
||||
if (node->OutputDefs().size() > 1)
|
||||
supported = false;
|
||||
}
|
||||
return supported;
|
||||
}
|
||||
|
||||
void CreateOrUpdateMklDnnNode(const Node* node,
|
||||
std::shared_ptr<mkl_dnn::Subgraph>& subgraph_ptr,
|
||||
mkl_dnn::Subgraph::SubgraphVariables& sub_var,
|
||||
bool fused,
|
||||
std::map<std::string, size_t>& output_to_source_node_map,
|
||||
NodeAttributes& subgraph_attributes) const;
|
||||
|
||||
// Create MklDnn node, update inputs, outputs and parent nodes
|
||||
// collect attribtes
|
||||
void CreateMetaDef(const onnxruntime::GraphViewer& graph_viewer,
|
||||
const NodeAttributes& subgraph_attributes,
|
||||
std::shared_ptr<mkl_dnn::Subgraph>& subgraph_ptr,
|
||||
mkl_dnn::Subgraph::SubgraphVariables& sub_var,
|
||||
std::vector<std::unique_ptr<ComputeCapability>>& result) const;
|
||||
|
||||
public:
|
||||
const std::shared_ptr<mkl_dnn::Subgraph> GetMklDnnSubgraph(const std::string& subgraph_id) {
|
||||
return mkl_subgraphs_[subgraph_id];
|
||||
}
|
||||
|
||||
private:
|
||||
// supported MklDnn Operators
|
||||
std::set<std::string> mkldnn_ops_ = {"Conv", "BatchNormalization", "Relu", "Sum",
|
||||
"AveragePool", "GlobalMaxPool", "GlobalAveragePool", "MaxPool", "LRN"};
|
||||
|
||||
mutable std::unordered_map<std::string, std::shared_ptr<mkl_dnn::Subgraph>> mkl_subgraphs_;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
152
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_activations.h
Normal file
152
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_activations.h
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
// Copyright(C) 2019 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
|
||||
#pragma once
|
||||
#include "core/util/math.h"
|
||||
#include "core/util/math_cpuonly.h"
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/providers/mkldnn/mkldnn_fwd.h"
|
||||
#include "core/providers/mkldnn/mkldnn_execution_provider.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_kernel.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
template <typename T>
|
||||
class MklDnnRelu : public MklDnnKernel {
|
||||
public:
|
||||
MklDnnRelu(const MklDnnNode& node,
|
||||
MKLDNNExecutionProvider* provider,
|
||||
const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") : MklDnnKernel(node, provider) {
|
||||
ORT_UNUSED_PARAMETER(attributes);
|
||||
ORT_UNUSED_PARAMETER(attributes_prefix);
|
||||
}
|
||||
|
||||
Status CreatePrimitives(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
mkldnn::engine& cpu_engine,
|
||||
std::vector<mkldnn::primitive>& net,
|
||||
mkldnn::memory::format& source_format) {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
TensorShape x_shape;
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
|
||||
mkldnn::memory::dims dims(xdim);
|
||||
|
||||
ort_source_format_ = GetSourceFormat(static_cast<int>(xdim));
|
||||
source_format = ort_source_format_;
|
||||
src_format_ = ort_source_format_;
|
||||
|
||||
x_shape = TensorShape(xshape, xdim);
|
||||
|
||||
mkldnn::memory::dims src_dims_mkl(
|
||||
x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), src_format_));
|
||||
src_mem_.reset(
|
||||
new mkldnn::memory({*src_md_, cpu_engine}, nullptr));
|
||||
} else {
|
||||
src_md_.reset(
|
||||
new mkldnn::memory::desc(parents_[0].get()->primitive_dst_mem_.get()->get_primitive_desc().desc()));
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
x_shape = parents_[0].get()->primitive_dst_shape_;
|
||||
ort_source_format_ = source_format;
|
||||
src_format_ = parents_[0].get()->primitive_dst_format_;
|
||||
}
|
||||
|
||||
primitive_dst_shape_ = TensorShape(x_shape);
|
||||
|
||||
mkldnn::memory::dims dst_dims_mkl(primitive_dst_shape_.GetDims().begin(), primitive_dst_shape_.GetDims().end());
|
||||
mkldnn::algorithm algo = mkldnn::algorithm::eltwise_relu;
|
||||
fwd_desc_.reset(new mkldnn::eltwise_forward::desc(
|
||||
mkldnn::prop_kind::forward_inference, algo, *src_md_, 0));
|
||||
|
||||
relu_fwd_pd_.reset(new mkldnn::eltwise_forward::primitive_desc(
|
||||
*fwd_desc_, cpu_engine));
|
||||
|
||||
primitive_src_format_ = static_cast<mkldnn::memory::format>(
|
||||
relu_fwd_pd_.get()->src_primitive_desc().desc().data.format);
|
||||
primitive_dst_format_ = static_cast<mkldnn::memory::format>(
|
||||
relu_fwd_pd_.get()->dst_primitive_desc().desc().data.format);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// last node of sub-graph. need to allocate memory for output_tensor
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
// reorder neded. Use primitive output as input to reorder and
|
||||
// allocate buffer for reorder output, final output of this subgraph
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(relu_fwd_pd_.get()->dst_primitive_desc()));
|
||||
} else {
|
||||
// Last node but re-order not needed. Allocate buffer to output of this node
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(relu_fwd_pd_.get()->dst_primitive_desc(), nullptr));
|
||||
}
|
||||
} else {
|
||||
// Intermediate node. Use mkldnn kernel internal memory for output and
|
||||
// use this as input to next node.
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(relu_fwd_pd_.get()->dst_primitive_desc()));
|
||||
}
|
||||
|
||||
relu_fwd_.reset(
|
||||
new mkldnn::eltwise_forward(*relu_fwd_pd_, *src_mem_, *primitive_dst_mem_));
|
||||
|
||||
net.push_back(*relu_fwd_);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// one of the end nodes. Allocate output buffer memory and
|
||||
// reorder is necessary
|
||||
mkldnn::memory::data_type t = MklDnnType<T>();
|
||||
InitDstReorderOutput(cpu_engine, t, net);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Bind(const OrtCustomOpApi* api, OrtKernelContext* context) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
// Sub-graph's first node. Read input from input buffer
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
src_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
auto& y_dims = primitive_dst_shape_.GetDims();
|
||||
// Allocate memory for output bufffer
|
||||
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
|
||||
T* dst_data = ort.GetTensorMutableData<T>(output);
|
||||
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(dst_data);
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
std::shared_ptr<mkldnn::memory> src_mem_;
|
||||
|
||||
std::unique_ptr<mkldnn::eltwise_forward::desc> fwd_desc_;
|
||||
std::unique_ptr<mkldnn::eltwise_forward::primitive_desc> relu_fwd_pd_;
|
||||
std::unique_ptr<mkldnn::primitive> relu_fwd_;
|
||||
|
||||
std::unique_ptr<mkldnn::memory::desc> src_md_;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
362
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_batchnorm.h
Normal file
362
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_batchnorm.h
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
// Copyright(C) 2019 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
|
||||
#pragma once
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/providers/mkldnn/mkldnn_fwd.h"
|
||||
#include "core/providers/mkldnn/mkldnn_execution_provider.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_kernel.h"
|
||||
#include "core/providers/mkldnn/memcpy_s.h"
|
||||
#include "core/util/math.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
class BatchNormHelper {
|
||||
public:
|
||||
static common::Status ValidateInputs(const TensorShape& xshape,
|
||||
const TensorShape& scale_shape,
|
||||
const TensorShape& b_shape,
|
||||
const TensorShape& mean_shape,
|
||||
const TensorShape& var_shape) {
|
||||
// defined as per spec and used for validation
|
||||
constexpr int kNumInputScaleDimensions = 1;
|
||||
constexpr int kNumInputBiasDimensions = 1;
|
||||
constexpr int kNumInputMeanDimensions = 1;
|
||||
constexpr int kNumInputVarianceDimensions = 1;
|
||||
|
||||
if (xshape.GetDims().empty()) {
|
||||
return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Invalid input X: Empty dimensions");
|
||||
}
|
||||
|
||||
int64_t num_channels = xshape.GetDims()[1];
|
||||
|
||||
if (scale_shape.NumDimensions() != kNumInputScaleDimensions) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input scale: NumDimensions() != ", kNumInputScaleDimensions);
|
||||
}
|
||||
if (scale_shape.GetDims()[0] != num_channels) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input scale: 0th dimension != ", num_channels);
|
||||
}
|
||||
|
||||
if (b_shape.NumDimensions() != kNumInputBiasDimensions) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input B: NumDimensions() != ", kNumInputBiasDimensions);
|
||||
}
|
||||
if (b_shape.GetDims()[0] != num_channels) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input B: 0th dimension != ", num_channels);
|
||||
}
|
||||
|
||||
if (mean_shape.NumDimensions() != kNumInputMeanDimensions) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input mean: NumDimensions() != ", kNumInputMeanDimensions);
|
||||
}
|
||||
if (mean_shape.GetDims()[0] != num_channels) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input mean: 0th dimension != ", num_channels);
|
||||
}
|
||||
|
||||
if (var_shape.NumDimensions() != kNumInputVarianceDimensions) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input var: NumDimensions() != ", kNumInputVarianceDimensions);
|
||||
}
|
||||
if (var_shape.GetDims()[0] != num_channels) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid input var: 0th dimension != ", num_channels);
|
||||
}
|
||||
return common::Status::OK();
|
||||
}
|
||||
|
||||
static void NormalizeDims(const TensorShape& x_shape, std::vector<int64_t>& new_dims) {
|
||||
new_dims.clear();
|
||||
auto& orig_dims = x_shape.GetDims();
|
||||
if (orig_dims.size() == 4 /*supported size by CUDA*/ ||
|
||||
orig_dims.size() == 5 /*supported size by CUDA*/) {
|
||||
new_dims = orig_dims;
|
||||
return;
|
||||
}
|
||||
|
||||
auto rank = x_shape.NumDimensions();
|
||||
auto num_samples = rank > 0 ? orig_dims[0] : 1; // NCHW
|
||||
auto num_channels = rank > 1 ? orig_dims[1] : 1;
|
||||
auto width = rank > 3 ? orig_dims[3] : 1;
|
||||
auto height = rank > 2 ? orig_dims[2] : 1;
|
||||
new_dims = {num_samples, num_channels, height, width};
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class MklDnnBatchNorm : public MklDnnKernel {
|
||||
public:
|
||||
explicit MklDnnBatchNorm(const MklDnnNode& node,
|
||||
MKLDNNExecutionProvider* provider,
|
||||
const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") : MklDnnKernel(node, provider) {
|
||||
ReadAttributes(attributes, attributes_prefix);
|
||||
}
|
||||
void ReadAttributes(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") override {
|
||||
auto attr = attributes.find(attributes_prefix + "epsilon");
|
||||
if (attr != attributes.end() &&
|
||||
attr->second.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT) {
|
||||
epsilon_ = attr->second.f();
|
||||
}
|
||||
}
|
||||
|
||||
Status CreatePrimitives(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
mkldnn::engine& cpu_engine,
|
||||
std::vector<mkldnn::primitive>& net,
|
||||
mkldnn::memory::format& source_format) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
TensorShape x_shape;
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
mkldnn::memory::dims dims(xdim);
|
||||
|
||||
ort_source_format_ = GetSourceFormat(static_cast<int>(xdim));
|
||||
source_format = ort_source_format_;
|
||||
src_format_ = ort_source_format_;
|
||||
x_shape = TensorShape(xshape, xdim);
|
||||
|
||||
mkldnn::memory::dims src_dims_mkl(
|
||||
x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), src_format_));
|
||||
} else {
|
||||
src_md_.reset(new mkldnn::memory::desc(parents_[0].get()->primitive_dst_mem_.get()->get_primitive_desc().desc()));
|
||||
x_shape = parents_[0].get()->primitive_dst_shape_;
|
||||
ort_source_format_ = source_format;
|
||||
src_format_ = parents_[0].get()->primitive_dst_format_;
|
||||
}
|
||||
|
||||
int num_dimensions = static_cast<int>(x_shape.NumDimensions());
|
||||
if (num_dimensions == 3) {
|
||||
primitive_created_ = Status(common::ONNXRUNTIME,
|
||||
common::NOT_IMPLEMENTED, "BatchNorm: Please call default CPU kernel.");
|
||||
return primitive_created_;
|
||||
}
|
||||
|
||||
const OrtValue* scale_input_tensor = ort.KernelContext_GetInput(context, input_index + 1);
|
||||
const OrtValue* b_input_tensor = ort.KernelContext_GetInput(context, input_index + 2);
|
||||
const OrtValue* mean_input_tensor = ort.KernelContext_GetInput(context, input_index + 3);
|
||||
const OrtValue* var_input_tensor = ort.KernelContext_GetInput(context, input_index + 4);
|
||||
|
||||
auto scale_tensor_info = ort.GetTensorTypeAndShape(scale_input_tensor);
|
||||
auto scale_tensor_shape = ort.GetTensorShape(scale_tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(scale_tensor_info);
|
||||
auto sshape = scale_tensor_shape.data();
|
||||
auto sdim = scale_tensor_shape.size();
|
||||
TensorShape scale_shape(sshape, sdim);
|
||||
|
||||
auto b_tensor_info = ort.GetTensorTypeAndShape(b_input_tensor);
|
||||
auto b_tensor_shape = ort.GetTensorShape(b_tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(b_tensor_info);
|
||||
auto bshape = b_tensor_shape.data();
|
||||
auto bdim = b_tensor_shape.size();
|
||||
TensorShape b_shape(bshape, bdim);
|
||||
|
||||
auto mean_tensor_info = ort.GetTensorTypeAndShape(mean_input_tensor);
|
||||
auto mean_tensor_shape = ort.GetTensorShape(mean_tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(mean_tensor_info);
|
||||
auto mshape = mean_tensor_shape.data();
|
||||
auto mdim = mean_tensor_shape.size();
|
||||
TensorShape mean_shape(mshape, mdim);
|
||||
|
||||
auto var_tensor_info = ort.GetTensorTypeAndShape(var_input_tensor);
|
||||
auto var_tensor_shape = ort.GetTensorShape(var_tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(var_tensor_info);
|
||||
auto vshape = var_tensor_shape.data();
|
||||
auto vdim = var_tensor_shape.size();
|
||||
TensorShape var_shape(vshape, vdim);
|
||||
|
||||
primitive_dst_shape_ = TensorShape(x_shape);
|
||||
|
||||
primitive_created_ = BatchNormHelper::ValidateInputs(x_shape, scale_shape, b_shape, mean_shape, var_shape);
|
||||
if (!primitive_created_.IsOK())
|
||||
return primitive_created_;
|
||||
|
||||
mkldnn::memory::dims src_dims_mkl(
|
||||
x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
mkldnn::memory::dims scale_dims_mkl(
|
||||
scale_shape.GetDims().begin(), scale_shape.GetDims().end());
|
||||
mkldnn::memory::dims b_dims_mkl(
|
||||
b_shape.GetDims().begin(), b_shape.GetDims().end());
|
||||
mkldnn::memory::dims mean_dims_mkl(
|
||||
mean_shape.GetDims().begin(), mean_shape.GetDims().end());
|
||||
mkldnn::memory::dims var_dims_mkl(
|
||||
var_shape.GetDims().begin(), var_shape.GetDims().end());
|
||||
|
||||
mkldnn::memory::dims dst_dims_mkl(
|
||||
primitive_dst_shape_.GetDims().begin(), primitive_dst_shape_.GetDims().end());
|
||||
|
||||
scale_shift_md_.reset(new mkldnn::memory::desc(
|
||||
{2, scale_dims_mkl[0]}, MklDnnType<T>(), mkldnn::memory::format::nc));
|
||||
mean_md_.reset(new mkldnn::memory::desc(
|
||||
{mean_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::x));
|
||||
var_md_.reset(new mkldnn::memory::desc(
|
||||
{var_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::x));
|
||||
primitive_dst_md_.reset(new mkldnn::memory::desc(
|
||||
{dst_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
|
||||
// scale_shift_mem will allocate 2*C*sizeof(float) buffer
|
||||
//
|
||||
scale_shift_mem_.reset(
|
||||
new mkldnn::memory({*scale_shift_md_, cpu_engine}));
|
||||
|
||||
mean_mem_.reset(
|
||||
new mkldnn::memory({*mean_md_, cpu_engine}, nullptr));
|
||||
var_mem_.reset(
|
||||
new mkldnn::memory({*var_md_, cpu_engine}, nullptr));
|
||||
|
||||
batchnorm_fwd_.reset(new mkldnn::batch_normalization_forward::desc(
|
||||
mkldnn::prop_kind::forward_inference, *src_md_, epsilon_,
|
||||
mkldnn::batch_normalization_flag::use_scale_shift |
|
||||
mkldnn::batch_normalization_flag::use_global_stats));
|
||||
|
||||
if (fuse_relu_) {
|
||||
mkldnn::primitive_attr attr;
|
||||
attr.set_int_output_round_mode(mkldnn::round_mode::round_nearest);
|
||||
// Execute RELU as Fuse PostOps
|
||||
const float ops_scale = 1.f;
|
||||
const float ops_alpha = 0.f; // relu negative slope
|
||||
const float ops_beta = 0.f;
|
||||
mkldnn::post_ops ops;
|
||||
ops.append_eltwise(ops_scale, mkldnn::algorithm::eltwise_relu, ops_alpha, ops_beta);
|
||||
attr.set_post_ops(ops);
|
||||
|
||||
batchnorm_fwd_pd_.reset(new mkldnn::batch_normalization_forward::primitive_desc(
|
||||
*batchnorm_fwd_, attr, cpu_engine));
|
||||
} else {
|
||||
batchnorm_fwd_pd_.reset(
|
||||
new mkldnn::batch_normalization_forward::primitive_desc(
|
||||
*batchnorm_fwd_, cpu_engine));
|
||||
}
|
||||
|
||||
// out format of this kernel
|
||||
primitive_dst_format_ = static_cast<mkldnn::memory::format>(
|
||||
batchnorm_fwd_pd_.get()->dst_primitive_desc().desc().data.format);
|
||||
primitive_src_format_ = static_cast<mkldnn::memory::format>(
|
||||
batchnorm_fwd_pd_.get()->dst_primitive_desc().desc().data.format);
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
src_mem_.reset(
|
||||
new mkldnn::memory(batchnorm_fwd_pd_.get()->src_primitive_desc(), nullptr));
|
||||
} else {
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// Use mkldnn's internal output buffer
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(batchnorm_fwd_pd_->dst_primitive_desc()));
|
||||
} else {
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(batchnorm_fwd_pd_->dst_primitive_desc(), nullptr));
|
||||
}
|
||||
} else {
|
||||
// last node of sub-graph. need to allocate memory for output_tensor
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(batchnorm_fwd_pd_->dst_primitive_desc()));
|
||||
}
|
||||
auto bn = mkldnn::batch_normalization_forward(
|
||||
*batchnorm_fwd_pd_,
|
||||
(const mkldnn::primitive::at)*src_mem_,
|
||||
(const mkldnn::primitive::at)*mean_mem_,
|
||||
(const mkldnn::primitive::at)*var_mem_,
|
||||
(const mkldnn::memory)*scale_shift_mem_,
|
||||
(const mkldnn::memory)*primitive_dst_mem_);
|
||||
net.push_back(bn);
|
||||
|
||||
// Allocate dst buffer if reorder is necessary
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// one of the end nodes. Allocate output buffer memory and
|
||||
// reorder is necessary
|
||||
mkldnn::memory::data_type t = MklDnnType<T>();
|
||||
InitDstReorderOutput(cpu_engine, t, net);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Bind(const OrtCustomOpApi* api, OrtKernelContext* context) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
if (!primitive_created_.IsOK()) {
|
||||
// abort as MKLDNN cannot execute this. but
|
||||
// ORT try to delete output_tensor buffer data. allocate memory so that it can delete
|
||||
// fix for test_averagepool_1d_default node test
|
||||
return primitive_created_;
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
src_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
}
|
||||
|
||||
const OrtValue* scale_input_tensor = ort.KernelContext_GetInput(context, input_index + 1);
|
||||
const T* scale_data = reinterpret_cast<const T*>(ort.GetTensorData<T>(scale_input_tensor));
|
||||
const OrtValue* b_input_tensor = ort.KernelContext_GetInput(context, input_index + 2);
|
||||
const T* b_data = reinterpret_cast<const T*>(ort.GetTensorData<T>(b_input_tensor));
|
||||
const OrtValue* mean_input_tensor = ort.KernelContext_GetInput(context, input_index + 3);
|
||||
const T* mean_data = reinterpret_cast<const T*>(ort.GetTensorData<T>(mean_input_tensor));
|
||||
const OrtValue* var_input_tensor = ort.KernelContext_GetInput(context, input_index + 4);
|
||||
const T* var_data = reinterpret_cast<const T*>(ort.GetTensorData<T>(var_input_tensor));
|
||||
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(scale_input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto sshape = tensor_shape.data();
|
||||
auto sdim = tensor_shape.size();
|
||||
|
||||
TensorShape scale_shape(sshape, sdim);
|
||||
mkldnn::memory::dims scale_dims_mkl(
|
||||
scale_shape.GetDims().begin(), scale_shape.GetDims().end());
|
||||
|
||||
mean_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(mean_data)));
|
||||
var_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(var_data)));
|
||||
|
||||
T* scale_shift_buf = static_cast<T*>(scale_shift_mem_->get_data_handle());
|
||||
|
||||
size_t src_bytes = sizeof(T) * scale_dims_mkl[0];
|
||||
size_t dst_bytes = sizeof(T) * scale_dims_mkl[0];
|
||||
|
||||
MEMCPY_S(scale_shift_buf, scale_data, src_bytes, dst_bytes);
|
||||
MEMCPY_S(&scale_shift_buf[scale_dims_mkl[0]], b_data, src_bytes, dst_bytes);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
auto& y_dims = primitive_dst_shape_.GetDims();
|
||||
// Allocate memory for output bufffer
|
||||
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
|
||||
T* dst_data = ort.GetTensorMutableData<T>(output);
|
||||
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(dst_data);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
std::shared_ptr<mkldnn::memory> src_mem_;
|
||||
std::unique_ptr<mkldnn::memory> scale_shift_mem_;
|
||||
std::unique_ptr<mkldnn::memory> mean_mem_;
|
||||
std::unique_ptr<mkldnn::memory> var_mem_;
|
||||
std::unique_ptr<mkldnn::memory> dst_mem_;
|
||||
|
||||
std::unique_ptr<mkldnn::memory::desc> src_md_;
|
||||
std::unique_ptr<mkldnn::memory::desc> scale_shift_md_;
|
||||
std::unique_ptr<mkldnn::memory::desc> mean_md_;
|
||||
std::unique_ptr<mkldnn::memory::desc> var_md_;
|
||||
std::unique_ptr<mkldnn::memory::desc> dst_md_;
|
||||
|
||||
std::unique_ptr<mkldnn::batch_normalization_forward::desc> batchnorm_fwd_;
|
||||
std::unique_ptr<mkldnn::batch_normalization_forward::primitive_desc> batchnorm_fwd_pd_;
|
||||
|
||||
protected:
|
||||
float epsilon_ = 1e-5f;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
637
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_conv.h
Normal file
637
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_conv.h
Normal file
|
|
@ -0,0 +1,637 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
#include "mkldnn_types.h"
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/providers/mkldnn/mkldnn_fwd.h"
|
||||
#include "core/providers/cpu/nn/autopad_type.h"
|
||||
#include "core/providers/mkldnn/mkldnn_execution_provider.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_kernel.h"
|
||||
#include "core/util/math.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
// helper function
|
||||
template <bool ForceSymmetricAutoPadding>
|
||||
Status ComputePadAndOutputShape(
|
||||
const int64_t in_dim,
|
||||
const int64_t stride,
|
||||
const int64_t kernel,
|
||||
const int64_t dilation,
|
||||
AutoPadType pad_type,
|
||||
int64_t* pad_head,
|
||||
int64_t* pad_tail,
|
||||
int64_t* out_dim) {
|
||||
const int64_t dkernel = dilation * (kernel - 1) + 1;
|
||||
|
||||
if (pad_type == AutoPadType::NOTSET) {
|
||||
*out_dim = static_cast<int64_t>(static_cast<float>(in_dim + *pad_head + *pad_tail - dkernel) / stride + 1);
|
||||
} else {
|
||||
switch (pad_type) {
|
||||
case AutoPadType::VALID:
|
||||
*pad_head = 0;
|
||||
*pad_tail = 0;
|
||||
*out_dim = (in_dim - dkernel) / stride + 1;
|
||||
break;
|
||||
case AutoPadType::SAME_UPPER:
|
||||
case AutoPadType::SAME_LOWER: {
|
||||
ORT_ENFORCE(dilation == 1, "Dilation not supported for AutoPadType::SAME_UPPER or AutoPadType::SAME_LOWER.");
|
||||
int64_t legacy_target_size = (in_dim + stride - 1) / stride;
|
||||
int64_t pad_needed = (legacy_target_size - 1) * stride + kernel - in_dim;
|
||||
*out_dim = (in_dim + pad_needed - dkernel) / stride + 1;
|
||||
|
||||
// make sure padding is symmetric
|
||||
if (ForceSymmetricAutoPadding)
|
||||
pad_needed = math::roundUpPow2<int64_t, 2>(pad_needed);
|
||||
|
||||
if (pad_type == AutoPadType::SAME_LOWER) {
|
||||
*pad_head = (pad_needed + 1) / 2;
|
||||
} else {
|
||||
*pad_head = pad_needed / 2;
|
||||
}
|
||||
*pad_tail = pad_needed - *pad_head;
|
||||
} break;
|
||||
default:
|
||||
return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "pad type not supported.");
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
class MklDnnConv : public MklDnnKernel {
|
||||
public:
|
||||
MklDnnConv(const MklDnnNode& node,
|
||||
MKLDNNExecutionProvider* provider,
|
||||
const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") : MklDnnKernel(node, provider) {
|
||||
ReadAttributes(attributes, attributes_prefix);
|
||||
}
|
||||
|
||||
Status CreatePrimitives(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
mkldnn::engine& cpu_engine,
|
||||
std::vector<mkldnn::primitive>& net,
|
||||
mkldnn::memory::format& source_format) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
const OrtValue* winput_tensor = ort.KernelContext_GetInput(context, input_index + 1);
|
||||
auto wtensor_info = ort.GetTensorTypeAndShape(winput_tensor);
|
||||
auto wtensor_shape = ort.GetTensorShape(wtensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(wtensor_info);
|
||||
auto wshape = wtensor_shape.data();
|
||||
auto wdim = wtensor_shape.size();
|
||||
|
||||
TensorShape w_shape(wshape, wdim);
|
||||
const int group_mkl = static_cast<int>(group_);
|
||||
|
||||
TensorShape x_shape;
|
||||
// std::unique_ptr<TensorShape> x_shape;
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
|
||||
x_shape = TensorShape(xshape, xdim);
|
||||
|
||||
mkldnn::memory::dims src_dims_mkl(x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
} else {
|
||||
// get the output of previous node (mkldnn block propagation).
|
||||
// TODO Sourcenode will set src of this node.
|
||||
x_shape = parents_[0].get()->primitive_dst_shape_;
|
||||
ort_source_format_ = source_format;
|
||||
src_format_ = parents_[0].get()->primitive_dst_format_;
|
||||
mkldnn::memory::dims src_dims_mkl(x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
}
|
||||
|
||||
primitive_created_ = ValidateInputShape(x_shape, w_shape);
|
||||
if (!primitive_created_.IsOK())
|
||||
return primitive_created_;
|
||||
|
||||
std::vector<int64_t> kernel_shape;
|
||||
primitive_created_ = ComputeKernelShape(w_shape, kernel_shape);
|
||||
if (!primitive_created_.IsOK())
|
||||
return primitive_created_;
|
||||
|
||||
const size_t kernel_rank = kernel_shape.size();
|
||||
|
||||
if (kernel_rank + 2 != wdim) {
|
||||
primitive_created_ = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "kernel_shape num_dims is not compatible with W num_dims.",
|
||||
" kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(),
|
||||
" W: ", w_shape.ToString().c_str());
|
||||
return primitive_created_;
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < kernel_rank; ++i) {
|
||||
if (kernel_shape[i] != w_shape[i + 2]) {
|
||||
primitive_created_ = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "kernel_shape is not compatible with W shape.",
|
||||
" kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(),
|
||||
" W: ", w_shape.ToString().c_str());
|
||||
return primitive_created_;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int64_t> pads(pads_);
|
||||
if (pads.empty()) {
|
||||
pads.resize(kernel_rank * 2, 0);
|
||||
}
|
||||
std::vector<int64_t> dilations(dilations_);
|
||||
if (dilations.empty()) {
|
||||
dilations.resize(kernel_rank, 1);
|
||||
}
|
||||
std::vector<int64_t> strides(strides_);
|
||||
if (strides.empty()) {
|
||||
strides.resize(kernel_rank, 1);
|
||||
}
|
||||
|
||||
const int64_t N = x_shape[0];
|
||||
const int64_t M = w_shape[0];
|
||||
std::vector<int64_t> y_dims;
|
||||
y_dims.insert(y_dims.begin(), {N, M});
|
||||
TensorShape input_shape = x_shape.Slice(2);
|
||||
primitive_created_ = InferOutputShape(input_shape, kernel_shape, strides, dilations, &pads, &y_dims);
|
||||
if (!primitive_created_.IsOK())
|
||||
return primitive_created_;
|
||||
|
||||
TensorShape y_shape(y_dims);
|
||||
primitive_dst_shape_ = TensorShape(y_dims);
|
||||
TensorShape output_shape = y_shape.Slice(2);
|
||||
mkldnn::memory::dims dst_dims_mkl(y_dims.begin(), y_dims.end());
|
||||
primitive_dst_md_.reset(new mkldnn::memory::desc(
|
||||
{dst_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
|
||||
mkldnn::memory::dims filter_dims_mkl;
|
||||
if (group_mkl == 1) {
|
||||
filter_dims_mkl.assign(w_shape.GetDims().begin(), w_shape.GetDims().end());
|
||||
} else {
|
||||
filter_dims_mkl.assign({group_mkl,
|
||||
static_cast<int>(w_shape[0] / group_mkl)});
|
||||
filter_dims_mkl.insert(filter_dims_mkl.end(), w_shape.GetDims().begin() + 1, w_shape.GetDims().end());
|
||||
}
|
||||
mkldnn::memory::dims strides_mkl(strides.begin(), strides.end());
|
||||
mkldnn::memory::dims dilations_mkl(dilations.begin(), dilations.end());
|
||||
// mkldnn dilations start from 0 so we need to subtract 1 from each dim.
|
||||
for (size_t dim = 0; dim < kernel_rank; dim++) {
|
||||
dilations_mkl[dim] -= 1;
|
||||
}
|
||||
|
||||
mkldnn::memory::dims padding_left_mkl(pads.begin(), pads.begin() + kernel_rank);
|
||||
mkldnn::memory::dims padding_right_mkl(pads.begin() + kernel_rank, pads.end());
|
||||
mkldnn::memory::dims bias_dims_mkl;
|
||||
if (mklnode_ptr_->num_inputs == 3) {
|
||||
const OrtValue* binput_tensor = ort.KernelContext_GetInput(context, input_index + 2);
|
||||
auto btensor_info = ort.GetTensorTypeAndShape(binput_tensor);
|
||||
auto btensor_shape = ort.GetTensorShape(btensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(btensor_info);
|
||||
auto bshape = btensor_shape.data();
|
||||
auto bdim = btensor_shape.size();
|
||||
TensorShape b_shape(bshape, bdim);
|
||||
bias_dims_mkl.assign(b_shape.GetDims().begin(), b_shape.GetDims().end());
|
||||
}
|
||||
|
||||
auto fmt = mkldnn::memory::format::any;
|
||||
if (kernel_rank == 1) {
|
||||
fmt = mkldnn::memory::format::ncw;
|
||||
if (group_mkl == 1) {
|
||||
filter_format_ = mkldnn::memory::format::oiw;
|
||||
} else {
|
||||
filter_format_ = mkldnn::memory::format::goiw;
|
||||
}
|
||||
} else if (kernel_rank == 2) {
|
||||
fmt = mkldnn::memory::format::nchw;
|
||||
if (group_mkl == 1) {
|
||||
filter_format_ = mkldnn::memory::format::oihw;
|
||||
} else {
|
||||
filter_format_ = mkldnn::memory::format::goihw;
|
||||
}
|
||||
} else {
|
||||
fmt = mkldnn::memory::format::ncdhw;
|
||||
if (group_mkl == 1) {
|
||||
filter_format_ = mkldnn::memory::format::oidhw;
|
||||
} else {
|
||||
filter_format_ = mkldnn::memory::format::goidhw;
|
||||
}
|
||||
}
|
||||
if (src_format_ == mkldnn::memory::format::any) {
|
||||
src_format_ = fmt;
|
||||
ort_source_format_ = fmt;
|
||||
source_format = fmt;
|
||||
}
|
||||
|
||||
// Set the memory descriptors to format::any to allow MKLDNN to decide what the optimal memory layout should be
|
||||
// for the computation given the input
|
||||
filter_md_.reset(new mkldnn::memory::desc(
|
||||
{filter_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
if (!bias_dims_mkl.empty())
|
||||
bias_md_.reset(new mkldnn::memory::desc(
|
||||
{bias_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
|
||||
if (!bias_dims_mkl.empty()) {
|
||||
fwd_desc_.reset(new mkldnn::convolution_forward::desc(
|
||||
mkldnn::prop_kind::forward, mkldnn::convolution_direct, *src_md_,
|
||||
*filter_md_, *bias_md_, *primitive_dst_md_,
|
||||
strides_mkl, dilations_mkl, padding_left_mkl,
|
||||
padding_right_mkl, mkldnn::padding_kind::zero));
|
||||
} else {
|
||||
fwd_desc_.reset(new mkldnn::convolution_forward::desc(
|
||||
mkldnn::prop_kind::forward, mkldnn::convolution_direct, *src_md_,
|
||||
*filter_md_, *primitive_dst_md_, strides_mkl,
|
||||
dilations_mkl, padding_left_mkl,
|
||||
padding_right_mkl, mkldnn::padding_kind::zero));
|
||||
}
|
||||
|
||||
if (fuse_relu_) {
|
||||
mkldnn::primitive_attr attr;
|
||||
attr.set_int_output_round_mode(mkldnn::round_mode::round_nearest);
|
||||
// Execute RELU as Fuse PostOps
|
||||
const float ops_scale = 1.f;
|
||||
const float ops_alpha = 0.f; // relu negative slope
|
||||
const float ops_beta = 0.f;
|
||||
mkldnn::post_ops ops;
|
||||
ops.append_eltwise(ops_scale, mkldnn::algorithm::eltwise_relu, ops_alpha, ops_beta);
|
||||
attr.set_post_ops(ops);
|
||||
|
||||
conv_fwd_pd_.reset(new mkldnn::convolution_forward::primitive_desc(
|
||||
*fwd_desc_, attr, cpu_engine));
|
||||
} else {
|
||||
conv_fwd_pd_.reset(new mkldnn::convolution_forward::primitive_desc(
|
||||
*fwd_desc_, cpu_engine));
|
||||
}
|
||||
|
||||
primitive_src_format_ = static_cast<mkldnn::memory::format>(
|
||||
conv_fwd_pd_.get()->src_primitive_desc().desc().data.format);
|
||||
|
||||
mkldnn_filter_format_ = static_cast<mkldnn::memory::format>(
|
||||
conv_fwd_pd_.get()->weights_primitive_desc().desc().data.format);
|
||||
|
||||
primitive_dst_format_ = static_cast<mkldnn::memory::format>(
|
||||
conv_fwd_pd_.get()->dst_primitive_desc().desc().data.format);
|
||||
|
||||
src_size_ = conv_fwd_pd_.get()->src_primitive_desc().get_size();
|
||||
filter_size_ = conv_fwd_pd_.get()->weights_primitive_desc().get_size();
|
||||
dst_size_ = conv_fwd_pd_.get()->dst_primitive_desc().get_size();
|
||||
|
||||
filter_mem_.reset(
|
||||
new mkldnn::memory(conv_fwd_pd_.get()->weights_primitive_desc(), nullptr));
|
||||
|
||||
if (primitive_src_format_ != src_format_) {
|
||||
mkldnn::memory::dims src_dims_mkl(x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
auto src_md = mkldnn::memory::desc(src_dims_mkl, MklDnnType<T>(), src_format_);
|
||||
auto pd = mkldnn::memory::primitive_desc(src_md, cpu_engine);
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty())
|
||||
src_mem_from_.reset(new mkldnn::memory(pd, nullptr));
|
||||
else
|
||||
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
|
||||
src_mem_.reset(new mkldnn::memory(conv_fwd_pd_->src_primitive_desc(), nullptr));
|
||||
net.push_back(mkldnn::reorder(*src_mem_from_, *src_mem_));
|
||||
} else {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
src_mem_.reset(new mkldnn::memory(conv_fwd_pd_->src_primitive_desc(), nullptr));
|
||||
} else {
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// Use mkldnn's internal output buffer
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(conv_fwd_pd_.get()->dst_primitive_desc()));
|
||||
} else {
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(conv_fwd_pd_.get()->dst_primitive_desc(), nullptr));
|
||||
}
|
||||
} else {
|
||||
// last node of sub-graph. need to allocate memory for output_tensor
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(conv_fwd_pd_.get()->dst_primitive_desc()));
|
||||
}
|
||||
|
||||
if (!bias_dims_mkl.empty()) {
|
||||
bias_mem_.reset(new mkldnn::memory(conv_fwd_pd_.get()->bias_primitive_desc(), nullptr));
|
||||
conv_fwd_.reset(new mkldnn::convolution_forward(*conv_fwd_pd_, *src_mem_, *filter_mem_,
|
||||
*bias_mem_, *primitive_dst_mem_));
|
||||
} else {
|
||||
conv_fwd_.reset(new mkldnn::convolution_forward(*conv_fwd_pd_, *src_mem_,
|
||||
*filter_mem_, *primitive_dst_mem_));
|
||||
}
|
||||
net.push_back(*conv_fwd_);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// one of the end nodes. Allocate output buffer memory and
|
||||
// reorder is necessary
|
||||
mkldnn::memory::data_type t = MklDnnType<T>();
|
||||
InitDstReorderOutput(cpu_engine, t, net);
|
||||
}
|
||||
primitive_created_ = Status::OK();
|
||||
return primitive_created_;
|
||||
}
|
||||
|
||||
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, mkldnn::engine& cpu_engine) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index + 1);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
|
||||
TensorShape W(xshape, xdim);
|
||||
const T* filter_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
|
||||
const int group_mkl = static_cast<int>(group_);
|
||||
|
||||
mkldnn::memory::dims filter_dims_mkl;
|
||||
if (group_mkl == 1) {
|
||||
filter_dims_mkl.assign(W.GetDims().begin(), W.GetDims().end());
|
||||
} else {
|
||||
filter_dims_mkl.assign({group_mkl,
|
||||
static_cast<int>(W[0] / group_mkl)});
|
||||
filter_dims_mkl.insert(filter_dims_mkl.end(), W.GetDims().begin() + 1, W.GetDims().end());
|
||||
}
|
||||
|
||||
{
|
||||
// lock to make sure reordering is done only once
|
||||
std::lock_guard<OrtMutex> lock(provider_->GetMutex());
|
||||
std::shared_ptr<mkldnn::memory> filter_dst_mem = provider_->GetWeightsMemoryBuffer(mklnode_ptr_->weight_name);
|
||||
|
||||
if (filter_dst_mem == nullptr) {
|
||||
auto pd = mkldnn::memory::primitive_desc(
|
||||
mkldnn::memory::desc(filter_dims_mkl, MklDnnType<T>(), filter_format_), cpu_engine);
|
||||
mkldnn::memory src = mkldnn::memory(pd, (void*)filter_data);
|
||||
IAllocatorUniquePtr<void> filter_reorder_buffer =
|
||||
IAllocator::MakeUniquePtr<void>(alloc_, filter_size_);
|
||||
filter_dst_mem.reset(
|
||||
new mkldnn::memory(conv_fwd_pd_->weights_primitive_desc(), filter_reorder_buffer.get()));
|
||||
|
||||
MemoryReorderParams params(src, *filter_dst_mem);
|
||||
DoReorder<T>(params);
|
||||
provider_->SaveAllocatedMemory(std::move(filter_reorder_buffer));
|
||||
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
|
||||
provider_->SetWeightsMemoryBuffer(mklnode_ptr_->weight_name, filter_dst_mem);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Status Bind(const OrtCustomOpApi* api, OrtKernelContext* context) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
if (!primitive_created_.IsOK()) {
|
||||
// abort as MKLDNN cannot execute this. but
|
||||
// ORT try to delete output_tensor buffer data. allocate memory so that it can delete
|
||||
// fix for test_averagepool_1d_default node test
|
||||
//auto xshape = input_tensors[input_index].shape;
|
||||
//auto xdim = input_tensors[input_index].ndim;
|
||||
//AllocateOutputTensor(output_tensors, mklnode_ptr_->output_index, xshape, xdim, input_tensors[0].dtype);
|
||||
return primitive_created_;
|
||||
}
|
||||
const OrtValue* winput_tensor = ort.KernelContext_GetInput(context, input_index + 1);
|
||||
const T* filter_data = const_cast<T*>(ort.GetTensorData<T>(winput_tensor));
|
||||
|
||||
const T* bias_data = nullptr;
|
||||
if (mklnode_ptr_->num_inputs == 3) {
|
||||
const OrtValue* binput_tensor = ort.KernelContext_GetInput(context, input_index + 2);
|
||||
bias_data = const_cast<T*>(ort.GetTensorData<T>(binput_tensor));
|
||||
}
|
||||
std::shared_ptr<mkldnn::memory> filter_dst_mem = provider_->GetWeightsMemoryBuffer(mklnode_ptr_->weight_name);
|
||||
if (filter_dst_mem == nullptr) {
|
||||
ReorderWeights(api, context, GetEngine());
|
||||
filter_dst_mem = provider_->GetWeightsMemoryBuffer(mklnode_ptr_->weight_name);
|
||||
}
|
||||
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
|
||||
|
||||
filter_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(filter_data)));
|
||||
if (bias_data != nullptr) {
|
||||
bias_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(bias_data)));
|
||||
}
|
||||
|
||||
if (primitive_src_format_ != src_format_) {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
src_mem_from_->set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
} else {
|
||||
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
|
||||
auto src_size = conv_fwd_pd_.get()->src_primitive_desc().get_size();
|
||||
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
|
||||
src_mem_->set_data_handle(src_reorder_buffer_.get());
|
||||
} else {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
src_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
} else {
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
auto& y_dims = primitive_dst_shape_.GetDims();
|
||||
// Allocate memory for output bufffer
|
||||
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
|
||||
T* dst_data = ort.GetTensorMutableData<T>(output);
|
||||
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(dst_data);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
void ReadAttributes(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") override {
|
||||
std::string auto_pad;
|
||||
auto attr = attributes.find(attributes_prefix + "auto_pad");
|
||||
if (attr != attributes.end() &&
|
||||
attr->second.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_STRING) {
|
||||
auto_pad = attr->second.s();
|
||||
}
|
||||
auto_pad_ = (auto_pad != "") ? StringToAutoPadType(auto_pad) : AutoPadType::NOTSET;
|
||||
|
||||
kernel_shape_specified_ = false;
|
||||
attr = attributes.find(attributes_prefix + "kernel_shape");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
Status status = GetIntsAttr(proto, kernel_shape_);
|
||||
kernel_shape_specified_ = true;
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "strides");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
Status status = GetIntsAttr(proto, strides_);
|
||||
}
|
||||
|
||||
bool attr_read = false;
|
||||
attr = attributes.find(attributes_prefix + "pads");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
if (GetIntsAttr(proto, pads_) == Status::OK())
|
||||
attr_read = true;
|
||||
}
|
||||
if (!attr_read) {
|
||||
pads_.resize(kernel_shape_.size() * 2, 0);
|
||||
}
|
||||
|
||||
attr_read = false;
|
||||
attr = attributes.find(attributes_prefix + "dilations");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
if (GetIntsAttr(proto, dilations_) == Status::OK())
|
||||
attr_read = true;
|
||||
}
|
||||
if (!attr_read) {
|
||||
dilations_.resize(kernel_shape_.size(), 1);
|
||||
}
|
||||
|
||||
attr_read = false;
|
||||
attr = attributes.find(attributes_prefix + "group");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
if (GetIntAttr(proto, group_) == Status::OK())
|
||||
attr_read = true;
|
||||
}
|
||||
if (!attr_read) {
|
||||
group_ = 1;
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
mkldnn::memory::format mkldnn_filter_format_;
|
||||
mkldnn::memory::format filter_format_;
|
||||
|
||||
std::shared_ptr<mkldnn::memory> src_mem_from_;
|
||||
std::unique_ptr<mkldnn::memory> src_mem_to_;
|
||||
|
||||
size_t src_size_;
|
||||
size_t filter_size_;
|
||||
size_t dst_size_;
|
||||
|
||||
std::shared_ptr<mkldnn::memory> src_mem_;
|
||||
std::unique_ptr<mkldnn::memory> filter_mem_;
|
||||
std::unique_ptr<mkldnn::memory> bias_mem_;
|
||||
|
||||
std::unique_ptr<mkldnn::convolution_forward::desc> fwd_desc_;
|
||||
|
||||
std::unique_ptr<mkldnn::memory::desc> src_md_;
|
||||
std::unique_ptr<mkldnn::memory::desc> filter_md_;
|
||||
std::unique_ptr<mkldnn::memory::desc> bias_md_;
|
||||
|
||||
std::unique_ptr<mkldnn::convolution_forward::primitive_desc> conv_fwd_pd_;
|
||||
std::unique_ptr<mkldnn::primitive> conv_fwd_;
|
||||
|
||||
private:
|
||||
IAllocatorUniquePtr<void> src_reorder_buffer_;
|
||||
IAllocatorUniquePtr<void> dst_reorder_buffer_;
|
||||
|
||||
private:
|
||||
Status ComputeKernelShape(const TensorShape& weight_shape, std::vector<int64_t>& kernel_shape) const {
|
||||
if (kernel_shape_specified_) {
|
||||
kernel_shape = kernel_shape_;
|
||||
if (kernel_shape.size() + 2 != weight_shape.NumDimensions()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "kernel_shape num_dims is not compatible with W num_dims.",
|
||||
" kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(),
|
||||
" W: ", weight_shape.ToString().c_str());
|
||||
}
|
||||
for (size_t i = 0; i < kernel_shape.size(); ++i) {
|
||||
if (kernel_shape[i] != weight_shape[i + 2]) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "kernel_shape is not compatible with W shape.",
|
||||
" kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(),
|
||||
" W: ", weight_shape.ToString().c_str());
|
||||
}
|
||||
}
|
||||
} else {
|
||||
auto& weight_dims = weight_shape.GetDims();
|
||||
kernel_shape = std::vector<int64_t>(weight_dims.begin() + 2, weight_dims.end());
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status ValidateInputShape(const TensorShape& X, const TensorShape& W) const {
|
||||
const int64_t C = X[1];
|
||||
const int64_t M = W[0];
|
||||
|
||||
if (X.NumDimensions() != W.NumDimensions()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "X num_dims does not match W num_dims.",
|
||||
" X: ", X.ToString().c_str(),
|
||||
" W: ", W.ToString().c_str());
|
||||
}
|
||||
|
||||
if (C != W[1] * group_) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input channels C is not equal to kernel channels * group.",
|
||||
" C: ", C,
|
||||
" kernel channels: ", W[1],
|
||||
" group: ", group_);
|
||||
}
|
||||
|
||||
if (M % group_ != 0) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Output channels M is not divisible by group.",
|
||||
" M: ", M,
|
||||
" group: ", group_);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template <bool ForceSymmetricAutoPadding = false>
|
||||
Status InferOutputShape(const TensorShape& input_shape,
|
||||
const std::vector<int64_t>& kernel_shape,
|
||||
const std::vector<int64_t>& strides,
|
||||
const std::vector<int64_t>& dilations,
|
||||
std::vector<int64_t>* pads,
|
||||
std::vector<int64_t>* output_shape) const {
|
||||
size_t rank = gsl::narrow_cast<int>(input_shape.NumDimensions());
|
||||
for (size_t dim = 0; dim < rank; ++dim) {
|
||||
if (dim >= strides.size() || dim >= kernel_shape.size() ||
|
||||
dim >= dilations.size() || dim >= pads->size() ||
|
||||
rank + dim >= pads->size()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Out of bound access to array");
|
||||
}
|
||||
int64_t dim_size = 0;
|
||||
ORT_RETURN_IF_ERROR(ComputePadAndOutputShape<ForceSymmetricAutoPadding>(
|
||||
input_shape[dim],
|
||||
strides[dim],
|
||||
kernel_shape[dim],
|
||||
dilations[dim],
|
||||
auto_pad_,
|
||||
&pads->at(dim),
|
||||
&pads->at(input_shape.NumDimensions() + dim),
|
||||
&dim_size));
|
||||
if (dim_size <= 0) {
|
||||
return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Invalid input shape: " + input_shape.ToString());
|
||||
}
|
||||
output_shape->push_back(dim_size);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<int64_t> kernel_shape_; // must use ComputeKernelShape(...), instead of kernel_shape_
|
||||
AutoPadType auto_pad_;
|
||||
int64_t group_;
|
||||
bool kernel_shape_specified_;
|
||||
std::vector<int64_t> strides_;
|
||||
std::vector<int64_t> pads_;
|
||||
std::vector<int64_t> dilations_;
|
||||
std::string activation_;
|
||||
float alpha_;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
257
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.cc
Normal file
257
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.cc
Normal file
|
|
@ -0,0 +1,257 @@
|
|||
// Copyright(C) 2018 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
#ifdef _MSC_VER
|
||||
#pragma warning(disable : 4505) //Unreferenced local function has been removed
|
||||
#endif
|
||||
|
||||
#include "mkldnn_func_kernel.h"
|
||||
#include "core/common/exceptions.h"
|
||||
#include "core/session/onnxruntime_cxx_api.h"
|
||||
#include "core/providers/mkldnn/mkldnn_common.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_conv.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_batchnorm.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_activations.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_pool.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_sum.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_lrn.h"
|
||||
#include "core/session/onnxruntime_cxx_api.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
namespace {
|
||||
template <typename T>
|
||||
class SubgraphPrimitive : public PrimitiveBase {
|
||||
public:
|
||||
SubgraphPrimitive(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
const SubgraphParams& params)
|
||||
: cpu_engine_(GetEngine()) {
|
||||
context_.stream.reset(new mkldnn::stream(mkldnn::stream::kind::eager));
|
||||
|
||||
if (context_.net.size() == 0) {
|
||||
CreateKernels(params);
|
||||
Initialize(api, context);
|
||||
}
|
||||
}
|
||||
|
||||
void UpdateProvider(const SubgraphParams& params) {
|
||||
if (context_.kernels.size() > 0 && context_.kernels[0]->GetProvider() != params.provider)
|
||||
for (auto& kernel : context_.kernels) {
|
||||
kernel->SetProvider(params.provider);
|
||||
}
|
||||
}
|
||||
|
||||
Status Compute(const OrtCustomOpApi* api, OrtKernelContext* context) {
|
||||
Status status;
|
||||
|
||||
for (auto& kernel : context_.kernels) {
|
||||
status = kernel->Bind(api, context);
|
||||
if (!status.IsOK())
|
||||
break;
|
||||
}
|
||||
if (status.IsOK())
|
||||
context_.stream->submit(context_.net);
|
||||
return status;
|
||||
}
|
||||
|
||||
~SubgraphPrimitive() = default;
|
||||
|
||||
private:
|
||||
void CreateKernels(const SubgraphParams& params) {
|
||||
for (const auto& mkldnn_node : params.subgraph->mkldnn_nodes) {
|
||||
if (mkldnn_node.name == "Conv") {
|
||||
std::ostringstream os;
|
||||
os << "Conv-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnConv<T>> kernel;
|
||||
kernel.reset(new MklDnnConv<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "Conv-Relu") {
|
||||
std::ostringstream os;
|
||||
os << "Conv-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnConv<T>> kernel;
|
||||
kernel.reset(new MklDnnConv<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
kernel->fuse_relu_ = true;
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "Relu") {
|
||||
std::ostringstream os;
|
||||
os << "Relu-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnRelu<T>> kernel;
|
||||
kernel.reset(new MklDnnRelu<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "BatchNormalization") {
|
||||
std::ostringstream os;
|
||||
os << "BatchNormalization-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnBatchNorm<T>> kernel;
|
||||
kernel.reset(new MklDnnBatchNorm<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "BatchNormalization-Relu") {
|
||||
std::ostringstream os;
|
||||
os << "BatchNormalization-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnBatchNorm<T>> kernel;
|
||||
kernel.reset(new MklDnnBatchNorm<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
kernel->fuse_relu_ = true;
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "MaxPool") {
|
||||
std::ostringstream os;
|
||||
os << "MaxPool-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnPool<T>> kernel;
|
||||
kernel.reset(new MklDnnPool<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "GlobalMaxPool") {
|
||||
std::ostringstream os;
|
||||
os << "GlobalMaxPool-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnPool<T>> kernel;
|
||||
kernel.reset(new MklDnnPool<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "AveragePool") {
|
||||
std::ostringstream os;
|
||||
os << "AveragePool-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnPool<T>> kernel;
|
||||
kernel.reset(new MklDnnPool<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "GlobalAveragePool") {
|
||||
std::ostringstream os;
|
||||
os << "GlobalAveragePool-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnPool<T>> kernel;
|
||||
kernel.reset(new MklDnnPool<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "Sum") {
|
||||
std::ostringstream os;
|
||||
os << "Sum-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnSum<T>> kernel;
|
||||
kernel.reset(new MklDnnSum<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
} else if (mkldnn_node.name == "LRN") {
|
||||
std::ostringstream os;
|
||||
os << "LRN-" << mkldnn_node.node_index << "-";
|
||||
std::shared_ptr<MklDnnLrn<T>> kernel;
|
||||
kernel.reset(new MklDnnLrn<T>(mkldnn_node, params.provider, params.attributes, os.str()));
|
||||
for (auto index : mkldnn_node.parent_nodes) {
|
||||
kernel->parents_.push_back(context_.kernels[index]);
|
||||
}
|
||||
context_.kernels.push_back(kernel);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
struct SubgraphContext {
|
||||
std::unique_ptr<mkldnn::stream> stream;
|
||||
std::vector<mkldnn::primitive> net;
|
||||
std::vector<std::shared_ptr<MklDnnKernel>> kernels;
|
||||
|
||||
SubgraphContext() : stream(nullptr) {}
|
||||
};
|
||||
|
||||
Status Initialize(const OrtCustomOpApi* api, OrtKernelContext* context) {
|
||||
// Propagate mkldnn block format
|
||||
// dst format of current node to src format of next node
|
||||
mkldnn::memory::format source_format = mkldnn::memory::format::any; // ONNXRuntime format
|
||||
for (auto& kernel : context_.kernels) {
|
||||
Status status = kernel->CreatePrimitives(api, context, cpu_engine_, context_.net, source_format);
|
||||
if (status.IsOK())
|
||||
kernel->ReorderWeights(api, context, cpu_engine_);
|
||||
else
|
||||
return status;
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
SubgraphContext context_;
|
||||
mkldnn::engine& cpu_engine_;
|
||||
};
|
||||
|
||||
// Pool which allows for reuse of MKLDNN Conv primitives which are expensive to instantiate.
|
||||
// To address thread safety, the primitives are stored in a map on thread local storage.
|
||||
template <typename T>
|
||||
class SubgraphPrimitivePool : public PrimitivePool<T> {
|
||||
public:
|
||||
static SubgraphPrimitive<T>* Get(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
const SubgraphParams& params) {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
std::string dims_str;
|
||||
for (auto i = 0; i < params.subgraph->mkldnn_nodes[0].num_inputs; i++) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, i);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto shape = tensor_shape.data();
|
||||
auto dim = tensor_shape.size();
|
||||
|
||||
TensorShape x_shape(shape, dim);
|
||||
mkldnn::memory::dims src_dims(x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
AddDimsToKey(dims_str, src_dims);
|
||||
}
|
||||
|
||||
SubgraphPrimitive<T>* primitive = dynamic_cast<SubgraphPrimitive<T>*>(
|
||||
SubgraphPrimitivePool<T>::GetInstance().GetPrimitive(params.subgraph_key + dims_str));
|
||||
|
||||
if (primitive == nullptr) {
|
||||
auto subgraph_primitive = std::make_unique<SubgraphPrimitive<T>>(api, context, params);
|
||||
primitive = subgraph_primitive.get();
|
||||
SubgraphPrimitivePool<T>::GetInstance().SetPrimitive(params.subgraph_key + dims_str, std::move(subgraph_primitive));
|
||||
}
|
||||
return primitive;
|
||||
}
|
||||
|
||||
private:
|
||||
SubgraphPrimitivePool() = default;
|
||||
~SubgraphPrimitivePool() = default;
|
||||
|
||||
static SubgraphPrimitivePool& GetInstance() {
|
||||
static SubgraphPrimitivePool pool;
|
||||
return pool;
|
||||
}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
template <typename T>
|
||||
Status MkldnnFuncKernel<T>::Compute(const OrtCustomOpApi* api, OrtKernelContext* context) const {
|
||||
Status status;
|
||||
try {
|
||||
SubgraphPrimitive<T>* primitive = SubgraphPrimitivePool<T>::Get(api, context, params_);
|
||||
primitive->UpdateProvider(params_);
|
||||
status = primitive->Compute(api, context);
|
||||
} catch (const mkldnn::error& e) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Status: ", e.status,
|
||||
", message: ", e.message.c_str());
|
||||
}
|
||||
return status;
|
||||
}
|
||||
|
||||
template class MkldnnFuncKernel<float>;
|
||||
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
234
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.h
Normal file
234
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.h
Normal file
|
|
@ -0,0 +1,234 @@
|
|||
// Copyright(C) 2018 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
#pragma once
|
||||
|
||||
#include "core/graph/onnx_protobuf.h"
|
||||
#include "core/providers/mkldnn/mkldnn_execution_provider.h"
|
||||
#include "core/session/onnxruntime_c_api.h"
|
||||
#include "core/framework/func_api.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
namespace {
|
||||
struct SubgraphParams {
|
||||
NodeAttributes attributes;
|
||||
MKLDNNExecutionProvider* provider;
|
||||
std::shared_ptr<Subgraph> subgraph;
|
||||
std::string subgraph_id;
|
||||
std::string subgraph_key;
|
||||
|
||||
SubgraphParams() {}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
template <typename T>
|
||||
class MkldnnFuncKernel {
|
||||
public:
|
||||
explicit MkldnnFuncKernel(const ComputeContext* context,
|
||||
const NodeAttributes& attributes,
|
||||
MKLDNNExecutionProvider* provider) {
|
||||
ORT_UNUSED_PARAMETER(context);
|
||||
|
||||
params_.provider = provider;
|
||||
params_.attributes = attributes;
|
||||
|
||||
auto sub_it = attributes.find("subgraph_id");
|
||||
if (sub_it->second.type() == ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_STRING) {
|
||||
params_.subgraph_id = sub_it->second.s();
|
||||
params_.subgraph = provider->GetMklDnnSubgraph(params_.subgraph_id);
|
||||
|
||||
std::ostringstream key_os;
|
||||
key_os << params_.subgraph->graph_name << "_" << params_.subgraph_id << "-";
|
||||
key_os << params_.subgraph->mkldnn_nodes.back().ToString() << "-";
|
||||
key_os << params_.subgraph->mkldnn_nodes.back().output_name;
|
||||
|
||||
if (params_.subgraph->mkldnn_nodes[0].name == "Conv") {
|
||||
std::ostringstream os;
|
||||
os << "Conv-" << params_.subgraph->mkldnn_nodes[0].node_index << "-";
|
||||
key_os << GetConvAttributeKey(attributes, os.str());
|
||||
}
|
||||
|
||||
if (params_.subgraph->mkldnn_nodes[0].name == "LRN") {
|
||||
std::ostringstream os;
|
||||
os << "LRN-" << params_.subgraph->mkldnn_nodes[0].node_index << "-";
|
||||
key_os << GetLrnAttributeKey(attributes, os.str());
|
||||
}
|
||||
|
||||
if (params_.subgraph->mkldnn_nodes[0].name.find("Pool") != std::string::npos) {
|
||||
std::ostringstream os;
|
||||
os << params_.subgraph->mkldnn_nodes[0].name << "-" << params_.subgraph->mkldnn_nodes[0].node_index << "-";
|
||||
key_os << GetPoolAttributesKey(attributes, os.str());
|
||||
}
|
||||
|
||||
params_.subgraph_key = key_os.str();
|
||||
}
|
||||
}
|
||||
|
||||
std::string GetPoolAttributesKey(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") {
|
||||
std::string key;
|
||||
|
||||
auto attr = attributes.find(attributes_prefix + "kernel_shape");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
key.append(std::to_string(proto.ints(i)));
|
||||
key.append(1, '_');
|
||||
}
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "auto_pad");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(attr->second.s());
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "pads");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
key.append(std::to_string(proto.ints(i)));
|
||||
key.append(1, '_');
|
||||
}
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "strides");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
key.append(std::to_string(proto.ints(i)));
|
||||
key.append(1, '_');
|
||||
}
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "count_include_pad");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
key.append(std::to_string(proto.i()));
|
||||
key.append(1, '_');
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "ceil_mode");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
key.append(std::to_string(proto.i()));
|
||||
key.append(1, '_');
|
||||
key.append(1, '#');
|
||||
}
|
||||
return key;
|
||||
}
|
||||
|
||||
std::string GetConvAttributeKey(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") {
|
||||
std::string key;
|
||||
|
||||
auto attr = attributes.find(attributes_prefix + "dilations");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
key.append(std::to_string(proto.ints(i)));
|
||||
key.append(1, '_');
|
||||
}
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "auto_pad");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
key.append(proto.s());
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "pads");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
key.append(std::to_string(proto.ints(i)));
|
||||
key.append(1, '_');
|
||||
}
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "strides");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
key.append(std::to_string(proto.ints(i)));
|
||||
key.append(1, '_');
|
||||
}
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "kernel_shape");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
key.append(std::to_string(proto.ints(i)));
|
||||
key.append(1, '_');
|
||||
}
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "group");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
key.append(std::to_string(proto.i()));
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
return key;
|
||||
}
|
||||
|
||||
std::string GetLrnAttributeKey(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") {
|
||||
std::string key;
|
||||
|
||||
auto attr = attributes.find(attributes_prefix + "alpha");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
key.append(std::to_string(proto.f()));
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "beta");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
key.append(std::to_string(proto.f()));
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "bias");
|
||||
if (attr != attributes.end()) {
|
||||
key.append(1, '#');
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
key.append(std::to_string(proto.f()));
|
||||
key.append(1, '#');
|
||||
}
|
||||
|
||||
return key;
|
||||
}
|
||||
|
||||
Status Compute(const OrtCustomOpApi* api, OrtKernelContext* context) const;
|
||||
|
||||
private:
|
||||
SubgraphParams params_;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
57
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.cc
Normal file
57
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.cc
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
// Copyright(C) 2019 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
|
||||
#include "mkldnn_kernel.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
void MklDnnKernel::InitDstReorderOutput(mkldnn::engine& cpu_engine,
|
||||
mkldnn::memory::data_type& data_type,
|
||||
std::vector<mkldnn::primitive>& net) {
|
||||
// Allocate dst buffer if reorder is necessary
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
// reorder to ONNXRuntime format
|
||||
mkldnn::memory::dims dst_dims_mkl(
|
||||
primitive_dst_shape_.GetDims().begin(), primitive_dst_shape_.GetDims().end());
|
||||
mkldnn::memory::desc dst_des = mkldnn::memory::desc(dst_dims_mkl,
|
||||
data_type, ort_source_format_);
|
||||
reorder_dst_mem_to_.reset(new mkldnn::memory({dst_des, cpu_engine}, nullptr));
|
||||
net.push_back(mkldnn::reorder(*primitive_dst_mem_, *reorder_dst_mem_to_));
|
||||
}
|
||||
}
|
||||
|
||||
mkldnn::memory::format MklDnnKernel::GetSourceFormat(int dim_size) {
|
||||
mkldnn::memory::format source_format = mkldnn::memory::format::any;
|
||||
switch (dim_size) {
|
||||
case 1: {
|
||||
source_format = mkldnn::memory::format::x;
|
||||
break;
|
||||
}
|
||||
case 2: {
|
||||
source_format = mkldnn::memory::format::nc;
|
||||
break;
|
||||
}
|
||||
case 3: {
|
||||
source_format = mkldnn::memory::format::ntc;
|
||||
break;
|
||||
}
|
||||
case 4: {
|
||||
source_format = mkldnn::memory::format::nchw;
|
||||
break;
|
||||
}
|
||||
case 5: {
|
||||
source_format = mkldnn::memory::format::ncdhw;
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
source_format = mkldnn::memory::format::any;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return source_format;
|
||||
}
|
||||
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
121
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.h
Normal file
121
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.h
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
// Copyright(C) 2019 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
|
||||
#pragma once
|
||||
#ifdef _WIN32
|
||||
#pragma warning(disable : 4244)
|
||||
#endif
|
||||
|
||||
#include "mkldnn.hpp"
|
||||
#include "core/common/cpuid_info.h"
|
||||
#include "core/session/onnxruntime_cxx_api.h"
|
||||
#include "core/providers/mkldnn/subgraph/subgraph.h"
|
||||
#include "core/providers/mkldnn/mkldnn_execution_provider.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
class MklDnnKernel {
|
||||
public:
|
||||
MklDnnKernel(const MklDnnNode& node,
|
||||
MKLDNNExecutionProvider* provider) {
|
||||
mklnode_ptr_ = std::make_shared<MklDnnNode>(node);
|
||||
provider_ = provider;
|
||||
alloc_ = provider_->GetAllocator(0, OrtMemTypeDefault);
|
||||
}
|
||||
virtual ~MklDnnKernel(){};
|
||||
|
||||
virtual Status CreatePrimitives(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
mkldnn::engine& cpu_engine,
|
||||
std::vector<mkldnn::primitive>& net,
|
||||
mkldnn::memory::format& src_fmt) = 0;
|
||||
|
||||
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, mkldnn::engine& cpu_engine) {
|
||||
ORT_UNUSED_PARAMETER(api);
|
||||
ORT_UNUSED_PARAMETER(context);
|
||||
ORT_UNUSED_PARAMETER(cpu_engine);
|
||||
}
|
||||
virtual void SetProvider(MKLDNNExecutionProvider* provider) {
|
||||
provider_ = provider;
|
||||
}
|
||||
MKLDNNExecutionProvider* GetProvider() {
|
||||
return provider_;
|
||||
}
|
||||
|
||||
virtual Status Bind(const OrtCustomOpApi* api, OrtKernelContext* context) = 0;
|
||||
|
||||
protected:
|
||||
virtual void ReadAttributes(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") {
|
||||
ORT_UNUSED_PARAMETER(attributes);
|
||||
ORT_UNUSED_PARAMETER(attributes_prefix);
|
||||
}
|
||||
|
||||
Status GetIntsAttr(ONNX_NAMESPACE::AttributeProto& proto, std::vector<int64_t>& values) {
|
||||
ORT_RETURN_IF_NOT(proto.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_INTS);
|
||||
values.reserve(proto.ints_size());
|
||||
for (int i = 0; i < proto.ints_size(); i++) {
|
||||
values.push_back(proto.ints(i));
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status GetIntAttr(ONNX_NAMESPACE::AttributeProto& proto, int64_t& value) {
|
||||
ORT_RETURN_IF_NOT(proto.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_INT);
|
||||
value = proto.i();
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status GetFloatAttr(ONNX_NAMESPACE::AttributeProto& proto, float& value) {
|
||||
ORT_RETURN_IF_NOT(proto.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT);
|
||||
value = proto.f();
|
||||
return Status::OK();
|
||||
}
|
||||
Status GetStringAttr(ONNX_NAMESPACE::AttributeProto& proto, std::string& value) {
|
||||
ORT_RETURN_IF_NOT(proto.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_STRING);
|
||||
value = proto.s();
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
void InitDstReorderOutput(mkldnn::engine& cpu_engine,
|
||||
mkldnn::memory::data_type& data_type,
|
||||
std::vector<mkldnn::primitive>& net);
|
||||
|
||||
//void AllocateMemoryAndReorderIfNeeded(const OrtCustomOpApi* api);
|
||||
|
||||
//void AllocateOutputTensor(const OrtCustomOpApi* api, int index, const int64_t* shape, size_t dim);
|
||||
|
||||
mkldnn::memory::format GetSourceFormat(int dim_size);
|
||||
|
||||
public:
|
||||
std::vector<std::shared_ptr<MklDnnKernel>> parents_;
|
||||
bool fuse_relu_ = false;
|
||||
bool fuse_sum_ = false;
|
||||
std::shared_ptr<mkldnn::memory> primitive_dst_mem_;
|
||||
std::unique_ptr<mkldnn::memory::desc> primitive_dst_md_;
|
||||
TensorShape primitive_dst_shape_;
|
||||
mkldnn::memory::format primitive_dst_format_ = mkldnn::memory::format::any;
|
||||
|
||||
protected:
|
||||
// ONNX Runtime format
|
||||
mkldnn::memory::format ort_source_format_ = mkldnn::memory::format::any;
|
||||
// input format.
|
||||
// It can be ORT format (nchw) or blocked memory format from parent nce
|
||||
mkldnn::memory::format src_format_ = mkldnn::memory::format::any;
|
||||
// Pointer to MklNode of subgraph IR
|
||||
std::shared_ptr<MklDnnNode> mklnode_ptr_;
|
||||
// input format expected by primitive object
|
||||
mkldnn::memory::format primitive_src_format_ = mkldnn::memory::format::any;
|
||||
|
||||
// memory used for reorders
|
||||
std::unique_ptr<mkldnn::memory> reorder_dst_mem_to_;
|
||||
|
||||
protected:
|
||||
Status primitive_created_;
|
||||
AllocatorPtr alloc_;
|
||||
MKLDNNExecutionProvider* provider_;
|
||||
};
|
||||
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
183
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_lrn.h
Normal file
183
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_lrn.h
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
#include "core/util/math.h"
|
||||
#include "core/util/math_cpuonly.h"
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/providers/mkldnn/mkldnn_fwd.h"
|
||||
#include "core/providers/mkldnn/mkldnn_execution_provider.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_kernel.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
template <typename T>
|
||||
class MklDnnLrn : public MklDnnKernel {
|
||||
public:
|
||||
MklDnnLrn(const MklDnnNode& node,
|
||||
MKLDNNExecutionProvider* provider,
|
||||
const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") : MklDnnKernel(node, provider) {
|
||||
ReadAttributes(attributes, attributes_prefix);
|
||||
}
|
||||
|
||||
Status CreatePrimitives(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
mkldnn::engine& cpu_engine,
|
||||
std::vector<mkldnn::primitive>& net,
|
||||
mkldnn::memory::format& source_format) {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
TensorShape x_shape;
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
|
||||
ort_source_format_ = GetSourceFormat(static_cast<int>(xdim));
|
||||
source_format = ort_source_format_;
|
||||
x_shape = TensorShape(xshape, xdim);
|
||||
|
||||
mkldnn::memory::dims src_dims_mkl(
|
||||
x_shape.GetDims().begin(), x_shape.GetDims().end());
|
||||
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), source_format));
|
||||
src_mem_.reset(
|
||||
new mkldnn::memory({*src_md_, cpu_engine}, nullptr));
|
||||
} else {
|
||||
src_md_.reset(
|
||||
new mkldnn::memory::desc(parents_[0].get()->primitive_dst_mem_.get()->get_primitive_desc().desc()));
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
x_shape = parents_[0].get()->primitive_dst_shape_;
|
||||
ort_source_format_ = source_format;
|
||||
}
|
||||
|
||||
primitive_dst_shape_ = TensorShape(x_shape);
|
||||
|
||||
mkldnn::algorithm algo = mkldnn::algorithm::lrn_across_channels;
|
||||
fwd_desc_.reset(new mkldnn::lrn_forward::desc(
|
||||
mkldnn::prop_kind::forward_scoring, algo, *src_md_,
|
||||
size_, alpha_, beta_, bias_));
|
||||
|
||||
fwd_primitive_desc_.reset(new mkldnn::lrn_forward::primitive_desc(
|
||||
*fwd_desc_, cpu_engine));
|
||||
|
||||
primitive_src_format_ = static_cast<mkldnn::memory::format>(
|
||||
fwd_primitive_desc_.get()->src_primitive_desc().desc().data.format);
|
||||
primitive_dst_format_ = static_cast<mkldnn::memory::format>(
|
||||
fwd_primitive_desc_.get()->dst_primitive_desc().desc().data.format);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// last node of sub-graph. need to allocate memory for output_tensor
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
// reorder neded. Use primitive output as input to reorder and
|
||||
// allocate buffer for reorder output, final output of this subgraph
|
||||
primitive_dst_mem_.reset(
|
||||
new mkldnn::memory(fwd_primitive_desc_.get()->dst_primitive_desc()));
|
||||
} else {
|
||||
// Last node but re-order not needed. Allocate buffer to output of this node
|
||||
primitive_dst_mem_.reset(
|
||||
new mkldnn::memory(fwd_primitive_desc_.get()->dst_primitive_desc(), nullptr));
|
||||
}
|
||||
} else {
|
||||
// Intermediate node. Use mkldnn kernel internal memory for output and
|
||||
// use this as input to next node.
|
||||
primitive_dst_mem_.reset(
|
||||
new mkldnn::memory(fwd_primitive_desc_.get()->dst_primitive_desc()));
|
||||
}
|
||||
|
||||
lrn_fwd_.reset(
|
||||
new mkldnn::lrn_forward(*fwd_primitive_desc_, *src_mem_, *primitive_dst_mem_));
|
||||
net.push_back(*lrn_fwd_);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// one of the end nodes. Allocate output buffer memory and
|
||||
// reorder is necessary
|
||||
mkldnn::memory::data_type t = MklDnnType<T>();
|
||||
InitDstReorderOutput(cpu_engine, t, net);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Bind(const OrtCustomOpApi* api, OrtKernelContext* context) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
// Sub-graph's first node. Read input from input buffer
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
src_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
auto& y_dims = primitive_dst_shape_.GetDims();
|
||||
// Allocate memory for output bufffer
|
||||
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
|
||||
T* dst_data = ort.GetTensorMutableData<T>(output);
|
||||
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(dst_data);
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
void ReadAttributes(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") override {
|
||||
auto attr = attributes.find(attributes_prefix + "size");
|
||||
if (attr != attributes.end() &&
|
||||
attr->second.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_INT) {
|
||||
size_ = attr->second.i();
|
||||
}
|
||||
ORT_ENFORCE(size_ > 0);
|
||||
ORT_ENFORCE(size_ % 2 == 1);
|
||||
|
||||
attr = attributes.find(attributes_prefix + "alpha");
|
||||
if (attr != attributes.end() &&
|
||||
attr->second.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT) {
|
||||
alpha_ = attr->second.f();
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "beta");
|
||||
if (attr != attributes.end() &&
|
||||
attr->second.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT) {
|
||||
beta_ = attr->second.f();
|
||||
}
|
||||
|
||||
bias_ = 1.0f;
|
||||
attr = attributes.find(attributes_prefix + "bias");
|
||||
if (attr != attributes.end() &&
|
||||
attr->second.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT) {
|
||||
bias_ = attr->second.f();
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
float alpha_ = 0;
|
||||
float beta_ = 0;
|
||||
float bias_ = 0;
|
||||
int size_ = 0;
|
||||
|
||||
private:
|
||||
std::shared_ptr<mkldnn::memory> src_mem_;
|
||||
|
||||
std::unique_ptr<mkldnn::lrn_forward::desc> fwd_desc_;
|
||||
std::unique_ptr<mkldnn::lrn_forward::primitive_desc> fwd_primitive_desc_;
|
||||
std::unique_ptr<mkldnn::primitive> lrn_fwd_;
|
||||
|
||||
std::unique_ptr<mkldnn::memory::desc> src_md_;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
427
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_pool.h
Normal file
427
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_pool.h
Normal file
|
|
@ -0,0 +1,427 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
#include "core/providers/mkldnn/mkldnn_fwd.h"
|
||||
#include "core/providers/cpu/nn/autopad_type.h"
|
||||
#include "core/providers/mkldnn/mkldnn_execution_provider.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_kernel.h"
|
||||
#include "core/util/math.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
template <typename T>
|
||||
class MklDnnPool : public MklDnnKernel {
|
||||
public:
|
||||
MklDnnPool(const MklDnnNode& node,
|
||||
MKLDNNExecutionProvider* provider,
|
||||
const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") : MklDnnKernel(node, provider) {
|
||||
op_name_ = node.name;
|
||||
ReadAttributes(attributes, attributes_prefix);
|
||||
}
|
||||
|
||||
Status CreatePrimitives(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
mkldnn::engine& cpu_engine, std::vector<mkldnn::primitive>& net,
|
||||
mkldnn::memory::format& source_format) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
|
||||
mkldnn::memory::dims dims(xdim);
|
||||
ort_source_format_ = GetSourceFormat(static_cast<int>(xdim));
|
||||
source_format = ort_source_format_;
|
||||
src_format_ = ort_source_format_;
|
||||
x_shape_ = TensorShape(xshape, xdim);
|
||||
|
||||
mkldnn::memory::dims src_dims_mkl(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
|
||||
|
||||
// reorder for better performance
|
||||
mkldnn::memory::format fmt = GetAVXFormat(src_dims_mkl);
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), fmt));
|
||||
} else {
|
||||
// get the output of previous node (mkldnn block propagation).
|
||||
// TODO Sourcenode will set src of this node.
|
||||
x_shape_ = parents_[0].get()->primitive_dst_shape_;
|
||||
ort_source_format_ = source_format;
|
||||
src_format_ = parents_[0].get()->primitive_dst_format_;
|
||||
mkldnn::memory::dims src_dims_mkl(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
|
||||
|
||||
if (src_format_ == ort_source_format_) {
|
||||
// reorder for better performance
|
||||
mkldnn::memory::format fmt = GetAVXFormat(src_dims_mkl);
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), fmt));
|
||||
} else {
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
parents_[0].get()->primitive_dst_mem_.get()->get_primitive_desc().desc()));
|
||||
}
|
||||
}
|
||||
|
||||
const auto& x_dims = x_shape_.GetDims();
|
||||
std::vector<int64_t> y_dims = SetOutputSize(x_shape_, x_shape_[1], &pads_);
|
||||
primitive_dst_shape_ = TensorShape(y_dims);
|
||||
|
||||
if (x_shape_.NumDimensions() <= 3) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Please call default CPU kernel.");
|
||||
}
|
||||
|
||||
if (global_pooling_) {
|
||||
kernel_shape_.assign(x_dims.begin() + 2, x_dims.end());
|
||||
pads_.assign(kernel_shape_.size() * 2, 0);
|
||||
strides_.assign(kernel_shape_.size(), 1);
|
||||
}
|
||||
|
||||
size_t num_outputs = 1; //OpKernel::Node().OutputDefs().size(); TODO
|
||||
if (num_outputs == 2) {
|
||||
ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "can not call cpu default op");
|
||||
}
|
||||
|
||||
mkldnn::memory::dims dst_dims_mkl(y_dims.begin(), y_dims.end());
|
||||
mkldnn::memory::dims kernel_mkl(kernel_shape_.begin(), kernel_shape_.end());
|
||||
mkldnn::memory::dims strides_mkl(strides_.begin(), strides_.end());
|
||||
mkldnn::memory::dims padding_left_mkl(pads_.begin(), pads_.begin() + (pads_.size() / 2));
|
||||
mkldnn::memory::dims padding_right_mkl(pads_.begin() + (pads_.size() / 2), pads_.end());
|
||||
|
||||
primitive_dst_md_.reset(new mkldnn::memory::desc(
|
||||
{dst_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
|
||||
mkldnn::algorithm algo = mkldnn::algorithm::pooling_max;
|
||||
if (op_name_ == "AveragePool" || op_name_ == "GlobalAveragePool") {
|
||||
algo = mkldnn::algorithm::pooling_avg_exclude_padding;
|
||||
if (count_include_pad_) {
|
||||
algo = mkldnn::algorithm::pooling_avg_include_padding;
|
||||
}
|
||||
}
|
||||
fwd_desc_.reset(new mkldnn::pooling_forward::desc(
|
||||
mkldnn::prop_kind::forward_inference, algo,
|
||||
*src_md_, *primitive_dst_md_,
|
||||
strides_mkl, kernel_mkl,
|
||||
padding_left_mkl, padding_right_mkl,
|
||||
mkldnn::padding_kind::zero));
|
||||
|
||||
fwd_primitive_desc_.reset(new mkldnn::pooling_forward::primitive_desc(
|
||||
*fwd_desc_, cpu_engine));
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
// Sub-graph's first node. Read input from input buffer
|
||||
src_mem_.reset(new mkldnn::memory(
|
||||
fwd_primitive_desc_.get()->src_primitive_desc(), nullptr));
|
||||
} else {
|
||||
// Sub-graph's inner node. set input to parent's output
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
|
||||
primitive_src_format_ = static_cast<mkldnn::memory::format>(
|
||||
fwd_primitive_desc_.get()->src_primitive_desc().desc().data.format);
|
||||
|
||||
primitive_dst_format_ = static_cast<mkldnn::memory::format>(
|
||||
fwd_primitive_desc_.get()->dst_primitive_desc().desc().data.format);
|
||||
|
||||
src_size_ = fwd_primitive_desc_.get()->src_primitive_desc().get_size();
|
||||
dst_size_ = fwd_primitive_desc_.get()->dst_primitive_desc().get_size();
|
||||
|
||||
// reorder source memory for best performance (AVX512);
|
||||
if (primitive_src_format_ != src_format_) {
|
||||
mkldnn::memory::dims src_dims_mkl(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
|
||||
auto src_md = mkldnn::memory::desc(src_dims_mkl, MklDnnType<T>(), src_format_);
|
||||
auto pd = mkldnn::memory::primitive_desc(src_md, cpu_engine);
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty())
|
||||
src_mem_from_.reset(new mkldnn::memory(pd, nullptr));
|
||||
else
|
||||
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
|
||||
src_mem_.reset(new mkldnn::memory(fwd_primitive_desc_->src_primitive_desc(), nullptr));
|
||||
net.push_back(mkldnn::reorder(*src_mem_from_, *src_mem_));
|
||||
} else {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
src_mem_.reset(new mkldnn::memory(fwd_primitive_desc_->src_primitive_desc(), nullptr));
|
||||
} else {
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// last node of sub-graph. need to allocate memory for output_tensor
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
// reorder neded. Use primitive output as input to reorder and
|
||||
// allocate buffer for reorder output, final output of this subgraph
|
||||
primitive_dst_mem_.reset(
|
||||
new mkldnn::memory(fwd_primitive_desc_.get()->dst_primitive_desc()));
|
||||
} else {
|
||||
// Last node but re-order not needed. Allocate buffer to output of this node
|
||||
primitive_dst_mem_.reset(
|
||||
new mkldnn::memory(fwd_primitive_desc_.get()->dst_primitive_desc(), nullptr));
|
||||
}
|
||||
} else {
|
||||
// Intermediate node. Use mkldnn kernel internal memory for output and
|
||||
// use this as input to next node.
|
||||
primitive_dst_mem_.reset(
|
||||
new mkldnn::memory(fwd_primitive_desc_.get()->dst_primitive_desc()));
|
||||
}
|
||||
pool_fwd_.reset(
|
||||
new mkldnn::pooling_forward(*fwd_primitive_desc_, *src_mem_, *primitive_dst_mem_));
|
||||
|
||||
net.push_back(*pool_fwd_);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// one of the end nodes. Allocate output buffer memory and
|
||||
// reorder is necessary
|
||||
mkldnn::memory::data_type t = MklDnnType<T>();
|
||||
InitDstReorderOutput(cpu_engine, t, net);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Bind(const OrtCustomOpApi* api, OrtKernelContext* context) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
if (x_shape_.NumDimensions() <= 3) {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
// abort as MKLDNN cannot execute this. but
|
||||
// ORT try to delete output_tensor buffer data. allocate memory so that it can delete
|
||||
// fix for test_averagepool_1d_default node test
|
||||
//auto xshape = input_tensors[input_index].shape;
|
||||
//auto xdim = input_tensors[input_index].ndim;
|
||||
//AllocateOutputTensor(output_tensors, mklnode_ptr_->output_index, xshape, xdim, input_tensors[0].dtype);
|
||||
}
|
||||
std::cout << "MKLDNN cannot compute shape with dim less than three." << std::endl;
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Please call default CPU kernel.");
|
||||
}
|
||||
|
||||
if (primitive_src_format_ != src_format_) {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
src_mem_from_->set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
} else {
|
||||
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
|
||||
auto src_size = fwd_primitive_desc_.get()->src_primitive_desc().get_size();
|
||||
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
|
||||
src_mem_->set_data_handle(src_reorder_buffer_.get());
|
||||
} else {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
src_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
} else {
|
||||
src_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// Last node of sub-graph. Allocate memory for output_buffer data
|
||||
// Reorder if needed
|
||||
auto& y_dims = primitive_dst_shape_.GetDims();
|
||||
// Allocate memory for output bufffer
|
||||
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
|
||||
T* dst_data = ort.GetTensorMutableData<T>(output);
|
||||
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(dst_data);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
void ReadAttributes(const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") override {
|
||||
global_pooling_ = (op_name_ == "GlobalAveragePool" || op_name_ == "GlobalMaxPool" || op_name_ == "GlobalLpPool");
|
||||
global_pooling_ = (op_name_ == "GlobalAveragePool" || op_name_ == "GlobalMaxPool" || op_name_ == "GlobalLpPool");
|
||||
|
||||
if (!global_pooling_) {
|
||||
bool attr_read = false;
|
||||
auto attr = attributes.find(attributes_prefix + "kernel_shape");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
GetIntsAttr(proto, kernel_shape_);
|
||||
attr_read = true;
|
||||
}
|
||||
ORT_ENFORCE(attr_read, "No kernel shape is set.");
|
||||
|
||||
std::string auto_padding;
|
||||
attr = attributes.find(attributes_prefix + "auto_pad");
|
||||
if (attr != attributes.end() &&
|
||||
attr->second.type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_STRING) {
|
||||
auto_padding = attr->second.s();
|
||||
}
|
||||
auto_pad_ = StringToAutoPadType(auto_padding);
|
||||
|
||||
attr_read = false;
|
||||
attr = attributes.find(attributes_prefix + "pads");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
if (GetIntsAttr(proto, pads_) == Status::OK())
|
||||
attr_read = true;
|
||||
}
|
||||
if (!attr_read) {
|
||||
pads_.resize(kernel_shape_.size() * 2, 0);
|
||||
}
|
||||
|
||||
attr_read = false;
|
||||
attr = attributes.find(attributes_prefix + "strides");
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
if (GetIntsAttr(proto, strides_) == Status::OK())
|
||||
attr_read = true;
|
||||
}
|
||||
if (!attr_read || strides_.empty()) {
|
||||
strides_.resize(kernel_shape_.size(), 1);
|
||||
}
|
||||
|
||||
attr = attributes.find(attributes_prefix + "count_include_pad");
|
||||
int64_t temp = 0;
|
||||
if (attr != attributes.end()) {
|
||||
ONNX_NAMESPACE::AttributeProto proto = attr->second;
|
||||
GetIntAttr(proto, temp);
|
||||
}
|
||||
count_include_pad_ = (temp != 0);
|
||||
|
||||
storage_order_ = 0;
|
||||
for (size_t dim = 0; dim < kernel_shape_.size(); ++dim) {
|
||||
ORT_ENFORCE(kernel_shape_[dim] > 0);
|
||||
ORT_ENFORCE(pads_[dim] < kernel_shape_[dim] && pads_[dim + kernel_shape_.size()] < kernel_shape_[dim],
|
||||
"Pad should be smaller than kernel.");
|
||||
}
|
||||
|
||||
ORT_ENFORCE(strides_.size() == kernel_shape_.size());
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
size_t src_size_;
|
||||
size_t dst_size_;
|
||||
|
||||
std::shared_ptr<mkldnn::memory> src_mem_;
|
||||
|
||||
std::unique_ptr<mkldnn::pooling_forward::desc> fwd_desc_;
|
||||
std::unique_ptr<mkldnn::memory::desc> src_md_;
|
||||
std::unique_ptr<mkldnn::pooling_forward::primitive_desc> fwd_primitive_desc_;
|
||||
std::unique_ptr<mkldnn::primitive> pool_fwd_;
|
||||
|
||||
std::shared_ptr<mkldnn::memory> src_mem_from_;
|
||||
std::unique_ptr<mkldnn::memory> src_mem_to_;
|
||||
|
||||
std::unique_ptr<mkldnn::memory> dst_mem_from_;
|
||||
std::unique_ptr<mkldnn::memory> dst_mem_to_;
|
||||
|
||||
private:
|
||||
mkldnn::memory::format GetAVXFormat(const mkldnn::memory::dims& src_dims_mkl) {
|
||||
bool is_2D = src_dims_mkl.size() == 4 ? true : false;
|
||||
mkldnn::memory::format fmt = mkldnn::memory::format::any;
|
||||
if (CPUIDInfo::GetCPUIDInfo().HasAVX512f()) {
|
||||
fmt = is_2D ? mkldnn::memory::format::nChw16c : mkldnn::memory::format::nCdhw16c;
|
||||
} else if (CPUIDInfo::GetCPUIDInfo().HasAVX2() && (src_dims_mkl[1] % 8 == 0)) {
|
||||
fmt = is_2D ? mkldnn::memory::format::nChw8c : mkldnn::memory::format::ncdhw;
|
||||
} else {
|
||||
fmt = is_2D ? mkldnn::memory::format::nchw : mkldnn::memory::format::ncdhw;
|
||||
}
|
||||
return fmt;
|
||||
}
|
||||
|
||||
std::vector<int64_t> SetOutputSize(const TensorShape& input_shape,
|
||||
int64_t output_channel,
|
||||
std::vector<int64_t>* pads) const {
|
||||
ORT_ENFORCE(input_shape.Size() > 0);
|
||||
std::vector<int64_t> output_dims;
|
||||
int64_t N = input_shape[0];
|
||||
InferOutputSize(input_shape.GetDims(), &output_dims, pads);
|
||||
|
||||
output_dims.insert(output_dims.begin(), {N, output_channel});
|
||||
|
||||
return output_dims;
|
||||
}
|
||||
|
||||
inline void InferOutputSize(const std::vector<int64_t>& input_dims,
|
||||
std::vector<int64_t>* output_dims,
|
||||
std::vector<int64_t>* pads) const {
|
||||
ORT_ENFORCE(input_dims.size() >= 2);
|
||||
if (global_pooling_) {
|
||||
output_dims->assign(input_dims.size() - 2, 1);
|
||||
} else {
|
||||
for (size_t dim = 0; dim < input_dims.size() - 2; ++dim) {
|
||||
int64_t dim_size = 0;
|
||||
ComputeSizeAndPad(static_cast<int>(input_dims[dim + 2]),
|
||||
strides_[dim],
|
||||
kernel_shape_[dim],
|
||||
&pads->at(dim),
|
||||
&pads->at(input_dims.size() + dim - 2),
|
||||
&dim_size);
|
||||
output_dims->push_back(dim_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
inline void ComputeSizeAndPad(const int64_t in_size,
|
||||
const int64_t stride,
|
||||
const int64_t kernel,
|
||||
int64_t* pad_head,
|
||||
int64_t* pad_tail,
|
||||
int64_t* out_size) const {
|
||||
if (auto_pad_ != AutoPadType::NOTSET) {
|
||||
switch (auto_pad_) {
|
||||
case AutoPadType::VALID:
|
||||
*pad_head = 0;
|
||||
*pad_tail = 0;
|
||||
*out_size = (in_size - kernel) / stride + 1;
|
||||
break;
|
||||
case AutoPadType::SAME_LOWER: {
|
||||
int64_t legacy_target_size = (in_size + stride - 1) / stride;
|
||||
int64_t pad_needed = (legacy_target_size - 1) * stride + kernel - in_size;
|
||||
*pad_head = (pad_needed + 1) / 2;
|
||||
*pad_tail = pad_needed - *pad_head;
|
||||
*out_size = (in_size + pad_needed - kernel) / stride + 1;
|
||||
break;
|
||||
}
|
||||
case AutoPadType::SAME_UPPER: {
|
||||
int64_t legacy_target_size = (in_size + stride - 1) / stride;
|
||||
int64_t pad_needed = (legacy_target_size - 1) * stride + kernel - in_size;
|
||||
*pad_head = pad_needed / 2;
|
||||
*pad_tail = pad_needed - *pad_head;
|
||||
*out_size = (in_size + pad_needed - kernel) / stride + 1;
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
ORT_THROW("Unsupported AutoPad Type.");
|
||||
}
|
||||
}
|
||||
} else {
|
||||
*out_size = static_cast<int64_t>(
|
||||
static_cast<float>(in_size + *pad_head + *pad_tail - kernel) / stride + 1);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
IAllocatorUniquePtr<void> src_reorder_buffer_;
|
||||
IAllocatorUniquePtr<void> dst_reorder_buffer_;
|
||||
|
||||
private:
|
||||
std::string op_name_;
|
||||
bool global_pooling_{};
|
||||
bool count_include_pad_{};
|
||||
int64_t storage_order_{0}; // MaxPool_8 only. 0 is row major, and 1 is column major. Default is 0.
|
||||
std::vector<int64_t> kernel_shape_;
|
||||
std::vector<int64_t> pads_;
|
||||
std::vector<int64_t> strides_;
|
||||
AutoPadType auto_pad_;
|
||||
|
||||
TensorShape x_shape_;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
172
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_sum.h
Normal file
172
onnxruntime/core/providers/mkldnn/subgraph/mkldnn_sum.h
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
// Copyright(C) 2019 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
|
||||
#pragma once
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/providers/mkldnn/mkldnn_fwd.h"
|
||||
#include "core/providers/mkldnn/mkldnn_common.h"
|
||||
#include "core/providers/mkldnn/subgraph/mkldnn_kernel.h"
|
||||
#include "core/util/math.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
template <typename T>
|
||||
class MklDnnSum : public MklDnnKernel {
|
||||
public:
|
||||
explicit MklDnnSum(const MklDnnNode& node,
|
||||
MKLDNNExecutionProvider* provider,
|
||||
const NodeAttributes& attributes,
|
||||
const std::string attributes_prefix = "") : MklDnnKernel(node, provider) {
|
||||
ReadAttributes(attributes, attributes_prefix);
|
||||
}
|
||||
|
||||
Status CreatePrimitives(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context,
|
||||
mkldnn::engine& cpu_engine,
|
||||
std::vector<mkldnn::primitive>& net,
|
||||
mkldnn::memory::format& source_format) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int num_inputs = mklnode_ptr_->num_inputs;
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
std::vector<float> coeff;
|
||||
TensorShape x_shape;
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
|
||||
ort_source_format_ = GetSourceFormat(static_cast<int>(xdim));
|
||||
source_format = ort_source_format_;
|
||||
src_format_ = ort_source_format_;
|
||||
x_shape = TensorShape(xshape, xdim);
|
||||
} else {
|
||||
x_shape = parents_[0].get()->primitive_dst_shape_;
|
||||
src_format_ = parents_[0].get()->primitive_dst_format_;
|
||||
}
|
||||
primitive_dst_shape_ = TensorShape(x_shape);
|
||||
|
||||
mkldnn::memory::dims dst_dims_mkl(
|
||||
primitive_dst_shape_.GetDims().begin(), primitive_dst_shape_.GetDims().end());
|
||||
|
||||
for (int i = 0; i < num_inputs; i++) {
|
||||
TensorShape x_shape1;
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
|
||||
auto tensor_shape = ort.GetTensorShape(tensor_info);
|
||||
ort.ReleaseTensorTypeAndShapeInfo(tensor_info);
|
||||
auto xshape = tensor_shape.data();
|
||||
auto xdim = tensor_shape.size();
|
||||
mkldnn::memory::dims dims(xdim);
|
||||
|
||||
ort_source_format_ = GetSourceFormat(static_cast<int>(xdim));
|
||||
x_shape1 = TensorShape(xshape, xdim);
|
||||
mkldnn::memory::dims src_dims_mkl(
|
||||
x_shape1.GetDims().begin(), x_shape1.GetDims().end());
|
||||
|
||||
src_md_.reset(new mkldnn::memory::desc(
|
||||
{src_dims_mkl}, MklDnnType<T>(), src_format_));
|
||||
|
||||
auto mpd = mkldnn::memory::primitive_desc(*src_md_, cpu_engine);
|
||||
auto src_memory = mkldnn::memory(mpd, nullptr);
|
||||
srcs_pd_.push_back(mpd);
|
||||
srcs_memory_.push_back(src_memory);
|
||||
coeff.push_back(1.0);
|
||||
} else {
|
||||
src_md_.reset(
|
||||
new mkldnn::memory::desc(parents_[i].get()->primitive_dst_mem_.get()->get_primitive_desc().desc()));
|
||||
auto mpd = mkldnn::memory::primitive_desc(*src_md_, cpu_engine);
|
||||
auto src_memory = *parents_[i].get()->primitive_dst_mem_; //mkldnn::memory(mpd);
|
||||
srcs_pd_.push_back(mpd);
|
||||
srcs_memory_.push_back(src_memory);
|
||||
coeff.push_back(1.0);
|
||||
ort_source_format_ = source_format;
|
||||
}
|
||||
}
|
||||
|
||||
primitive_dst_md_.reset(new mkldnn::memory::desc(
|
||||
{dst_dims_mkl}, MklDnnType<T>(), mkldnn::memory::format::any));
|
||||
sum_pd_.reset(new mkldnn::sum::primitive_desc(
|
||||
*primitive_dst_md_, coeff, srcs_pd_));
|
||||
primitive_dst_format_ = static_cast<mkldnn::memory::format>(sum_pd_->dst_primitive_desc().desc().data.format);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// last node of sub-graph. need to allocate memory for output_tensor
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
// reorder neded. Use primitive output as input to reorder and
|
||||
// allocate buffer for reorder output, final output of this subgraph
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(sum_pd_->dst_primitive_desc()));
|
||||
} else {
|
||||
// Last node but re-order not needed. Allocate buffer to output of this node
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(sum_pd_->dst_primitive_desc(), nullptr));
|
||||
}
|
||||
} else {
|
||||
// Intermediate node. Use mkldnn kernel internal memory for output and
|
||||
// use this as input to next node.
|
||||
primitive_dst_mem_.reset(new mkldnn::memory(sum_pd_->dst_primitive_desc()));
|
||||
}
|
||||
primitive_dst_format_ = static_cast<mkldnn::memory::format>(sum_pd_->dst_primitive_desc().desc().data.format);
|
||||
|
||||
std::vector<mkldnn::primitive::at> inputs;
|
||||
for (int i = 0; i < num_inputs; i++) {
|
||||
inputs.push_back(srcs_memory_[i]);
|
||||
}
|
||||
auto c = mkldnn::sum(*sum_pd_, inputs, *primitive_dst_mem_);
|
||||
net.push_back(c);
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// one of the end nodes. Allocate output buffer memory and
|
||||
// reorder is necessary
|
||||
mkldnn::memory::data_type t = MklDnnType<T>();
|
||||
InitDstReorderOutput(cpu_engine, t, net);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Bind(const OrtCustomOpApi* api, OrtKernelContext* context) override {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
|
||||
int num_inputs = mklnode_ptr_->num_inputs;
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
for (int i = 0; i < num_inputs; i++) {
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index + i);
|
||||
const T* src_data = const_cast<T*>(ort.GetTensorData<T>(input_tensor));
|
||||
srcs_memory_[i].set_data_handle(static_cast<void*>(const_cast<T*>(src_data)));
|
||||
}
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
// Last node. Allocate output buffer memory and reorder if needed
|
||||
const auto& y_dims = primitive_dst_shape_.GetDims();
|
||||
// Allocate memory for output bufffer
|
||||
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
|
||||
T* dst_data = ort.GetTensorMutableData<T>(output);
|
||||
|
||||
if (primitive_dst_format_ != ort_source_format_) {
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(dst_data);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<mkldnn::memory::desc> src_md_;
|
||||
std::vector<mkldnn::memory> srcs_memory_;
|
||||
|
||||
std::vector<mkldnn::memory::primitive_desc> srcs_pd_;
|
||||
std::unique_ptr<mkldnn::memory::primitive_desc> src_mpd_;
|
||||
std::unique_ptr<mkldnn::memory::primitive_desc> dst_pd_;
|
||||
std::unique_ptr<mkldnn::sum::primitive_desc> sum_pd_;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
73
onnxruntime/core/providers/mkldnn/subgraph/subgraph.h
Normal file
73
onnxruntime/core/providers/mkldnn/subgraph/subgraph.h
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
// Copyright(C) 2019 Intel Corporation
|
||||
// Licensed under the MIT License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include "core/framework/op_node_proto_helper.h"
|
||||
#include "core/graph/graph.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace mkl_dnn {
|
||||
|
||||
struct MklDnnNode {
|
||||
std::string name;
|
||||
int node_index = -1;
|
||||
int input_start_index = -1; // start index in inputs()
|
||||
int num_inputs = 0; // and how many inputs
|
||||
int output_index = -1; // index in output()
|
||||
std::string weight_name;
|
||||
std::string output_name;
|
||||
std::vector<size_t> parent_nodes; // index to parents in vector mklnodes
|
||||
|
||||
std::string ToString() const {
|
||||
std::string key;
|
||||
key.reserve(128);
|
||||
key.append(name);
|
||||
key.append("-");
|
||||
key.append(std::to_string(input_start_index));
|
||||
key.append("-");
|
||||
key.append(std::to_string(num_inputs));
|
||||
key.append("-");
|
||||
key.append(std::to_string(output_index));
|
||||
key.append("-");
|
||||
key.append(output_name);
|
||||
key.append("-");
|
||||
for (auto& out : parent_nodes)
|
||||
key.append(std::to_string(out) + ",");
|
||||
key.append(";");
|
||||
return key;
|
||||
}
|
||||
};
|
||||
|
||||
struct Subgraph {
|
||||
struct SubgraphVariables {
|
||||
std::vector<std::string> inputs;
|
||||
std::vector<std::string> outputs;
|
||||
std::vector<std::string> outputs_as_input_other_node;
|
||||
std::vector<onnxruntime::NodeIndex> subgraph_node_indexes;
|
||||
int subgraph_index = 0;
|
||||
|
||||
SubgraphVariables() {
|
||||
subgraph_index = 0;
|
||||
}
|
||||
void Reset() {
|
||||
subgraph_node_indexes.clear();
|
||||
inputs.clear();
|
||||
outputs.clear();
|
||||
outputs_as_input_other_node.clear();
|
||||
}
|
||||
};
|
||||
|
||||
Subgraph(const std::string& name) {
|
||||
graph_name = name;
|
||||
}
|
||||
|
||||
std::string graph_name;
|
||||
std::string subgraph_id;
|
||||
std::vector<MklDnnNode> mkldnn_nodes;
|
||||
};
|
||||
} // namespace mkl_dnn
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -579,6 +579,39 @@ def run_onnx_tests(build_dir, configs, onnx_test_data_dir, provider, enable_para
|
|||
run_subprocess([exe,'-x'] + cmd, cwd=cwd)
|
||||
|
||||
|
||||
# mkldnn temporary function for running onnx tests and model tests separately.
|
||||
def mkldnn_run_onnx_tests(build_dir, configs, onnx_test_data_dir):
|
||||
for config in configs:
|
||||
cwd = get_config_build_dir(build_dir, config)
|
||||
if is_windows():
|
||||
exe = os.path.join(cwd, config, 'onnx_test_runner')
|
||||
model_dir = os.path.join(cwd, "models")
|
||||
else:
|
||||
exe = os.path.join(cwd, 'onnx_test_runner')
|
||||
model_dir = os.path.join(build_dir, "models")
|
||||
cmd_base = ['-e', 'mkldnn', '-c', '1', '-j', '1']
|
||||
if os.path.exists(onnx_test_data_dir):
|
||||
onnxdata_cmd = cmd_base + [onnx_test_data_dir]
|
||||
# /data/onnx
|
||||
run_subprocess([exe] + onnxdata_cmd, cwd=cwd)
|
||||
run_subprocess([exe,'-x'] + onnxdata_cmd, cwd=cwd)
|
||||
|
||||
# models/opset7, models/opset8, models/opset9
|
||||
if config != 'Debug' and os.path.exists(model_dir):
|
||||
opset7_model_dir = os.path.join(model_dir, 'opset7')
|
||||
opset7_cmd = cmd_base + [opset7_model_dir]
|
||||
opset8_model_dir = os.path.join(model_dir, 'opset8')
|
||||
opset8_cmd = cmd_base + [opset8_model_dir]
|
||||
opset9_model_dir = os.path.join(model_dir, 'opset9')
|
||||
opset9_cmd = cmd_base + [opset9_model_dir]
|
||||
run_subprocess([exe] + opset7_cmd, cwd=cwd)
|
||||
run_subprocess([exe, '-x'] + opset7_cmd, cwd=cwd)
|
||||
run_subprocess([exe] + opset8_cmd, cwd=cwd)
|
||||
run_subprocess([exe, '-x'] + opset8_cmd, cwd=cwd)
|
||||
run_subprocess([exe] + opset9_cmd, cwd=cwd)
|
||||
run_subprocess([exe, '-x'] + opset9_cmd, cwd=cwd)
|
||||
|
||||
|
||||
def split_server_binary_and_symbol(build_dir, configs):
|
||||
if is_windows():
|
||||
# TODO: Windows support
|
||||
|
|
@ -839,7 +872,7 @@ def main():
|
|||
run_onnx_tests(build_dir, configs, onnx_test_data_dir, None, True, 0)
|
||||
|
||||
if args.use_mkldnn:
|
||||
run_onnx_tests(build_dir, configs, onnx_test_data_dir, 'mkldnn', True, 1)
|
||||
mkldnn_run_onnx_tests(build_dir, configs, onnx_test_data_dir)
|
||||
|
||||
if args.build_server:
|
||||
split_server_binary_and_symbol(build_dir, configs)
|
||||
|
|
|
|||
Loading…
Reference in a new issue