From 435b904f0e8dd845ff4cd9f5efbe7dd66b1d9219 Mon Sep 17 00:00:00 2001 From: jeyblu <63478620+jeyblu@users.noreply.github.com> Date: Thu, 12 Nov 2020 20:17:54 -0800 Subject: [PATCH] add dnnl gpu engine (#5788) --- BUILD.md | 16 +- cmake/external/dnnl.cmake | 16 +- cmake/onnxruntime_providers.cmake | 5 +- onnxruntime/core/providers/dnnl/dnnl_common.h | 49 +++- .../dnnl/subgraph/dnnl_activations.h | 88 +++++-- .../providers/dnnl/subgraph/dnnl_batchnorm.h | 129 ++++++++--- .../core/providers/dnnl/subgraph/dnnl_conv.h | 218 +++++++++++++----- .../dnnl/subgraph/dnnl_conv_batchnorm.h | 199 ++++++++++++---- .../dnnl/subgraph/dnnl_func_kernel.cc | 18 +- .../providers/dnnl/subgraph/dnnl_kernel.cc | 5 +- .../providers/dnnl/subgraph/dnnl_kernel.h | 7 +- .../core/providers/dnnl/subgraph/dnnl_lrn.h | 91 ++++++-- .../core/providers/dnnl/subgraph/dnnl_pool.h | 199 +++++++++++----- .../core/providers/dnnl/subgraph/dnnl_sum.h | 104 ++++++--- tools/ci_build/build.py | 9 + 15 files changed, 870 insertions(+), 283 deletions(-) diff --git a/BUILD.md b/BUILD.md index f5f428a550..6f4d79d597 100644 --- a/BUILD.md +++ b/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"` --- diff --git a/cmake/external/dnnl.cmake b/cmake/external/dnnl.cmake index e9a24fccc7..aa194c7714 100644 --- a/cmake/external/dnnl.cmake +++ b/cmake/external/dnnl.cmake @@ -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() diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index 05d6a4344b..8ef3262906 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -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") diff --git a/onnxruntime/core/providers/dnnl/dnnl_common.h b/onnxruntime/core/providers/dnnl/dnnl_common.h index 90c22a9702..7b7ce54f3e 100644 --- a/onnxruntime/core/providers/dnnl/dnnl_common.h +++ b/onnxruntime/core/providers/dnnl/dnnl_common.h @@ -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 #include @@ -20,10 +21,50 @@ dnnl::memory::data_type DnnnType() { 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_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& getEngineMap() { + return dnnl_engine_map; + } + + const dnnl::engine& getEngine(dnnl::engine::kind kind) { + std::unordered_map::iterator iter = dnnl_engine_map.find(kind); + return iter->second; + } +}; static void AddDimsToKey(std::string& key, const dnnl::memory::dims& dims) { key.append(1, '#'); diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_activations.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_activations.h index e19d795c1f..b96cc1b7ea 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_activations.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_activations.h @@ -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, std::vector& net, std::vector>& net_args) { + dnnl::engine cpu_engine; + dnnl::engine engine_to_use; + std::unordered_map::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(), ort_source_format_)); src_mem_ = onnxruntime::make_unique( dnnl::memory({{src_dims}, DnnnType(), ort_source_format_}, cpu_engine, nullptr)); + if (gpu_available_) { + src_mem_gpu_ = onnxruntime::make_unique(*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(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::prop_kind::forward_inference, algo, *src_md_, 0)); relu_fwd_pd_ = onnxruntime::make_unique( - 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(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(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(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(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(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(relu_fwd_pd_.get()->dst_desc(), cpu_engine)); + } else { // gpu_available_ + primitive_dst_mem_ = std::make_shared(dnnl::memory(relu_fwd_pd_.get()->dst_desc(), gpu_engine)); } relu_fwd_ = onnxruntime::make_unique( 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(); - 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(primitive_dst_shape_.GetDims().size())); T* dst_data = ort.GetTensorMutableData(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 src_mem_; + std::shared_ptr src_mem_gpu_; std::unique_ptr fwd_desc_; std::unique_ptr relu_fwd_pd_; std::unique_ptr relu_fwd_; std::unique_ptr src_md_; + + bool gpu_available_; }; } // namespace ort_dnnl } // namespace onnxruntime diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_batchnorm.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_batchnorm.h index 5ad9f307ae..0b2f5e7a88 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_batchnorm.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_batchnorm.h @@ -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, std::vector& net, std::vector>& net_args) override { + dnnl::engine cpu_engine; + dnnl::engine engine_to_use; + std::unordered_map::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(*var_md_, cpu_engine, nullptr)); + if (gpu_available_) { + scale_shift_mem_gpu_ = onnxruntime::make_unique( + dnnl::memory({*scale_shift_md_, gpu_engine})); + mean_mem_gpu_ = onnxruntime::make_unique( + dnnl::memory(*mean_md_, gpu_engine)); + var_mem_gpu_ = onnxruntime::make_unique( + dnnl::memory(*var_md_, gpu_engine)); + } + batchnorm_fwd_ = onnxruntime::make_unique( 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(*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( - *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( batchnorm_fwd_pd_.get()->dst_desc()); - if (mklnode_ptr_->parent_nodes.empty()) { - src_mem_ = onnxruntime::make_unique( - 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(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(batchnorm_fwd_pd_.get()->src_desc(), cpu_engine, nullptr)); + src_mem_gpu_ = onnxruntime::make_unique( + 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(batchnorm_fwd_pd_->dst_desc(), cpu_engine)); + } else { + primitive_dst_mem_ = onnxruntime::make_unique( + 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(batchnorm_fwd_pd_->dst_desc(), cpu_engine)); - } else { - primitive_dst_mem_ = onnxruntime::make_unique( - 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(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(); - 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(primitive_dst_shape_.GetDims().size())); T* dst_data = ort.GetTensorMutableData(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 var_mem_; std::unique_ptr dst_mem_; + std::shared_ptr src_mem_gpu_; + std::unique_ptr scale_shift_mem_gpu_; + std::unique_ptr mean_mem_gpu_; + std::unique_ptr var_mem_gpu_; + std::unique_ptr src_md_; std::unique_ptr scale_shift_md_; std::unique_ptr mean_md_; @@ -359,6 +432,8 @@ class DnnlBatchNorm : public DnnlKernel { std::unique_ptr batchnorm_fwd_; std::unique_ptr batchnorm_fwd_pd_; + bool gpu_available_; + protected: float epsilon_ = 1e-5f; }; diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h index 1baf516a9c..ce2a895440 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h @@ -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, std::vector& net, std::vector>& net_args) override { + dnnl::engine cpu_engine; + dnnl::engine engine_to_use; + std::unordered_map::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(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(*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(*fwd_desc_, cpu_engine)); + dnnl::convolution_forward::primitive_desc(*fwd_desc_, engine_to_use)); } primitive_src_desc_ = static_cast( @@ -288,44 +306,84 @@ class DnnlConv : public DnnlKernel { filter_mem_ = onnxruntime::make_unique( dnnl::memory(conv_fwd_pd_.get()->weights_desc(), cpu_engine, nullptr)); + if (gpu_available_) { + filter_mem_gpu_ = onnxruntime::make_unique( + 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(), 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(), ort_source_format_}); - if (mklnode_ptr_->parent_nodes.empty()) - src_mem_from_ = onnxruntime::make_unique( - 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(pd, cpu_engine, nullptr)); + else + src_mem_from_ = parents_[0].get()->primitive_dst_mem_; - src_mem_ = onnxruntime::make_unique( - 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(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(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(), ort_source_format_}); + + if (mklnode_ptr_->parent_nodes.empty()) { + src_mem_from_ = onnxruntime::make_unique( + dnnl::memory(pd, cpu_engine, nullptr)); + } else { + src_mem_from_ = parents_[0].get()->primitive_dst_mem_; + } + src_mem_gpu_ = onnxruntime::make_unique( + 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(conv_fwd_pd_->src_desc(), cpu_engine, nullptr)); + src_mem_gpu_ = onnxruntime::make_unique( + 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(conv_fwd_pd_.get()->dst_desc(), cpu_engine)); + } else { + primitive_dst_mem_ = onnxruntime::make_unique( + 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(conv_fwd_pd_.get()->dst_desc(), cpu_engine)); - } else { - primitive_dst_mem_ = onnxruntime::make_unique( - 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(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(*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(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(*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(); - 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(), filter_format_}, cpu_engine, (void*)filter_data); IAllocatorUniquePtr filter_reorder_buffer = IAllocator::MakeUniquePtr(alloc_, filter_size_); - filter_dst_mem = onnxruntime::make_unique( - dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get())); + if (!gpu_available_) { + filter_dst_mem = onnxruntime::make_unique( + 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(filter_dst_mem->get_data_handle()); + } else { // gpu_available_ + filter_dst_mem = onnxruntime::make_unique( + 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(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 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(filter_dst_mem->get_data_handle()); + if (!gpu_available_) { + filter_data = static_cast(filter_dst_mem->get_data_handle()); + filter_mem_->set_data_handle(static_cast(const_cast(filter_data))); + } else { // gpu_available_ +#ifdef USE_DNNL_GPU_OCL + std::lock_guard 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(const_cast(filter_data))); if (bias_data != nullptr) { bias_mem_->set_data_handle(static_cast(const_cast(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(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(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(primitive_dst_shape_.GetDims().size())); T* dst_data = ort.GetTensorMutableData(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 filter_mem_; std::unique_ptr bias_mem_; + std::shared_ptr src_mem_gpu_; + std::unique_ptr filter_mem_gpu_; + std::unique_ptr bias_mem_gpu_; + std::unique_ptr fwd_desc_; std::unique_ptr src_md_; @@ -546,6 +651,11 @@ class DnnlConv : public DnnlKernel { std::unique_ptr conv_fwd_pd_; std::unique_ptr conv_fwd_; + dnnl::engine dnnl_engine_cpu_; + dnnl::engine dnnl_engine_gpu_; + + bool gpu_available_; + private: IAllocatorUniquePtr src_reorder_buffer_; IAllocatorUniquePtr dst_reorder_buffer_; diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h index 9b1a8b2104..7e2654d6fc 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h @@ -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, std::vector& net, std::vector>& net_args) override { + dnnl::engine cpu_engine; + dnnl::engine engine_to_use; + std::unordered_map::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(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(*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(*fwd_desc_, cpu_engine)); + dnnl::convolution_forward::primitive_desc(*fwd_desc_, engine_to_use)); } primitive_src_desc_ = static_cast( @@ -241,61 +259,114 @@ class DnnlConvBatchNorm : public DnnlKernel { filter_mem_ = onnxruntime::make_unique( dnnl::memory(conv_fwd_pd_.get()->weights_desc(), cpu_engine, nullptr)); + if (gpu_available_) { + filter_mem_gpu_ = onnxruntime::make_unique( + 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(), 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(), ort_source_format_}); - if (mklnode_ptr_->parent_nodes.empty()) - src_mem_from_ = onnxruntime::make_unique( - 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(pd, cpu_engine, nullptr)); + else + src_mem_from_ = parents_[0].get()->primitive_dst_mem_; - src_mem_ = onnxruntime::make_unique( - 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(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(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(), ort_source_format_}); + + if (mklnode_ptr_->parent_nodes.empty()) { + src_mem_from_ = onnxruntime::make_unique( + dnnl::memory(pd, cpu_engine, nullptr)); + } else { + src_mem_from_ = parents_[0].get()->primitive_dst_mem_; + } + src_mem_gpu_ = onnxruntime::make_unique( + 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(conv_fwd_pd_->src_desc(), cpu_engine, nullptr)); + src_mem_gpu_ = onnxruntime::make_unique( + 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(conv_fwd_pd_.get()->dst_desc(), cpu_engine)); + } else { + primitive_dst_mem_ = onnxruntime::make_unique( + 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(conv_fwd_pd_.get()->dst_desc(), cpu_engine)); - } else { - primitive_dst_mem_ = onnxruntime::make_unique( - 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(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(conv_fwd_pd_.get()->bias_desc(), cpu_engine, nullptr)); conv_fwd_ = onnxruntime::make_unique( 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(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(); - 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(), filter_format_}, cpu_engine, (void*)weights_scaled_by_axis.data()); IAllocatorUniquePtr filter_reorder_buffer = IAllocator::MakeUniquePtr(alloc_, filter_size_); - filter_dst_mem = onnxruntime::make_unique( - dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get())); + if (!gpu_available_) { + filter_dst_mem = onnxruntime::make_unique( + 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(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 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(filter_dst_mem->get_data_handle()); - filter_mem_->set_data_handle(static_cast(const_cast(filter_data))); + if (!gpu_available_) { + filter_data = static_cast(filter_dst_mem->get_data_handle()); + filter_mem_->set_data_handle(static_cast(const_cast(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 bias_mem = provider_->GetBiasMemoryBuffer(mklnode_ptr_->weight_name); const T* bias_data = static_cast(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(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(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(primitive_dst_shape_.GetDims().size())); T* dst_data = ort.GetTensorMutableData(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 filter_mem_; std::unique_ptr bias_mem_; + std::shared_ptr src_mem_gpu_; + std::unique_ptr filter_mem_gpu_; + std::unique_ptr bias_mem_gpu_; + std::unique_ptr fwd_desc_; std::unique_ptr src_md_; @@ -593,6 +689,11 @@ class DnnlConvBatchNorm : public DnnlKernel { std::unique_ptr conv_fwd_pd_; std::unique_ptr conv_fwd_; + dnnl::engine dnnl_engine_cpu_; + dnnl::engine dnnl_engine_gpu_; + + bool gpu_available_; + private: IAllocatorUniquePtr src_reorder_buffer_; IAllocatorUniquePtr dst_reorder_buffer_; diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_func_kernel.cc b/onnxruntime/core/providers/dnnl/subgraph/dnnl_func_kernel.cc index de21724e47..1047560f2c 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_func_kernel.cc +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_func_kernel.cc @@ -18,6 +18,7 @@ namespace onnxruntime { namespace ort_dnnl { +DnnlEngineInstance* DnnlEngineInstance::instance = 0; namespace { template @@ -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(cpu_engine_)); + const SubgraphParams& params) { + dnnl_engine_instance_ = DnnlEngineInstance::getInstance(); + std::unordered_map::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_engine_instance_->getEngine(dnnl::engine::kind::gpu))); + } else { + context_.stream = onnxruntime::make_unique(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. diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.cc b/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.cc index 9c07aab299..a5d758d046 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.cc +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.cc @@ -10,9 +10,10 @@ namespace ort_dnnl { void DnnlKernel::InitDstReorderOutput(dnnl::engine& cpu_engine, dnnl::memory::data_type& data_type, std::vector& net, - std::vector>& net_args) { + std::vector>& 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()); diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h index f215509b21..2aefaf081b 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h @@ -26,11 +26,11 @@ class DnnlKernel { virtual void CreatePrimitives(const OrtCustomOpApi* api, OrtKernelContext* context, - dnnl::engine& cpu_engine, + const std::unordered_map& dnnl_engine, std::vector& net, std::vector>& 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& net, - std::vector>& net_args); + std::vector>& net_args, + bool gpu_available = false); dnnl::memory::format_tag GetSourceFormat(int dim_size); diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_lrn.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_lrn.h index dc07853991..38ac28db99 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_lrn.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_lrn.h @@ -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, std::vector& net, std::vector>& net_args) override { + dnnl::engine cpu_engine; + dnnl::engine engine_to_use; + std::unordered_map::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(), ort_source_format_)); src_mem_ = onnxruntime::make_unique( dnnl::memory(*src_md_, cpu_engine, nullptr)); + if (gpu_available_) { + src_mem_gpu_ = onnxruntime::make_unique(*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(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(*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(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(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(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(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(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(*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(); - 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(primitive_dst_shape_.GetDims().size())); T* dst_data = ort.GetTensorMutableData(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 src_mem_; + std::shared_ptr src_mem_gpu_; std::unique_ptr fwd_desc_; std::unique_ptr fwd_primitive_desc_; std::unique_ptr lrn_fwd_; std::unique_ptr src_md_; + + bool gpu_available_; }; } // namespace ort_dnnl } // namespace onnxruntime diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h index 1e396eff93..0b8b0d1cce 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h @@ -21,8 +21,24 @@ class DnnlPool : public DnnlKernel { void CreatePrimitives(const OrtCustomOpApi* api, OrtKernelContext* context, - dnnl::engine& cpu_engine, std::vector& net, + const std::unordered_map& dnnl_engine, std::vector& net, std::vector>& net_args) override { + dnnl::engine cpu_engine; + dnnl::engine engine_to_use; + std::unordered_map::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({src_dims_mkl}, DnnnType(), 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({src_dims_mkl}, DnnnType(), fmt)); + } else { + src_md_ = onnxruntime::make_unique( + dnnl::memory::desc(parents_[0].get()->primitive_dst_mem_->get_desc())); + } + } else { // gpu_available_ src_md_ = onnxruntime::make_unique( 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(*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(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(*fwd_desc_, cpu_engine)); + } else { // gpu_available_ + fwd_primitive_desc_ = onnxruntime::make_unique( + 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(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(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(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(pd, cpu_engine, nullptr)); + else + src_mem_from_ = parents_[0].get()->primitive_dst_mem_; - src_mem_ = onnxruntime::make_unique( - 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(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(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(pd, cpu_engine, nullptr)); + } else { + src_mem_from_ = parents_[0].get()->primitive_dst_mem_; + } + src_mem_gpu_ = onnxruntime::make_unique( + 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(fwd_primitive_desc_->src_desc(), cpu_engine, nullptr)); + src_mem_gpu_ = onnxruntime::make_unique( + 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(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(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(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(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(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(*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(); - 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(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(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(primitive_dst_shape_.GetDims().size())); T* dst_data = ort.GetTensorMutableData(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 src_mem_; + std::shared_ptr src_mem_gpu_; std::unique_ptr fwd_desc_; std::unique_ptr src_md_; @@ -309,13 +388,15 @@ class DnnlPool : public DnnlKernel { std::unique_ptr dst_mem_from_; std::unique_ptr 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; diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_sum.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_sum.h index e1edd95d9f..bfeefb95c5 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_sum.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_sum.h @@ -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, std::vector& net, std::vector>& net_args) override { + dnnl::engine cpu_engine; + dnnl::engine engine_to_use; + std::unordered_map::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(), 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({dst_dims_mkl}, DnnnType(), dnnl::memory::format_tag::any)); - sum_pd_ = onnxruntime::make_unique( - dnnl::sum::primitive_desc(*primitive_dst_md_, coeff, srcs_pd_, cpu_engine)); + if (!gpu_available_) { + sum_pd_ = onnxruntime::make_unique( + dnnl::sum::primitive_desc(*primitive_dst_md_, coeff, srcs_pd_, cpu_engine)); + } else { // gpu_available_ + sum_pd_ = onnxruntime::make_unique( + 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(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(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(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(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(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 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 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(); - 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(primitive_dst_shape_.GetDims().size())); T* dst_data = ort.GetTensorMutableData(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 src_md_; std::vector src_mds_; std::vector srcs_memory_; + std::vector srcs_memory_gpu_; std::vector srcs_pd_; std::unique_ptr src_mpd_; std::unique_ptr dst_pd_; std::unique_ptr sum_pd_; + + bool gpu_available_; }; } // namespace ort_dnnl } // namespace onnxruntime diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 3bae668432..7465872f0b 100755 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -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"),