diff --git a/docs/execution_providers/MKL-DNN-ExecutionProvider.md b/docs/execution_providers/MKL-DNN-ExecutionProvider.md new file mode 100644 index 0000000000..716ef27ed5 --- /dev/null +++ b/docs/execution_providers/MKL-DNN-ExecutionProvider.md @@ -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).. diff --git a/docs/execution_providers/MKL-DNN-Subgraphs.md b/docs/execution_providers/MKL-DNN-Subgraphs.md new file mode 100644 index 0000000000..5a82a1726f --- /dev/null +++ b/docs/execution_providers/MKL-DNN-Subgraphs.md @@ -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. + +

+ + +## Subgraph Classes +Primitive like MklDnnConv, MklDnnPool, etc are derived from MklDnnKernel base class. + +The following UML diagram captures Subgraph classes. + +

+ + +## Subgraph Execution + +MklDnnExecutionProvicer::Compute() function creates MklDnnFuncKernel and call it’s Compute Function. + + +MklDnnFuncKernel::Compute function creates SubgraphPrimitve pool and add the object to a map. + +SubgraphPrimitve constructor calls the following member functions +``` +SubgraphPrimitve::CreatePrimitives() + for (auto& mklnode : mklnodes) { + if (mklnode.name == "Conv") { + kernel.reset(new MklDnnConv()); + kernels.push_back(kernel); + } else if (mklnode.name == "BatchNormalization-Relu") { + kernel.reset(new MklDnnBatchNorm()); + context_.kernels.push_back(kernel); + } else if (mklnode.name == "MaxPool") { + kernel.reset(new MklDnnPool()); + context_.kernels.push_back(kernel); + } + . + . + . +``` +In CreatePrimitives method, we iterate MklDnnNodes and creates MklDnnKernel objects and add MKL-DNN primitive to a vector. It also reads attributes. This is done only once, at first iteration. + +``` +SubgraphPrimitve::Compute() + for (auto& kernel : kernels) { + kernel->Bind(input_tensors, output_tensors); + } + stream->submit(net); +``` + +In SubgraphPrimitve::Compute() method, we iterate thru MklDnn Kernels and bind input data. Then we submit the vector of Primitives to MKL-DNN stream. + diff --git a/docs/execution_providers/images/mkl-dnn_node.png b/docs/execution_providers/images/mkl-dnn_node.png new file mode 100644 index 0000000000..d2863f8938 Binary files /dev/null and b/docs/execution_providers/images/mkl-dnn_node.png differ diff --git a/docs/execution_providers/images/mkl-dnn_subgraph.png b/docs/execution_providers/images/mkl-dnn_subgraph.png new file mode 100644 index 0000000000..a4b5587ae7 Binary files /dev/null and b/docs/execution_providers/images/mkl-dnn_subgraph.png differ diff --git a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc index aae5959264..307cdb7454 100644 --- a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc +++ b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc @@ -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 MKLDNNExecutionProvider::GetKernelRegistry() con static std::shared_ptr kernel_registry = onnxruntime::mkl_dnn::GetMklDnnKernelRegistry(); return kernel_registry; } + +bool MKLDNNExecutionProvider::UseSubgraph(const onnxruntime::GraphViewer& graph_viewer, + const std::vector& kernel_registries, + std::vector>& 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& subgraph_ptr, + mkl_dnn::Subgraph::SubgraphVariables& sub_var, + bool fused, + std::map& 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(node->InputDefs().size()); + mkldnn_node.input_start_index = static_cast(sub_var.inputs.size()) - 1; + mkldnn_node.node_index = static_cast(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 att(key, att_it->second); + subgraph_attributes[key] = att_it->second; + } + } +} + +std::vector> MKLDNNExecutionProvider::GetCapability( + const onnxruntime::GraphViewer& graph_viewer, + const std::vector& kernel_registries) const { + ORT_UNUSED_PARAMETER(kernel_registries); + std::vector> 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 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 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& subgraph_ptr, + mkl_dnn::Subgraph::SubgraphVariables& sub_var, + std::vector>& 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 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(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 sub_graph = std::make_unique(); + sub_graph->nodes = sub_var.subgraph_node_indexes; + sub_graph->SetMetaDef(meta_def); + result.push_back(std::make_unique(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& fused_nodes, + std::vector& 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(context, attributes, this); + *state = p; + return 0; + }; + + compute_info.release_state_func = [](FunctionState state) { + if (state) + delete static_cast*>(state); + }; + + compute_info.compute_func = [](FunctionState state, const OrtCustomOpApi* api, OrtKernelContext* context) { + onnxruntime::mkl_dnn::MkldnnFuncKernel* custom_op = reinterpret_cast*>(state); + return custom_op->Compute(api, context); + }; + + node_compute_funcs.push_back(compute_info); + } + return Status::OK(); +} } // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h index 0e7ba75eae..7e844adce5 100644 --- a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h +++ b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h @@ -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& filter_dst_mem) { + const std::shared_ptr& 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> + GetCapability(const onnxruntime::GraphViewer& graph, + const std::vector& /*kernel_registries*/) const override; + + common::Status Compile(const std::vector& fused_nodes, + std::vector& 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> 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& kernel_registries, + std::vector>& 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& subgraph_ptr, + mkl_dnn::Subgraph::SubgraphVariables& sub_var, + bool fused, + std::map& 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& subgraph_ptr, + mkl_dnn::Subgraph::SubgraphVariables& sub_var, + std::vector>& result) const; + + public: + const std::shared_ptr GetMklDnnSubgraph(const std::string& subgraph_id) { + return mkl_subgraphs_[subgraph_id]; + } + + private: + // supported MklDnn Operators + std::set mkldnn_ops_ = {"Conv", "BatchNormalization", "Relu", "Sum", + "AveragePool", "GlobalMaxPool", "GlobalAveragePool", "MaxPool", "LRN"}; + + mutable std::unordered_map> mkl_subgraphs_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_activations.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_activations.h new file mode 100644 index 0000000000..406542fd15 --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_activations.h @@ -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 +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& 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(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(), 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( + relu_fwd_pd_.get()->src_primitive_desc().desc().data.format); + primitive_dst_format_ = static_cast( + 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(); + 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(ort.GetTensorData(input_tensor)); + src_mem_->set_data_handle(static_cast(const_cast(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(primitive_dst_shape_.GetDims().size())); + T* dst_data = ort.GetTensorMutableData(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 src_mem_; + + std::unique_ptr fwd_desc_; + std::unique_ptr relu_fwd_pd_; + std::unique_ptr relu_fwd_; + + std::unique_ptr src_md_; +}; +} // namespace mkl_dnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_batchnorm.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_batchnorm.h new file mode 100644 index 0000000000..17728b927f --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_batchnorm.h @@ -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& 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 +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& 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(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(), 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(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(), mkldnn::memory::format::nc)); + mean_md_.reset(new mkldnn::memory::desc( + {mean_dims_mkl}, MklDnnType(), mkldnn::memory::format::x)); + var_md_.reset(new mkldnn::memory::desc( + {var_dims_mkl}, MklDnnType(), mkldnn::memory::format::x)); + primitive_dst_md_.reset(new mkldnn::memory::desc( + {dst_dims_mkl}, MklDnnType(), 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( + batchnorm_fwd_pd_.get()->dst_primitive_desc().desc().data.format); + primitive_src_format_ = static_cast( + 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(); + 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(ort.GetTensorData(input_tensor)); + src_mem_->set_data_handle(static_cast(const_cast(src_data))); + } + + const OrtValue* scale_input_tensor = ort.KernelContext_GetInput(context, input_index + 1); + const T* scale_data = reinterpret_cast(ort.GetTensorData(scale_input_tensor)); + const OrtValue* b_input_tensor = ort.KernelContext_GetInput(context, input_index + 2); + const T* b_data = reinterpret_cast(ort.GetTensorData(b_input_tensor)); + const OrtValue* mean_input_tensor = ort.KernelContext_GetInput(context, input_index + 3); + const T* mean_data = reinterpret_cast(ort.GetTensorData(mean_input_tensor)); + const OrtValue* var_input_tensor = ort.KernelContext_GetInput(context, input_index + 4); + const T* var_data = reinterpret_cast(ort.GetTensorData(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(const_cast(mean_data))); + var_mem_->set_data_handle(static_cast(const_cast(var_data))); + + T* scale_shift_buf = static_cast(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(primitive_dst_shape_.GetDims().size())); + T* dst_data = ort.GetTensorMutableData(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 src_mem_; + std::unique_ptr scale_shift_mem_; + std::unique_ptr mean_mem_; + std::unique_ptr var_mem_; + std::unique_ptr dst_mem_; + + std::unique_ptr src_md_; + std::unique_ptr scale_shift_md_; + std::unique_ptr mean_md_; + std::unique_ptr var_md_; + std::unique_ptr dst_md_; + + std::unique_ptr batchnorm_fwd_; + std::unique_ptr batchnorm_fwd_pd_; + + protected: + float epsilon_ = 1e-5f; +}; +} // namespace mkl_dnn +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_conv.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_conv.h new file mode 100644 index 0000000000..ceafa69e5f --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_conv.h @@ -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 +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(static_cast(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(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 +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& 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(group_); + + TensorShape x_shape; + // std::unique_ptr 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(), 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(), mkldnn::memory::format::any)); + } + + primitive_created_ = ValidateInputShape(x_shape, w_shape); + if (!primitive_created_.IsOK()) + return primitive_created_; + + std::vector 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 pads(pads_); + if (pads.empty()) { + pads.resize(kernel_rank * 2, 0); + } + std::vector dilations(dilations_); + if (dilations.empty()) { + dilations.resize(kernel_rank, 1); + } + std::vector 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 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(), 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(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(), mkldnn::memory::format::any)); + if (!bias_dims_mkl.empty()) + bias_md_.reset(new mkldnn::memory::desc( + {bias_dims_mkl}, MklDnnType(), 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( + conv_fwd_pd_.get()->src_primitive_desc().desc().data.format); + + mkldnn_filter_format_ = static_cast( + conv_fwd_pd_.get()->weights_primitive_desc().desc().data.format); + + primitive_dst_format_ = static_cast( + 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(), 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(); + 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(ort.GetTensorData(input_tensor)); + + const int group_mkl = static_cast(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(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 lock(provider_->GetMutex()); + std::shared_ptr 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(), filter_format_), cpu_engine); + mkldnn::memory src = mkldnn::memory(pd, (void*)filter_data); + IAllocatorUniquePtr filter_reorder_buffer = + IAllocator::MakeUniquePtr(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(params); + provider_->SaveAllocatedMemory(std::move(filter_reorder_buffer)); + filter_data = static_cast(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(ort.GetTensorData(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(ort.GetTensorData(binput_tensor)); + } + std::shared_ptr 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(filter_dst_mem->get_data_handle()); + + filter_mem_->set_data_handle(static_cast(const_cast(filter_data))); + if (bias_data != nullptr) { + bias_mem_->set_data_handle(static_cast(const_cast(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(ort.GetTensorData(input_tensor)); + src_mem_from_->set_data_handle(static_cast(const_cast(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(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(ort.GetTensorData(input_tensor)); + src_mem_->set_data_handle(static_cast(const_cast(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(primitive_dst_shape_.GetDims().size())); + T* dst_data = ort.GetTensorMutableData(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 src_mem_from_; + std::unique_ptr src_mem_to_; + + size_t src_size_; + size_t filter_size_; + size_t dst_size_; + + std::shared_ptr src_mem_; + std::unique_ptr filter_mem_; + std::unique_ptr bias_mem_; + + std::unique_ptr fwd_desc_; + + std::unique_ptr src_md_; + std::unique_ptr filter_md_; + std::unique_ptr bias_md_; + + std::unique_ptr conv_fwd_pd_; + std::unique_ptr conv_fwd_; + + private: + IAllocatorUniquePtr src_reorder_buffer_; + IAllocatorUniquePtr dst_reorder_buffer_; + + private: + Status ComputeKernelShape(const TensorShape& weight_shape, std::vector& 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(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 + Status InferOutputShape(const TensorShape& input_shape, + const std::vector& kernel_shape, + const std::vector& strides, + const std::vector& dilations, + std::vector* pads, + std::vector* output_shape) const { + size_t rank = gsl::narrow_cast(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( + 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 kernel_shape_; // must use ComputeKernelShape(...), instead of kernel_shape_ + AutoPadType auto_pad_; + int64_t group_; + bool kernel_shape_specified_; + std::vector strides_; + std::vector pads_; + std::vector dilations_; + std::string activation_; + float alpha_; +}; +} // namespace mkl_dnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.cc b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.cc new file mode 100644 index 0000000000..15a7510505 --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.cc @@ -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 +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> kernel; + kernel.reset(new MklDnnConv(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> kernel; + kernel.reset(new MklDnnConv(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> kernel; + kernel.reset(new MklDnnRelu(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> kernel; + kernel.reset(new MklDnnBatchNorm(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> kernel; + kernel.reset(new MklDnnBatchNorm(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> kernel; + kernel.reset(new MklDnnPool(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> kernel; + kernel.reset(new MklDnnPool(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> kernel; + kernel.reset(new MklDnnPool(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> kernel; + kernel.reset(new MklDnnPool(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> kernel; + kernel.reset(new MklDnnSum(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> kernel; + kernel.reset(new MklDnnLrn(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 stream; + std::vector net; + std::vector> 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 +class SubgraphPrimitivePool : public PrimitivePool { + public: + static SubgraphPrimitive* 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* primitive = dynamic_cast*>( + SubgraphPrimitivePool::GetInstance().GetPrimitive(params.subgraph_key + dims_str)); + + if (primitive == nullptr) { + auto subgraph_primitive = std::make_unique>(api, context, params); + primitive = subgraph_primitive.get(); + SubgraphPrimitivePool::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 +Status MkldnnFuncKernel::Compute(const OrtCustomOpApi* api, OrtKernelContext* context) const { + Status status; + try { + SubgraphPrimitive* primitive = SubgraphPrimitivePool::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; + +} // namespace mkl_dnn +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.h new file mode 100644 index 0000000000..59105a3b8c --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_func_kernel.h @@ -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; + std::string subgraph_id; + std::string subgraph_key; + + SubgraphParams() {} +}; +} // namespace + +template +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 \ No newline at end of file diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.cc b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.cc new file mode 100644 index 0000000000..19007e99ef --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.cc @@ -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& 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 \ No newline at end of file diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.h new file mode 100644 index 0000000000..0a4bc68eea --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_kernel.h @@ -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(node); + provider_ = provider; + alloc_ = provider_->GetAllocator(0, OrtMemTypeDefault); + } + virtual ~MklDnnKernel(){}; + + virtual Status CreatePrimitives(const OrtCustomOpApi* api, + OrtKernelContext* context, + mkldnn::engine& cpu_engine, + std::vector& 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& 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& 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> parents_; + bool fuse_relu_ = false; + bool fuse_sum_ = false; + std::shared_ptr primitive_dst_mem_; + std::unique_ptr 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 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 reorder_dst_mem_to_; + + protected: + Status primitive_created_; + AllocatorPtr alloc_; + MKLDNNExecutionProvider* provider_; +}; + +} // namespace mkl_dnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_lrn.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_lrn.h new file mode 100644 index 0000000000..c27914e79a --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_lrn.h @@ -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 +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& 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(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(), 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( + fwd_primitive_desc_.get()->src_primitive_desc().desc().data.format); + primitive_dst_format_ = static_cast( + 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(); + 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(ort.GetTensorData(input_tensor)); + src_mem_->set_data_handle(static_cast(const_cast(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(primitive_dst_shape_.GetDims().size())); + T* dst_data = ort.GetTensorMutableData(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 src_mem_; + + std::unique_ptr fwd_desc_; + std::unique_ptr fwd_primitive_desc_; + std::unique_ptr lrn_fwd_; + + std::unique_ptr src_md_; +}; +} // namespace mkl_dnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_pool.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_pool.h new file mode 100644 index 0000000000..e4ac5099fc --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_pool.h @@ -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 +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& 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(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(), 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(), 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 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(), 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( + fwd_primitive_desc_.get()->src_primitive_desc().desc().data.format); + + primitive_dst_format_ = static_cast( + 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(), 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(); + 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(ort.GetTensorData(input_tensor)); + src_mem_from_->set_data_handle(static_cast(const_cast(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(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(ort.GetTensorData(input_tensor)); + src_mem_->set_data_handle(static_cast(const_cast(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(primitive_dst_shape_.GetDims().size())); + T* dst_data = ort.GetTensorMutableData(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 src_mem_; + + std::unique_ptr fwd_desc_; + std::unique_ptr src_md_; + std::unique_ptr fwd_primitive_desc_; + std::unique_ptr pool_fwd_; + + std::shared_ptr src_mem_from_; + std::unique_ptr src_mem_to_; + + std::unique_ptr dst_mem_from_; + std::unique_ptr 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 SetOutputSize(const TensorShape& input_shape, + int64_t output_channel, + std::vector* pads) const { + ORT_ENFORCE(input_shape.Size() > 0); + std::vector 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& input_dims, + std::vector* output_dims, + std::vector* 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(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( + static_cast(in_size + *pad_head + *pad_tail - kernel) / stride + 1); + } + } + + private: + IAllocatorUniquePtr src_reorder_buffer_; + IAllocatorUniquePtr 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 kernel_shape_; + std::vector pads_; + std::vector strides_; + AutoPadType auto_pad_; + + TensorShape x_shape_; +}; +} // namespace mkl_dnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_sum.h b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_sum.h new file mode 100644 index 0000000000..a45af46acc --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/mkldnn_sum.h @@ -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 +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& 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 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(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(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(), 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(), mkldnn::memory::format::any)); + sum_pd_.reset(new mkldnn::sum::primitive_desc( + *primitive_dst_md_, coeff, srcs_pd_)); + primitive_dst_format_ = static_cast(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(sum_pd_->dst_primitive_desc().desc().data.format); + + std::vector 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(); + 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(ort.GetTensorData(input_tensor)); + srcs_memory_[i].set_data_handle(static_cast(const_cast(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(primitive_dst_shape_.GetDims().size())); + T* dst_data = ort.GetTensorMutableData(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 src_md_; + std::vector srcs_memory_; + + std::vector srcs_pd_; + std::unique_ptr src_mpd_; + std::unique_ptr dst_pd_; + std::unique_ptr sum_pd_; +}; +} // namespace mkl_dnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/subgraph/subgraph.h b/onnxruntime/core/providers/mkldnn/subgraph/subgraph.h new file mode 100644 index 0000000000..6e3f967dab --- /dev/null +++ b/onnxruntime/core/providers/mkldnn/subgraph/subgraph.h @@ -0,0 +1,73 @@ +// Copyright(C) 2019 Intel Corporation +// Licensed under the MIT License + +#pragma once + +#include +#include +#include +#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 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 inputs; + std::vector outputs; + std::vector outputs_as_input_other_node; + std::vector 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 mkldnn_nodes; +}; +} // namespace mkl_dnn +} // namespace onnxruntime \ No newline at end of file diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index cab2006497..6a23c292a2 100755 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -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)