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:
Sreekanth Yalachigere 2019-06-10 20:18:56 -07:00 committed by jywu-msft
parent 3c3186c761
commit 24d6b0f5c4
18 changed files with 3239 additions and 2 deletions

View 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)..

View 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 its 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.

Binary file not shown.

After

Width:  |  Height:  |  Size: 50 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 42 KiB

View file

@ -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

View file

@ -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

View 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

View 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

View 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

View 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

View 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

View 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

View 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

View 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

View 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

View 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

View 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

View file

@ -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)