mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-13 18:08:13 +00:00
add dnnl gpu engine (#5788)
This commit is contained in:
parent
0ea998134a
commit
435b904f0e
15 changed files with 870 additions and 283 deletions
16
BUILD.md
16
BUILD.md
|
|
@ -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"`
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
16
cmake/external/dnnl.cmake
vendored
16
cmake/external/dnnl.cmake
vendored
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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, '#');
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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_;
|
||||
|
|
|
|||
|
|
@ -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_;
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
Loading…
Reference in a new issue