add dnnl gpu engine (#5788)

This commit is contained in:
jeyblu 2020-11-12 20:17:54 -08:00 committed by GitHub
parent 0ea998134a
commit 435b904f0e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 870 additions and 283 deletions

View file

@ -308,9 +308,23 @@ These instructions are for JetPack SDK 4.4.
See more information on DNNL and MKL-ML [here](./docs/execution_providers/DNNL-ExecutionProvider.md).
#### Build Instructions
The DNNL execution provider can be built for Intel CPU or GPU. To build for Intel GPU, install [Intel SDK for OpenCL Applications](https://software.intel.com/content/www/us/en/develop/tools/opencl-sdk.html). Install the latest GPU driver - [Windows graphics driver](https://downloadcenter.intel.com/product/80939/Graphics), [Linux graphics compute runtime and OpenCL driver](https://github.com/intel/compute-runtime/releases).
##### Windows
`.\build.bat --use_dnnl`
##### Linux
`./build.sh --use_dnnl`
To build for Intel GPU, replace dnnl_opencl_root with the path of the Intel SDK for OpenCL Applications.
##### Windows
`.\build.bat --use_dnnl --dnnl_gpu_runtime ocl --dnnl_opencl_root "c:\program files (x86)\intelswtools\sw_dev_tools\opencl\sdk"`
##### Linux
DNNL: `./build.sh --use_dnnl`
`./build.sh --use_dnnl --dnnl_gpu_runtime ocl --dnnl_opencl_root "/opt/intel/sw_dev_tools/opencl-sdk"`
---

View file

@ -1,8 +1,8 @@
include (ExternalProject)
set(DNNL_URL https://github.com/intel/mkl-dnn.git)
# If DNNL_TAG is updated, check if platform.cmake.patch need to be updated.
set(DNNL_TAG v1.1.1)
set(DNNL_URL https://github.com/oneapi-src/onednn)
# If DNNL_TAG is updated, check if MKLML_VERSION and platform.cmake.patch need to be updated.
set(DNNL_TAG v1.7)
if(WIN32)
set(MKLML_OS_VERSION_STR "win")
@ -20,6 +20,14 @@ else()
endif()
endif()
if (onnxruntime_USE_DNNL AND onnxruntime_DNNL_GPU_RUNTIME STREQUAL "ocl" AND onnxruntime_DNNL_OPENCL_ROOT STREQUAL "")
message(FATAL_ERROR "onnxruntime_DNNL_OPENCL_ROOT required for onnxruntime_DNNL_GPU_RUNTIME")
elseif(onnxruntime_USE_DNNL AND onnxruntime_DNNL_GPU_RUNTIME STREQUAL "ocl")
file(TO_CMAKE_PATH ${onnxruntime_DNNL_OPENCL_ROOT} onnxruntime_DNNL_OPENCL_ROOT)
set(DNNL_OCL_INCLUDE_DIR ${onnxruntime_DNNL_OPENCL_ROOT}/include)
set(DNNL_GPU_CMAKE_ARGS "-DDNNL_GPU_RUNTIME=OCL " "-DOPENCLROOT=${onnxruntime_DNNL_OPENCL_ROOT}")
endif()
if (onnxruntime_USE_DNNL)
set(DNNL_SOURCE ${CMAKE_CURRENT_BINARY_DIR}/dnnl/src/dnnl/src)
set(DNNL_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/dnnl/install)
@ -46,7 +54,7 @@ if (onnxruntime_USE_DNNL)
GIT_TAG ${DNNL_TAG}
# PATCH_COMMAND ${MKLDNN_PATCH_DISCARD_COMMAND} COMMAND ${DNNL_PATCH_COMMAND}
SOURCE_DIR ${DNNL_SOURCE}
CMAKE_ARGS -DDNNL_BUILD_TESTS=OFF -DDNNL_BUILD_EXAMPLES=OFF -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} -DCMAKE_INSTALL_PREFIX=${DNNL_INSTALL}
CMAKE_ARGS -DDNNL_BUILD_TESTS=OFF -DDNNL_BUILD_EXAMPLES=OFF -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} -DCMAKE_INSTALL_PREFIX=${DNNL_INSTALL} ${DNNL_GPU_CMAKE_ARGS}
)
link_directories(${DNNL_LIB_DIR})
endif()

View file

@ -364,12 +364,15 @@ if (onnxruntime_USE_DNNL)
target_link_directories(onnxruntime_providers_dnnl PRIVATE ${DNNL_LIB_DIR})
onnxruntime_add_include_to_target(onnxruntime_providers_dnnl onnxruntime_common onnx) # onnx needed for stl_backports.h
add_dependencies(onnxruntime_providers_dnnl onnxruntime_providers_shared project_dnnl ${onnxruntime_EXTERNAL_DEPENDENCIES})
target_include_directories(onnxruntime_providers_dnnl PRIVATE ${ONNXRUNTIME_ROOT} ${CMAKE_CURRENT_BINARY_DIR} ${eigen_INCLUDE_DIRS} ${DNNL_INCLUDE_DIR})
target_include_directories(onnxruntime_providers_dnnl PRIVATE ${ONNXRUNTIME_ROOT} ${CMAKE_CURRENT_BINARY_DIR} ${eigen_INCLUDE_DIRS} ${DNNL_INCLUDE_DIR} ${DNNL_OCL_INCLUDE_DIR})
# ${CMAKE_CURRENT_BINARY_DIR} is so that #include "onnxruntime_config.h" inside tensor_shape.h is found
target_link_libraries(onnxruntime_providers_dnnl PRIVATE dnnl onnxruntime_providers_shared)
install(DIRECTORY ${PROJECT_SOURCE_DIR}/../include/onnxruntime/core/providers/dnnl DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/onnxruntime/core/providers)
set_target_properties(onnxruntime_providers_dnnl PROPERTIES FOLDER "ONNXRuntime")
set_target_properties(onnxruntime_providers_dnnl PROPERTIES LINKER_LANGUAGE CXX)
if (onnxruntime_DNNL_GPU_RUNTIME STREQUAL "ocl")
target_compile_definitions(onnxruntime_providers_dnnl PRIVATE USE_DNNL_GPU_OCL=1)
endif()
if(APPLE)
set_property(TARGET onnxruntime_providers_dnnl APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker -exported_symbols_list ${ONNXRUNTIME_ROOT}/core/providers/dnnl/exported_symbols.lst")

View file

@ -5,6 +5,7 @@
#include "core/providers/shared_library/provider_api.h"
#include "gsl/gsl-lite.hpp"
#include "dnnl.hpp"
#include "dnnl_debug.h"
#include <unordered_map>
#include <list>
@ -20,10 +21,50 @@ dnnl::memory::data_type DnnnType<float>() {
return dnnl::memory::data_type::f32;
}
static dnnl::engine& GetEngine() {
static dnnl::engine cpu_engine = dnnl::engine(dnnl::engine::kind::cpu, 0);
return cpu_engine;
}
class DnnlEngineInstance {
private:
static DnnlEngineInstance* instance;
std::unordered_map<dnnl::engine::kind, dnnl::engine> dnnl_engine_map;
DnnlEngineInstance() {
dnnl::engine engine;
try {
engine = dnnl::engine(dnnl::engine::kind::cpu, 0);
if (engine) {
dnnl_engine_map.insert(std::make_pair(dnnl::engine::kind::cpu, engine));
}
} catch (const std::exception& e) {
LOGS_DEFAULT(ERROR) << e.what() << std::endl;
throw;
}
try {
engine = dnnl::engine(dnnl::engine::kind::gpu, 0);
if (engine) {
dnnl_engine_map.insert(std::make_pair(dnnl::engine::kind::gpu, engine));
}
} catch (const std::exception& e) {
LOGS_DEFAULT(INFO) << e.what() << std::endl;
}
}
public:
DnnlEngineInstance(DnnlEngineInstance& other) = delete;
void operator=(const DnnlEngineInstance&) = delete;
static DnnlEngineInstance* getInstance() {
if (!instance)
instance = new DnnlEngineInstance();
return instance;
}
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& getEngineMap() {
return dnnl_engine_map;
}
const dnnl::engine& getEngine(dnnl::engine::kind kind) {
std::unordered_map<dnnl::engine::kind, dnnl::engine>::iterator iter = dnnl_engine_map.find(kind);
return iter->second;
}
};
static void AddDimsToKey(std::string& key, const dnnl::memory::dims& dims) {
key.append(1, '#');

View file

@ -22,9 +22,25 @@ class DnnlRelu : public DnnlKernel {
void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) {
dnnl::engine cpu_engine;
dnnl::engine engine_to_use;
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine.find(dnnl::engine::kind::cpu);
if (iter != dnnl_engine.end()) {
cpu_engine = (dnnl::engine)iter->second;
engine_to_use = cpu_engine;
}
gpu_available_ = false;
dnnl::engine gpu_engine;
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine.end()) {
gpu_engine = (dnnl::engine)(iter->second);
gpu_available_ = true;
engine_to_use = gpu_engine;
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
}
Ort::CustomOpApi ort{*api};
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
@ -59,10 +75,20 @@ class DnnlRelu : public DnnlKernel {
dnnl::memory::desc({src_dims}, DnnnType<T>(), ort_source_format_));
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory({{src_dims}, DnnnType<T>(), ort_source_format_}, cpu_engine, nullptr));
if (gpu_available_) {
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(*src_md_, gpu_engine);
net.push_back(mkldnn::reorder(*src_mem_, *src_mem_gpu_));
net_args.push_back({{MKLDNN_ARG_SRC, *src_mem_},
{MKLDNN_ARG_DST, *src_mem_gpu_}});
}
} else {
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
dnnl::memory::desc(parents_[0].get()->primitive_dst_desc_));
src_mem_ = parents_[0].get()->primitive_dst_mem_;
if (!gpu_available_) {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
} else { // gpu_available_
src_mem_gpu_ = parents_[0].get()->primitive_dst_mem_;
}
x_shape = parents_[0].get()->primitive_dst_shape_;
ort_source_format_ = parents_[0].get()->ort_source_format_;
ort_source_desc_ = parents_[0].get()->ort_source_desc_;
@ -76,40 +102,49 @@ class DnnlRelu : public DnnlKernel {
fwd_desc_ = onnxruntime::make_unique<dnnl::eltwise_forward::desc>(
dnnl::eltwise_forward::desc(dnnl::prop_kind::forward_inference, algo, *src_md_, 0));
relu_fwd_pd_ = onnxruntime::make_unique<dnnl::eltwise_forward::primitive_desc>(
dnnl::eltwise_forward::primitive_desc(*fwd_desc_, cpu_engine));
dnnl::eltwise_forward::primitive_desc(*fwd_desc_, engine_to_use));
primitive_src_desc_ = relu_fwd_pd_.get()->src_desc();
primitive_dst_desc_ = relu_fwd_pd_.get()->dst_desc();
if (mklnode_ptr_->output_index >= 0) {
// last node of sub-graph. need to allocate memory for output_tensor
if (primitive_dst_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), cpu_engine));
if (!gpu_available_) {
if (mklnode_ptr_->output_index >= 0) {
// last node of sub-graph. need to allocate memory for output_tensor
if (primitive_dst_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), cpu_engine));
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), cpu_engine, nullptr));
// Intermediate node. Use dnnl kernel internal memory for output and
// use this as input to next node.
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), cpu_engine));
}
} else {
// Intermediate node. Use dnnl kernel internal memory for output and
// use this as input to next node.
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), cpu_engine));
} else { // gpu_available_
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), gpu_engine));
}
relu_fwd_ = onnxruntime::make_unique<dnnl::eltwise_forward>(
dnnl::eltwise_forward(*relu_fwd_pd_));
net.push_back(*relu_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
if (!gpu_available_) {
net.push_back(*relu_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
} else { // gpu_available_
net.push_back(*relu_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
}
if (mklnode_ptr_->output_index >= 0) {
// one of the end nodes. Allocate output buffer memory and
// reorder is necessary
dnnl::memory::data_type t = DnnnType<T>();
InitDstReorderOutput(cpu_engine, t, net, net_args);
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
}
}
@ -133,10 +168,14 @@ class DnnlRelu : public DnnlKernel {
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
T* dst_data = ort.GetTensorMutableData<T>(output);
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (primitive_dst_desc_ != ort_source_desc_) {
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
} else { // gpu_available_
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
}
@ -145,12 +184,15 @@ class DnnlRelu : public DnnlKernel {
private:
std::shared_ptr<dnnl::memory> src_mem_;
std::shared_ptr<dnnl::memory> src_mem_gpu_;
std::unique_ptr<dnnl::eltwise_forward::desc> fwd_desc_;
std::unique_ptr<dnnl::eltwise_forward::primitive_desc> relu_fwd_pd_;
std::unique_ptr<dnnl::primitive> relu_fwd_;
std::unique_ptr<dnnl::memory::desc> src_md_;
bool gpu_available_;
};
} // namespace ort_dnnl
} // namespace onnxruntime

View file

@ -97,9 +97,25 @@ class DnnlBatchNorm : public DnnlKernel {
void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
dnnl::engine cpu_engine;
dnnl::engine engine_to_use;
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine.find(dnnl::engine::kind::cpu);
if (iter != dnnl_engine.end()) {
cpu_engine = (dnnl::engine)iter->second;
engine_to_use = cpu_engine;
}
gpu_available_ = false;
dnnl::engine gpu_engine;
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine.end()) {
gpu_engine = (dnnl::engine)(iter->second);
gpu_available_ = true;
engine_to_use = gpu_engine;
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
}
Ort::CustomOpApi ort{*api};
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
@ -213,6 +229,15 @@ class DnnlBatchNorm : public DnnlKernel {
var_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(*var_md_, cpu_engine, nullptr));
if (gpu_available_) {
scale_shift_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory({*scale_shift_md_, gpu_engine}));
mean_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(*mean_md_, gpu_engine));
var_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(*var_md_, gpu_engine));
}
batchnorm_fwd_ = onnxruntime::make_unique<dnnl::batch_normalization_forward::desc>(
dnnl::batch_normalization_forward::desc(
dnnl::prop_kind::forward_inference, *src_md_, epsilon_,
@ -231,11 +256,11 @@ class DnnlBatchNorm : public DnnlKernel {
attr.set_post_ops(ops);
batchnorm_fwd_pd_ = onnxruntime::make_unique<dnnl::batch_normalization_forward::primitive_desc>(
dnnl::batch_normalization_forward::primitive_desc(*batchnorm_fwd_, attr, cpu_engine));
dnnl::batch_normalization_forward::primitive_desc(*batchnorm_fwd_, attr, engine_to_use));
} else {
batchnorm_fwd_pd_ = onnxruntime::make_unique<dnnl::batch_normalization_forward::primitive_desc>(
dnnl::batch_normalization_forward::primitive_desc(
*batchnorm_fwd_, cpu_engine));
*batchnorm_fwd_, engine_to_use));
}
// out format of this kernel
@ -244,42 +269,81 @@ class DnnlBatchNorm : public DnnlKernel {
primitive_src_desc_ = static_cast<dnnl::memory::desc>(
batchnorm_fwd_pd_.get()->dst_desc());
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_.get()->src_desc(), cpu_engine, nullptr));
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
if (!gpu_available_) {
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_.get()->src_desc(), cpu_engine, nullptr));
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
}
} else { // gpu_available_
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_.get()->src_desc(), cpu_engine, nullptr));
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_.get()->src_desc(), gpu_engine));
net.push_back(dnnl::reorder(*src_mem_, *src_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *src_mem_gpu_}});
} else {
src_mem_gpu_ = parents_[0].get()->primitive_dst_mem_;
}
}
if (mklnode_ptr_->output_index >= 0) {
// Use Dnnl's internal output buffer
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (mklnode_ptr_->output_index >= 0) {
// Use Dnnl's internal output buffer
if (primitive_dst_desc_ != ort_source_desc_) {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_->dst_desc(), cpu_engine));
} else {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_->dst_desc(), cpu_engine, nullptr));
}
} else {
// last node of sub-graph. need to allocate memory for output_tensor
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_->dst_desc(), cpu_engine));
} else {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_->dst_desc(), cpu_engine, nullptr));
}
} else {
// last node of sub-graph. need to allocate memory for output_tensor
} else { // gpu_available_
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(batchnorm_fwd_pd_->dst_desc(), cpu_engine));
dnnl::memory(batchnorm_fwd_pd_->dst_desc(), gpu_engine));
}
if (gpu_available_) {
net.push_back(dnnl::reorder(*mean_mem_, *mean_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *mean_mem_},
{DNNL_ARG_DST, *mean_mem_gpu_}});
net.push_back(dnnl::reorder(*var_mem_, *var_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *var_mem_},
{DNNL_ARG_DST, *var_mem_gpu_}});
net.push_back(dnnl::reorder(*scale_shift_mem_, *scale_shift_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *scale_shift_mem_},
{DNNL_ARG_DST, *scale_shift_mem_gpu_}});
}
auto bn = dnnl::batch_normalization_forward(
*batchnorm_fwd_pd_);
net.push_back(bn);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_MEAN, *mean_mem_},
{DNNL_ARG_VARIANCE, *var_mem_},
{DNNL_ARG_SCALE_SHIFT, *scale_shift_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
if (!gpu_available_) {
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_MEAN, *mean_mem_},
{DNNL_ARG_VARIANCE, *var_mem_},
{DNNL_ARG_SCALE_SHIFT, *scale_shift_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
} else { // gpu_available_
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
{DNNL_ARG_MEAN, *mean_mem_gpu_},
{DNNL_ARG_VARIANCE, *var_mem_gpu_},
{DNNL_ARG_SCALE_SHIFT, *scale_shift_mem_gpu_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
}
// 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
dnnl::memory::data_type t = DnnnType<T>();
InitDstReorderOutput(cpu_engine, t, net, net_args);
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
}
}
@ -334,10 +398,14 @@ class DnnlBatchNorm : public DnnlKernel {
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
T* dst_data = ort.GetTensorMutableData<T>(output);
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (primitive_dst_desc_ != ort_source_desc_) {
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
} else { // gpu_available_
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
}
return Status::OK();
@ -350,6 +418,11 @@ class DnnlBatchNorm : public DnnlKernel {
std::unique_ptr<dnnl::memory> var_mem_;
std::unique_ptr<dnnl::memory> dst_mem_;
std::shared_ptr<dnnl::memory> src_mem_gpu_;
std::unique_ptr<dnnl::memory> scale_shift_mem_gpu_;
std::unique_ptr<dnnl::memory> mean_mem_gpu_;
std::unique_ptr<dnnl::memory> var_mem_gpu_;
std::unique_ptr<dnnl::memory::desc> src_md_;
std::unique_ptr<dnnl::memory::desc> scale_shift_md_;
std::unique_ptr<dnnl::memory::desc> mean_md_;
@ -359,6 +432,8 @@ class DnnlBatchNorm : public DnnlKernel {
std::unique_ptr<dnnl::batch_normalization_forward::desc> batchnorm_fwd_;
std::unique_ptr<dnnl::batch_normalization_forward::primitive_desc> batchnorm_fwd_pd_;
bool gpu_available_;
protected:
float epsilon_ = 1e-5f;
};

View file

@ -69,9 +69,27 @@ class DnnlConv : public DnnlKernel {
void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
dnnl::engine cpu_engine;
dnnl::engine engine_to_use;
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine.find(dnnl::engine::kind::cpu);
if (iter != dnnl_engine.end()) {
dnnl_engine_cpu_ = (dnnl::engine)iter->second;
cpu_engine = (dnnl::engine)iter->second;
engine_to_use = cpu_engine;
}
gpu_available_ = false;
dnnl::engine gpu_engine;
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine.end()) {
dnnl_engine_gpu_ = (dnnl::engine)iter->second;
gpu_engine = (dnnl::engine)(iter->second);
gpu_available_ = true;
engine_to_use = gpu_engine;
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
}
Ort::CustomOpApi ort{*api};
stream_ = onnxruntime::make_unique<dnnl::stream>(dnnl::stream(cpu_engine));
@ -267,10 +285,10 @@ class DnnlConv : public DnnlKernel {
attr.set_post_ops(ops);
conv_fwd_pd_ = onnxruntime::make_unique<dnnl::convolution_forward::primitive_desc>(
dnnl::convolution_forward::primitive_desc(*fwd_desc_, attr, cpu_engine));
dnnl::convolution_forward::primitive_desc(*fwd_desc_, attr, engine_to_use));
} else {
conv_fwd_pd_ = onnxruntime::make_unique<dnnl::convolution_forward::primitive_desc>(
dnnl::convolution_forward::primitive_desc(*fwd_desc_, cpu_engine));
dnnl::convolution_forward::primitive_desc(*fwd_desc_, engine_to_use));
}
primitive_src_desc_ = static_cast<dnnl::memory::desc>(
@ -288,44 +306,84 @@ class DnnlConv : public DnnlKernel {
filter_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->weights_desc(), cpu_engine, nullptr));
if (gpu_available_) {
filter_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->weights_desc(), gpu_engine, nullptr));
}
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape.GetDims().begin(), x_shape.GetDims().end());
auto pd = dnnl::memory::desc({{src_dims}, DnnnType<T>(), ort_source_format_});
if (!gpu_available_) {
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape.GetDims().begin(), x_shape.GetDims().end());
auto pd = dnnl::memory::desc({{src_dims}, DnnnType<T>(), ort_source_format_});
if (mklnode_ptr_->parent_nodes.empty())
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
else
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
if (mklnode_ptr_->parent_nodes.empty())
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
else
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_}});
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_}});
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
}
}
} else { // gpu_available_
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape.GetDims().begin(), x_shape.GetDims().end());
auto pd = dnnl::memory::desc({{src_dims}, DnnnType<T>(), ort_source_format_});
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
} else {
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
}
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), gpu_engine));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_gpu_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_gpu_}});
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), gpu_engine));
net.push_back(dnnl::reorder(*src_mem_, *src_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *src_mem_gpu_}});
} else {
src_mem_gpu_ = parents_[0].get()->primitive_dst_mem_;
}
}
}
if (mklnode_ptr_->output_index >= 0) {
// Use Dnnl's internal output buffer
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (mklnode_ptr_->output_index >= 0) {
// Use Dnnl's internal output buffer
if (primitive_dst_desc_ != ort_source_desc_) {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine));
} else {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// last node of sub-graph. need to allocate memory for output_tensor
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine));
} else {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// last node of sub-graph. need to allocate memory for output_tensor
} else { // gpu_available_
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine));
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), gpu_engine));
}
if (!bias_dims_mkl.empty()) {
@ -333,28 +391,49 @@ class DnnlConv : public DnnlKernel {
dnnl::memory(conv_fwd_pd_.get()->bias_desc(), cpu_engine, nullptr));
conv_fwd_ = onnxruntime::make_unique<dnnl::convolution_forward>(
dnnl::convolution_forward(*conv_fwd_pd_));
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_WEIGHTS, *filter_mem_},
{DNNL_ARG_BIAS, *bias_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
if (!gpu_available_) {
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_WEIGHTS, *filter_mem_},
{DNNL_ARG_BIAS, *bias_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
} else { // gpu_available_
bias_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->bias_desc(), gpu_engine));
net.push_back(dnnl::reorder(*bias_mem_, *bias_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *bias_mem_},
{DNNL_ARG_DST, *bias_mem_gpu_}});
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
{DNNL_ARG_WEIGHTS, *filter_mem_gpu_},
{DNNL_ARG_BIAS, *bias_mem_gpu_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
}
} else {
conv_fwd_ = onnxruntime::make_unique<dnnl::convolution_forward>(
dnnl::convolution_forward(*conv_fwd_pd_));
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_WEIGHTS, *filter_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
if (!gpu_available_) {
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_WEIGHTS, *filter_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
} else { // gpu_available_
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
{DNNL_ARG_WEIGHTS, *filter_mem_gpu_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
}
}
if (mklnode_ptr_->output_index >= 0) {
// one of the end nodes. Allocate output buffer memory and
// reorder is necessary
dnnl::memory::data_type t = DnnnType<T>();
InitDstReorderOutput(cpu_engine, t, net, net_args);
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
}
}
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, dnnl::engine& cpu_engine) override {
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, const dnnl::engine& cpu_engine) override {
Ort::CustomOpApi ort{*api};
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
@ -387,14 +466,23 @@ class DnnlConv : public DnnlKernel {
if (filter_dst_mem == nullptr) {
dnnl::memory src = dnnl::memory({{filter_dims_mkl}, DnnnType<T>(), filter_format_}, cpu_engine, (void*)filter_data);
IAllocatorUniquePtr<void> filter_reorder_buffer = IAllocator::MakeUniquePtr<void>(alloc_, filter_size_);
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get()));
if (!gpu_available_) {
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get()));
dnnl::reorder(src, *filter_dst_mem)
.execute(cpu_engine, src, *filter_dst_mem);
dnnl::reorder(src, *filter_dst_mem)
.execute(cpu_engine, src, *filter_dst_mem);
provider_->SaveAllocatedMemory(std::move(filter_reorder_buffer));
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
} else { // gpu_available_
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), dnnl_engine_gpu_));
dnnl::reorder(src, *filter_dst_mem)
.execute(dnnl_engine_gpu_, src, *filter_dst_mem);
}
provider_->SaveAllocatedMemory(std::move(filter_reorder_buffer));
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
provider_->SetWeightsMemoryBuffer(mklnode_ptr_->weight_name, filter_dst_mem);
}
}
@ -416,12 +504,19 @@ class DnnlConv : public DnnlKernel {
}
std::shared_ptr<dnnl::memory> filter_dst_mem = provider_->GetWeightsMemoryBuffer(mklnode_ptr_->weight_name);
if (filter_dst_mem == nullptr) {
ReorderWeights(api, context, GetEngine());
ReorderWeights(api, context, dnnl_engine_cpu_);
filter_dst_mem = provider_->GetWeightsMemoryBuffer(mklnode_ptr_->weight_name);
}
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
if (!gpu_available_) {
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
filter_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(filter_data)));
} else { // gpu_available_
#ifdef USE_DNNL_GPU_OCL
std::lock_guard<OrtMutex> lock(provider_->GetMutex());
filter_mem_gpu_->set_ocl_mem_object(filter_dst_mem->get_ocl_mem_object());
#endif
}
filter_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(filter_data)));
if (bias_data != nullptr) {
bias_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(bias_data)));
}
@ -435,9 +530,11 @@ class DnnlConv : public DnnlKernel {
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
}
auto src_size = conv_fwd_pd_.get()->src_desc().get_size();
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
if (!gpu_available_) {
auto src_size = conv_fwd_pd_.get()->src_desc().get_size();
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
}
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
@ -454,10 +551,14 @@ class DnnlConv : public DnnlKernel {
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
T* dst_data = ort.GetTensorMutableData<T>(output);
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (primitive_dst_desc_ != ort_source_desc_) {
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
} else { // gpu_available_
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
}
return Status::OK();
@ -537,6 +638,10 @@ class DnnlConv : public DnnlKernel {
std::unique_ptr<dnnl::memory> filter_mem_;
std::unique_ptr<dnnl::memory> bias_mem_;
std::shared_ptr<dnnl::memory> src_mem_gpu_;
std::unique_ptr<dnnl::memory> filter_mem_gpu_;
std::unique_ptr<dnnl::memory> bias_mem_gpu_;
std::unique_ptr<dnnl::convolution_forward::desc> fwd_desc_;
std::unique_ptr<dnnl::memory::desc> src_md_;
@ -546,6 +651,11 @@ class DnnlConv : public DnnlKernel {
std::unique_ptr<dnnl::convolution_forward::primitive_desc> conv_fwd_pd_;
std::unique_ptr<dnnl::primitive> conv_fwd_;
dnnl::engine dnnl_engine_cpu_;
dnnl::engine dnnl_engine_gpu_;
bool gpu_available_;
private:
IAllocatorUniquePtr<void> src_reorder_buffer_;
IAllocatorUniquePtr<void> dst_reorder_buffer_;

View file

@ -23,9 +23,27 @@ class DnnlConvBatchNorm : public DnnlKernel {
void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
dnnl::engine cpu_engine;
dnnl::engine engine_to_use;
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine.find(dnnl::engine::kind::cpu);
if (iter != dnnl_engine.end()) {
dnnl_engine_cpu_ = (dnnl::engine)iter->second;
cpu_engine = (dnnl::engine)iter->second;
engine_to_use = cpu_engine;
}
gpu_available_ = false;
dnnl::engine gpu_engine;
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine.end()) {
dnnl_engine_gpu_ = (dnnl::engine)iter->second;
gpu_engine = (dnnl::engine)(iter->second);
gpu_available_ = true;
engine_to_use = gpu_engine;
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
}
Ort::CustomOpApi ort{*api};
stream_ = onnxruntime::make_unique<dnnl::stream>(dnnl::stream(cpu_engine));
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
@ -220,10 +238,10 @@ class DnnlConvBatchNorm : public DnnlKernel {
attr.set_post_ops(ops);
conv_fwd_pd_ = onnxruntime::make_unique<dnnl::convolution_forward::primitive_desc>(
dnnl::convolution_forward::primitive_desc(*fwd_desc_, attr, cpu_engine));
dnnl::convolution_forward::primitive_desc(*fwd_desc_, attr, engine_to_use));
} else {
conv_fwd_pd_ = onnxruntime::make_unique<dnnl::convolution_forward::primitive_desc>(
dnnl::convolution_forward::primitive_desc(*fwd_desc_, cpu_engine));
dnnl::convolution_forward::primitive_desc(*fwd_desc_, engine_to_use));
}
primitive_src_desc_ = static_cast<dnnl::memory::desc>(
@ -241,61 +259,114 @@ class DnnlConvBatchNorm : public DnnlKernel {
filter_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->weights_desc(), cpu_engine, nullptr));
if (gpu_available_) {
filter_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->weights_desc(), gpu_engine));
}
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape.GetDims().begin(), x_shape.GetDims().end());
auto pd = dnnl::memory::desc({{src_dims}, DnnnType<T>(), ort_source_format_});
if (!gpu_available_) {
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape.GetDims().begin(), x_shape.GetDims().end());
auto pd = dnnl::memory::desc({{src_dims}, DnnnType<T>(), ort_source_format_});
if (mklnode_ptr_->parent_nodes.empty())
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
else
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
if (mklnode_ptr_->parent_nodes.empty())
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
else
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_}});
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_}});
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
}
}
} else { // gpu_available_
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape.GetDims().begin(), x_shape.GetDims().end());
auto pd = dnnl::memory::desc({{src_dims}, DnnnType<T>(), ort_source_format_});
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
} else {
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
}
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), gpu_engine));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_gpu_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_gpu_}});
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), cpu_engine, nullptr));
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->src_desc(), gpu_engine));
net.push_back(dnnl::reorder(*src_mem_, *src_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *src_mem_gpu_}});
} else {
src_mem_gpu_ = parents_[0].get()->primitive_dst_mem_;
}
}
}
if (mklnode_ptr_->output_index >= 0) {
// Use Dnnl's internal output buffer
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (mklnode_ptr_->output_index >= 0) {
// Use Dnnl's internal output buffer
if (primitive_dst_desc_ != ort_source_desc_) {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine));
} else {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// last node of sub-graph. need to allocate memory for output_tensor
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine));
} else {
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// last node of sub-graph. need to allocate memory for output_tensor
} else { // gpu_available_
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), cpu_engine));
dnnl::memory(conv_fwd_pd_.get()->dst_desc(), gpu_engine));
}
bias_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->bias_desc(), cpu_engine, nullptr));
conv_fwd_ = onnxruntime::make_unique<dnnl::convolution_forward>(
dnnl::convolution_forward(*conv_fwd_pd_));
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_WEIGHTS, *filter_mem_},
{DNNL_ARG_BIAS, *bias_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
if (!gpu_available_) {
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_WEIGHTS, *filter_mem_},
{DNNL_ARG_BIAS, *bias_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
} else { // gpu_available_
bias_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_.get()->bias_desc(), gpu_engine));
net.push_back(dnnl::reorder(*bias_mem_, *bias_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *bias_mem_},
{DNNL_ARG_DST, *bias_mem_gpu_}});
net.push_back(*conv_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
{DNNL_ARG_WEIGHTS, *filter_mem_gpu_},
{DNNL_ARG_BIAS, *bias_mem_gpu_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
}
if (mklnode_ptr_->output_index >= 0) {
// one of the end nodes. Allocate output buffer memory and
// reorder is necessary
dnnl::memory::data_type t = DnnnType<T>();
InitDstReorderOutput(cpu_engine, t, net, net_args);
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
}
}
@ -355,7 +426,7 @@ class DnnlConvBatchNorm : public DnnlKernel {
}
}
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, dnnl::engine& cpu_engine) override {
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, const dnnl::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);
@ -396,13 +467,22 @@ class DnnlConvBatchNorm : public DnnlKernel {
if (filter_dst_mem == nullptr) {
dnnl::memory src = dnnl::memory({{filter_dims_mkl}, DnnnType<T>(), filter_format_}, cpu_engine, (void*)weights_scaled_by_axis.data());
IAllocatorUniquePtr<void> filter_reorder_buffer = IAllocator::MakeUniquePtr<void>(alloc_, filter_size_);
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get()));
if (!gpu_available_) {
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get()));
dnnl::reorder(src, *filter_dst_mem)
.execute(cpu_engine, src, *filter_dst_mem);
dnnl::reorder(src, *filter_dst_mem)
.execute(cpu_engine, src, *filter_dst_mem);
provider_->SaveAllocatedMemory(std::move(filter_reorder_buffer));
} else { // gpu_available_
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), dnnl_engine_gpu_));
dnnl::reorder(src, *filter_dst_mem)
.execute(dnnl_engine_gpu_, src, *filter_dst_mem);
}
provider_->SaveAllocatedMemory(std::move(filter_reorder_buffer));
provider_->SetWeightsMemoryBuffer(mklnode_ptr_->weight_name, filter_dst_mem);
}
@ -458,11 +538,17 @@ class DnnlConvBatchNorm : public DnnlKernel {
std::shared_ptr<dnnl::memory> filter_dst_mem = provider_->GetWeightsMemoryBuffer(mklnode_ptr_->weight_name);
if (filter_dst_mem == nullptr) {
ReorderWeights(api, context, GetEngine());
ReorderWeights(api, context, dnnl_engine_cpu_);
filter_dst_mem = provider_->GetWeightsMemoryBuffer(mklnode_ptr_->weight_name);
}
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
filter_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(filter_data)));
if (!gpu_available_) {
filter_data = static_cast<T*>(filter_dst_mem->get_data_handle());
filter_mem_->set_data_handle(static_cast<void*>(const_cast<T*>(filter_data)));
} else { // gpu_available_
#ifdef USE_DNNL_GPU_OCL
filter_mem_gpu_->set_ocl_mem_object(filter_dst_mem->get_ocl_mem_object());
#endif
}
std::shared_ptr<dnnl::memory> bias_mem = provider_->GetBiasMemoryBuffer(mklnode_ptr_->weight_name);
const T* bias_data = static_cast<T*>(bias_mem->get_data_handle());
@ -477,9 +563,11 @@ class DnnlConvBatchNorm : public DnnlKernel {
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
}
auto src_size = conv_fwd_pd_.get()->src_desc().get_size();
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
if (!gpu_available_) {
auto src_size = conv_fwd_pd_.get()->src_desc().get_size();
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
}
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
@ -496,10 +584,14 @@ class DnnlConvBatchNorm : public DnnlKernel {
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
T* dst_data = ort.GetTensorMutableData<T>(output);
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (primitive_dst_desc_ != ort_source_desc_) {
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
} else { // gpu_available_
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
}
return Status::OK();
@ -584,6 +676,10 @@ class DnnlConvBatchNorm : public DnnlKernel {
std::unique_ptr<dnnl::memory> filter_mem_;
std::unique_ptr<dnnl::memory> bias_mem_;
std::shared_ptr<dnnl::memory> src_mem_gpu_;
std::unique_ptr<dnnl::memory> filter_mem_gpu_;
std::unique_ptr<dnnl::memory> bias_mem_gpu_;
std::unique_ptr<dnnl::convolution_forward::desc> fwd_desc_;
std::unique_ptr<dnnl::memory::desc> src_md_;
@ -593,6 +689,11 @@ class DnnlConvBatchNorm : public DnnlKernel {
std::unique_ptr<dnnl::convolution_forward::primitive_desc> conv_fwd_pd_;
std::unique_ptr<dnnl::primitive> conv_fwd_;
dnnl::engine dnnl_engine_cpu_;
dnnl::engine dnnl_engine_gpu_;
bool gpu_available_;
private:
IAllocatorUniquePtr<void> src_reorder_buffer_;
IAllocatorUniquePtr<void> dst_reorder_buffer_;

View file

@ -18,6 +18,7 @@
namespace onnxruntime {
namespace ort_dnnl {
DnnlEngineInstance* DnnlEngineInstance::instance = 0;
namespace {
template <typename T>
@ -25,9 +26,14 @@ class SubgraphPrimitive : public PrimitiveBase {
public:
SubgraphPrimitive(const OrtCustomOpApi* api,
OrtKernelContext* context,
const SubgraphParams& params)
: cpu_engine_(GetEngine()) {
context_.stream = onnxruntime::make_unique<dnnl::stream>(dnnl::stream(cpu_engine_));
const SubgraphParams& params) {
dnnl_engine_instance_ = DnnlEngineInstance::getInstance();
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine_instance_->getEngineMap().find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine_instance_->getEngineMap().end()) {
context_.stream = onnxruntime::make_unique<dnnl::stream>(dnnl::stream(dnnl_engine_instance_->getEngine(dnnl::engine::kind::gpu)));
} else {
context_.stream = onnxruntime::make_unique<dnnl::stream>(dnnl::stream(dnnl_engine_instance_->getEngine(dnnl::engine::kind::cpu)));
}
if (context_.net.size() == 0) {
CreateKernels(params);
@ -196,15 +202,15 @@ class SubgraphPrimitive : public PrimitiveBase {
// Propagate Dnnl block format
// dst format of current node to src format of next node
for (auto& kernel : context_.kernels) {
kernel->CreatePrimitives(api, context, cpu_engine_, context_.net, context_.net_args);
kernel->CreatePrimitives(api, context, dnnl_engine_instance_->getEngineMap(), context_.net, context_.net_args);
if (kernel->primitive_created_status_.IsOK()) {
kernel->ReorderWeights(api, context, cpu_engine_);
kernel->ReorderWeights(api, context, dnnl_engine_instance_->getEngine(dnnl::engine::kind::cpu));
}
}
}
SubgraphContext context_;
dnnl::engine& cpu_engine_;
DnnlEngineInstance* dnnl_engine_instance_;
};
// Pool which allows for reuse of DNNL Conv primitives which are expensive to instantiate.

View file

@ -10,9 +10,10 @@ namespace ort_dnnl {
void DnnlKernel::InitDstReorderOutput(dnnl::engine& cpu_engine,
dnnl::memory::data_type& data_type,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) {
std::vector<std::unordered_map<int, dnnl::memory>>& net_args,
bool gpu_available) {
// Allocate dst buffer if reorder is necessary
if (primitive_dst_desc_ != ort_source_desc_) {
if (primitive_dst_desc_ != ort_source_desc_ || gpu_available) {
// reorder to ONNXRuntime format
dnnl::memory::dims dst_dims_mkl(
primitive_dst_shape_.GetDims().begin(), primitive_dst_shape_.GetDims().end());

View file

@ -26,11 +26,11 @@ class DnnlKernel {
virtual void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) = 0;
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, dnnl::engine& cpu_engine) {
virtual void ReorderWeights(const OrtCustomOpApi* api, OrtKernelContext* context, const dnnl::engine& cpu_engine) {
ORT_UNUSED_PARAMETER(api);
ORT_UNUSED_PARAMETER(context);
ORT_UNUSED_PARAMETER(cpu_engine);
@ -80,7 +80,8 @@ class DnnlKernel {
void InitDstReorderOutput(dnnl::engine& cpu_engine,
dnnl::memory::data_type& data_type,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args);
std::vector<std::unordered_map<int, dnnl::memory>>& net_args,
bool gpu_available = false);
dnnl::memory::format_tag GetSourceFormat(int dim_size);

View file

@ -21,9 +21,25 @@ class DnnlLrn : public DnnlKernel {
void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
dnnl::engine cpu_engine;
dnnl::engine engine_to_use;
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine.find(dnnl::engine::kind::cpu);
if (iter != dnnl_engine.end()) {
cpu_engine = (dnnl::engine)iter->second;
engine_to_use = cpu_engine;
}
gpu_available_ = false;
dnnl::engine gpu_engine;
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine.end()) {
gpu_engine = (dnnl::engine)(iter->second);
gpu_available_ = true;
engine_to_use = gpu_engine;
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
}
Ort::CustomOpApi ort{*api};
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
@ -48,10 +64,20 @@ class DnnlLrn : public DnnlKernel {
dnnl::memory::desc({src_dims}, DnnnType<T>(), ort_source_format_));
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(*src_md_, cpu_engine, nullptr));
if (gpu_available_) {
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(*src_md_, gpu_engine);
net.push_back(mkldnn::reorder(*src_mem_, *src_mem_gpu_));
net_args.push_back({{MKLDNN_ARG_SRC, *src_mem_},
{MKLDNN_ARG_DST, *src_mem_gpu_}});
}
} else {
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
dnnl::memory::desc(parents_[0].get()->primitive_dst_desc_));
src_mem_ = parents_[0].get()->primitive_dst_mem_;
if (!gpu_available_) {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
} else { // gpu_available_
src_mem_gpu_ = parents_[0].get()->primitive_dst_mem_;
}
x_shape = parents_[0].get()->primitive_dst_shape_;
ort_source_format_ = parents_[0].get()->ort_source_format_;
ort_source_desc_ = parents_[0].get()->ort_source_desc_;
@ -66,41 +92,51 @@ class DnnlLrn : public DnnlKernel {
size_, alpha_, beta_, bias_));
fwd_primitive_desc_ = onnxruntime::make_unique<dnnl::lrn_forward::primitive_desc>(
dnnl::lrn_forward::primitive_desc(*fwd_desc_, cpu_engine));
dnnl::lrn_forward::primitive_desc(*fwd_desc_, engine_to_use));
primitive_src_desc_ = fwd_primitive_desc_.get()->src_desc();
primitive_dst_desc_ = fwd_primitive_desc_.get()->dst_desc();
if (mklnode_ptr_->output_index >= 0) {
// last node of sub-graph. need to allocate memory for output_tensor
if (primitive_dst_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
if (!gpu_available_) {
if (mklnode_ptr_->output_index >= 0) {
// last node of sub-graph. need to allocate memory for output_tensor
if (primitive_dst_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine));
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// Intermediate node. Use Dnnl kernel internal memory for output and
// use this as input to next node.
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine));
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// Intermediate node. Use Dnnl kernel internal memory for output and
// use this as input to next node.
} else { // gpu_available_
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine));
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), gpu_engine));
}
lrn_fwd_ = onnxruntime::make_unique<dnnl::lrn_forward>(
dnnl::lrn_forward(*fwd_primitive_desc_));
net.push_back(*lrn_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
if (!gpu_available_) {
net.push_back(*lrn_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
} else { // gpu_available_
net.push_back(*lrn_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
}
if (mklnode_ptr_->output_index >= 0) {
// one of the end nodes. Allocate output buffer memory and
// reorder is necessary
dnnl::memory::data_type t = DnnnType<T>();
InitDstReorderOutput(cpu_engine, t, net, net_args);
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
}
}
@ -121,10 +157,14 @@ class DnnlLrn : public DnnlKernel {
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
T* dst_data = ort.GetTensorMutableData<T>(output);
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (primitive_dst_desc_ != ort_source_desc_) {
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
} else { // gpu_available_
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
}
@ -170,12 +210,15 @@ class DnnlLrn : public DnnlKernel {
private:
std::shared_ptr<dnnl::memory> src_mem_;
std::shared_ptr<dnnl::memory> src_mem_gpu_;
std::unique_ptr<dnnl::lrn_forward::desc> fwd_desc_;
std::unique_ptr<dnnl::lrn_forward::primitive_desc> fwd_primitive_desc_;
std::unique_ptr<dnnl::primitive> lrn_fwd_;
std::unique_ptr<dnnl::memory::desc> src_md_;
bool gpu_available_;
};
} // namespace ort_dnnl
} // namespace onnxruntime

View file

@ -21,8 +21,24 @@ class DnnlPool : public DnnlKernel {
void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine, std::vector<dnnl::primitive>& net,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine, std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
dnnl::engine cpu_engine;
dnnl::engine engine_to_use;
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine.find(dnnl::engine::kind::cpu);
if (iter != dnnl_engine.end()) {
cpu_engine = (dnnl::engine)iter->second;
engine_to_use = cpu_engine;
}
gpu_available_ = false;
dnnl::engine gpu_engine;
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine.end()) {
gpu_engine = (dnnl::engine)(iter->second);
gpu_available_ = true;
engine_to_use = gpu_engine;
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
}
Ort::CustomOpApi ort{*api};
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
@ -59,12 +75,17 @@ class DnnlPool : public DnnlKernel {
ort_source_desc_ = parents_[0].get()->ort_source_desc_;
source_desc_ = parents_[0].get()->primitive_dst_desc_;
if (source_desc_ == ort_source_desc_) {
// reorder for better performance
dnnl::memory::format_tag fmt = GetAVXFormat(src_dims_mkl);
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
dnnl::memory::desc({src_dims_mkl}, DnnnType<T>(), fmt));
} else {
if (!gpu_available_) {
if (source_desc_ == ort_source_desc_) {
// reorder for better performance
dnnl::memory::format_tag fmt = GetAVXFormat(src_dims_mkl);
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
dnnl::memory::desc({src_dims_mkl}, DnnnType<T>(), fmt));
} else {
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
dnnl::memory::desc(parents_[0].get()->primitive_dst_mem_->get_desc()));
}
} else { // gpu_available_
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
dnnl::memory::desc(parents_[0].get()->primitive_dst_mem_->get_desc()));
}
@ -107,78 +128,129 @@ class DnnlPool : public DnnlKernel {
strides_mkl, kernel_mkl,
padding_left_mkl, padding_right_mkl));
fwd_primitive_desc_ = onnxruntime::make_unique<dnnl::pooling_forward::primitive_desc>(
dnnl::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_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->src_desc(), cpu_engine, nullptr));
} else {
// Sub-graph's inner node. set input to parent's output
src_mem_ = parents_[0].get()->primitive_dst_mem_;
if (!gpu_available_) {
fwd_primitive_desc_ = onnxruntime::make_unique<dnnl::pooling_forward::primitive_desc>(
dnnl::pooling_forward::primitive_desc(*fwd_desc_, cpu_engine));
} else { // gpu_available_
fwd_primitive_desc_ = onnxruntime::make_unique<dnnl::pooling_forward::primitive_desc>(
dnnl::pooling_forward::primitive_desc(*fwd_desc_, gpu_engine));
}
//if (mklnode_ptr_->parent_nodes.empty()) {
// // Sub-graph's first node. Read input from input buffer
// src_mem_ = onnxruntime::make_unique<dnnl::memory>(
// dnnl::memory(fwd_primitive_desc_.get()->src_desc(), cpu_engine, nullptr));
//} else {
// // Sub-graph's inner node. set input to parent's output
// src_mem_ = parents_[0].get()->primitive_dst_mem_;
//}
//if (gpu_available_) {
// src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
// dnnl::memory(fwd_primitive_desc_.get()->src_desc(), gpu_engine));
//}
primitive_src_desc_ = fwd_primitive_desc_.get()->src_desc();
primitive_dst_desc_ = fwd_primitive_desc_.get()->dst_desc();
src_size_ = fwd_primitive_desc_.get()->src_desc().get_size();
dst_size_ = fwd_primitive_desc_.get()->dst_desc().get_size();
// reorder source memory for best performance (AVX512);
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
auto pd = dnnl::memory::desc(source_desc_);
if (!gpu_available_) {
// reorder source memory for best performance (AVX512);
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
auto pd = dnnl::memory::desc(source_desc_);
if (mklnode_ptr_->parent_nodes.empty())
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
else
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
if (mklnode_ptr_->parent_nodes.empty())
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
else
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_->src_desc(), cpu_engine, nullptr));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_}});
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_->src_desc(), cpu_engine, nullptr));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_}});
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_->src_desc(), cpu_engine, nullptr));
} else {
src_mem_ = parents_[0].get()->primitive_dst_mem_;
}
}
} else { // gpu_available_
if (primitive_src_desc_ != source_desc_) {
dnnl::memory::dims src_dims(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
auto pd = dnnl::memory::desc(source_desc_);
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(pd, cpu_engine, nullptr));
} else {
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
}
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_->src_desc(), gpu_engine));
net.push_back(dnnl::reorder(*src_mem_from_, *src_mem_gpu_));
net_args.push_back({{DNNL_ARG_FROM, *src_mem_from_},
{DNNL_ARG_TO, *src_mem_gpu_}});
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_->src_desc(), cpu_engine, nullptr));
src_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_->src_desc(), gpu_engine));
net.push_back(dnnl::reorder(*src_mem_, *src_mem_gpu_));
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *src_mem_gpu_}});
} else {
src_mem_gpu_ = 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_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
if (!gpu_available_) {
if (mklnode_ptr_->output_index >= 0) {
// last node of sub-graph. need to allocate memory for output_tensor
if (primitive_dst_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine));
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// Intermediate node. Use Dnnl kernel internal memory for output and
// use this as input to next node.
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine));
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine, nullptr));
}
} else {
// Intermediate node. Use Dnnl kernel internal memory for output and
// use this as input to next node.
} else { // gpu_available_
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), cpu_engine));
dnnl::memory(fwd_primitive_desc_.get()->dst_desc(), gpu_engine));
}
pool_fwd_ = onnxruntime::make_unique<dnnl::pooling_forward>(
dnnl::pooling_forward(*fwd_primitive_desc_));
net.push_back(*pool_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
if (!gpu_available_) {
net.push_back(*pool_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
} else { // gpu_available_
net.push_back(*pool_fwd_);
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
{DNNL_ARG_DST, *primitive_dst_mem_}});
}
if (mklnode_ptr_->output_index >= 0) {
// one of the end nodes. Allocate output buffer memory and
// reorder is necessary
dnnl::memory::data_type t = DnnnType<T>();
InitDstReorderOutput(cpu_engine, t, net, net_args);
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
}
}
@ -197,9 +269,11 @@ class DnnlPool : public DnnlKernel {
src_mem_from_ = parents_[0].get()->primitive_dst_mem_;
}
auto src_size = fwd_primitive_desc_.get()->src_desc().get_size();
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
if (!gpu_available_) {
auto src_size = fwd_primitive_desc_.get()->src_desc().get_size();
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
}
} else {
if (mklnode_ptr_->parent_nodes.empty()) {
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index);
@ -218,10 +292,14 @@ class DnnlPool : public DnnlKernel {
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
T* dst_data = ort.GetTensorMutableData<T>(output);
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (primitive_dst_desc_ != ort_source_desc_) {
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
} else { // gpu_available_
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
}
return Status::OK();
@ -297,6 +375,7 @@ class DnnlPool : public DnnlKernel {
size_t dst_size_;
std::shared_ptr<dnnl::memory> src_mem_;
std::shared_ptr<dnnl::memory> src_mem_gpu_;
std::unique_ptr<dnnl::pooling_forward::desc> fwd_desc_;
std::unique_ptr<dnnl::memory::desc> src_md_;
@ -309,13 +388,15 @@ class DnnlPool : public DnnlKernel {
std::unique_ptr<dnnl::memory> dst_mem_from_;
std::unique_ptr<dnnl::memory> dst_mem_to_;
bool gpu_available_;
private:
dnnl::memory::format_tag GetAVXFormat(const dnnl::memory::dims& src_dims_mkl) {
bool is_2D = src_dims_mkl.size() == 4 ? true : false;
dnnl::memory::format_tag fmt = dnnl::memory::format_tag::any;
if (CPUIDInfo::GetCPUIDInfo().HasAVX512f()) {
if (!gpu_available_ && CPUIDInfo::GetCPUIDInfo().HasAVX512f()) {
fmt = is_2D ? dnnl::memory::format_tag::nChw16c : dnnl::memory::format_tag::nCdhw16c;
} else if (CPUIDInfo::GetCPUIDInfo().HasAVX2() && (src_dims_mkl[1] % 8 == 0)) {
} else if (!gpu_available_ && CPUIDInfo::GetCPUIDInfo().HasAVX2() && (src_dims_mkl[1] % 8 == 0)) {
fmt = is_2D ? dnnl::memory::format_tag::nChw8c : dnnl::memory::format_tag::ncdhw;
} else {
fmt = is_2D ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::ncdhw;

View file

@ -21,9 +21,25 @@ class DnnlSum : public DnnlKernel {
void CreatePrimitives(const OrtCustomOpApi* api,
OrtKernelContext* context,
dnnl::engine& cpu_engine,
const std::unordered_map<dnnl::engine::kind, dnnl::engine>& dnnl_engine,
std::vector<dnnl::primitive>& net,
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
dnnl::engine cpu_engine;
dnnl::engine engine_to_use;
std::unordered_map<dnnl::engine::kind, dnnl::engine>::const_iterator iter = dnnl_engine.find(dnnl::engine::kind::cpu);
if (iter != dnnl_engine.end()) {
cpu_engine = (dnnl::engine)iter->second;
engine_to_use = cpu_engine;
}
gpu_available_ = false;
dnnl::engine gpu_engine;
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
if (iter != dnnl_engine.end()) {
gpu_engine = (dnnl::engine)(iter->second);
gpu_available_ = true;
engine_to_use = gpu_engine;
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
}
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;
@ -74,50 +90,79 @@ class DnnlSum : public DnnlKernel {
x_shape1.GetDims().begin(), x_shape1.GetDims().end());
auto mpd = dnnl::memory::desc({src_dims}, DnnnType<T>(), ort_source_format_);
auto src_memory = dnnl::memory(mpd, cpu_engine, nullptr);
if (gpu_available_) {
auto src_memory_gpu = dnnl::memory(mpd, gpu_engine);
net.push_back(dnnl::reorder(src_memory, src_memory_gpu));
net_args.push_back({{DNNL_ARG_SRC, src_memory},
{DNNL_ARG_DST, src_memory_gpu}});
srcs_memory_gpu_.push_back(src_memory_gpu);
}
srcs_pd_.push_back(mpd);
srcs_memory_.push_back(src_memory);
coeff.push_back(1.0);
} else {
x_shape = parents_[0].get()->primitive_dst_shape_;
auto mpd = dnnl::memory::desc(parents_[i].get()->primitive_dst_desc_);
auto src_memory = *parents_[i].get()->primitive_dst_mem_;
if (!gpu_available_) {
auto src_memory = *parents_[i].get()->primitive_dst_mem_;
srcs_memory_.push_back(src_memory);
} else { // gpu_available_
auto src_memory_gpu = *parents_[i].get()->primitive_dst_mem_;
srcs_memory_gpu_.push_back(src_memory_gpu);
}
srcs_pd_.push_back(mpd);
srcs_memory_.push_back(src_memory);
coeff.push_back(1.0);
}
}
primitive_dst_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
dnnl::memory::desc({dst_dims_mkl}, DnnnType<T>(), dnnl::memory::format_tag::any));
sum_pd_ = onnxruntime::make_unique<dnnl::sum::primitive_desc>(
dnnl::sum::primitive_desc(*primitive_dst_md_, coeff, srcs_pd_, cpu_engine));
if (!gpu_available_) {
sum_pd_ = onnxruntime::make_unique<dnnl::sum::primitive_desc>(
dnnl::sum::primitive_desc(*primitive_dst_md_, coeff, srcs_pd_, cpu_engine));
} else { // gpu_available_
sum_pd_ = onnxruntime::make_unique<dnnl::sum::primitive_desc>(
dnnl::sum::primitive_desc(*primitive_dst_md_, coeff, srcs_pd_, gpu_engine));
}
if (mklnode_ptr_->output_index >= 0) {
// last node of sub-graph. need to allocate memory for output_tensor
if (primitive_dst_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
if (!gpu_available_) {
if (mklnode_ptr_->output_index >= 0) {
// last node of sub-graph. need to allocate memory for output_tensor
if (primitive_dst_desc_ != ort_source_desc_) {
// reorder neded. Use primitive output as input to reorder and
// allocate buffer for reorder output, final output of this subgraph
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(sum_pd_->dst_desc(), cpu_engine));
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(sum_pd_->dst_desc(), cpu_engine, nullptr));
}
} else {
// Intermediate node. Use Dnnl kernel internal memory for output and
// use this as input to next node.
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(sum_pd_->dst_desc(), cpu_engine));
} else {
// Last node but re-order not needed. Allocate buffer to output of this node
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(sum_pd_->dst_desc(), cpu_engine, nullptr));
}
} else {
// Intermediate node. Use Dnnl kernel internal memory for output and
// use this as input to next node.
} else { // gpu_available_
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(sum_pd_->dst_desc(), cpu_engine));
dnnl::memory(sum_pd_->dst_desc(), gpu_engine));
}
primitive_dst_desc_ = sum_pd_->dst_desc();
auto c = dnnl::sum(*sum_pd_);
net.push_back(c);
std::unordered_map<int, dnnl::memory> args{
{DNNL_ARG_DST, *primitive_dst_mem_}};
for (int i = 0; i < (int)num_inputs; i++) {
args.insert({DNNL_ARG_MULTIPLE_SRC + i, srcs_memory_[i]});
std::unordered_map<int, dnnl::memory> args;
if (!gpu_available_) {
args.insert({DNNL_ARG_DST, *primitive_dst_mem_});
for (int i = 0; i < (int)num_inputs; i++) {
args.insert({DNNL_ARG_MULTIPLE_SRC + i, srcs_memory_[i]});
}
} else { // gpu_available_
args.insert({DNNL_ARG_DST, *primitive_dst_mem_});
for (int i = 0; i < (int)num_inputs; i++) {
args.insert({DNNL_ARG_MULTIPLE_SRC + i, srcs_memory_gpu_[i]});
}
}
net_args.push_back(args);
@ -125,7 +170,7 @@ class DnnlSum : public DnnlKernel {
// one of the end nodes. Allocate output buffer memory and
// reorder is necessary
dnnl::memory::data_type t = DnnnType<T>();
InitDstReorderOutput(cpu_engine, t, net, net_args);
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
}
}
@ -150,10 +195,14 @@ class DnnlSum : public DnnlKernel {
OrtValue* output = ort.KernelContext_GetOutput(context, mklnode_ptr_->output_index, &y_dims[0], static_cast<int>(primitive_dst_shape_.GetDims().size()));
T* dst_data = ort.GetTensorMutableData<T>(output);
if (primitive_dst_desc_ != ort_source_desc_) {
if (!gpu_available_) {
if (primitive_dst_desc_ != ort_source_desc_) {
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
} else { // gpu_available_
reorder_dst_mem_to_->set_data_handle(dst_data);
} else {
primitive_dst_mem_->set_data_handle(dst_data);
}
}
return Status::OK();
@ -163,11 +212,14 @@ class DnnlSum : public DnnlKernel {
std::unique_ptr<dnnl::memory::desc> src_md_;
std::vector<dnnl::memory::desc> src_mds_;
std::vector<dnnl::memory> srcs_memory_;
std::vector<dnnl::memory> srcs_memory_gpu_;
std::vector<dnnl::memory::desc> srcs_pd_;
std::unique_ptr<dnnl::memory::desc> src_mpd_;
std::unique_ptr<dnnl::memory::desc> dst_pd_;
std::unique_ptr<dnnl::sum::primitive_desc> sum_pd_;
bool gpu_available_;
};
} // namespace ort_dnnl
} // namespace onnxruntime

View file

@ -300,6 +300,13 @@ def parse_arguments():
"--use_openblas", action='store_true', help="Build with OpenBLAS.")
parser.add_argument(
"--use_dnnl", action='store_true', help="Build with DNNL.")
parser.add_argument(
"--dnnl_gpu_runtime", action='store', default='', type=str.lower,
help="e.g. --dnnl_gpu_runtime ocl")
parser.add_argument(
"--dnnl_opencl_root", action='store', default='',
help="Path to OpenCL SDK. "
"e.g. --dnnl_opencl_root \"C:/Program Files (x86)/IntelSWTools/sw_dev_tools/OpenCL/sdk\"")
parser.add_argument(
"--use_featurizers", action='store_true',
help="Build with ML Featurizer support.")
@ -666,6 +673,8 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
"OFF" if args.use_openblas else "ON"),
"-Donnxruntime_USE_OPENBLAS=" + ("ON" if args.use_openblas else "OFF"),
"-Donnxruntime_USE_DNNL=" + ("ON" if args.use_dnnl else "OFF"),
"-Donnxruntime_DNNL_GPU_RUNTIME=" + (args.dnnl_gpu_runtime if args.use_dnnl else ""),
"-Donnxruntime_DNNL_OPENCL_ROOT=" + (args.dnnl_opencl_root if args.use_dnnl else ""),
"-Donnxruntime_USE_NGRAPH=" + ("ON" if args.use_ngraph else "OFF"),
"-Donnxruntime_USE_NNAPI_BUILTIN=" + ("ON" if args.use_nnapi else "OFF"),
"-Donnxruntime_USE_RKNPU=" + ("ON" if args.use_rknpu else "OFF"),