mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Add GPU support for DNNL endpoint (#6741)
* Added code for Relugrad with GPU support. Signed-off-by: Chethan Palangotu Keshava <chethan.palangotu.keshava@intel.com> * Add GPU support for DNNL ConvGrad Signed-off-by: George Nash <george.nash@intel.com> * Add GPU support for DNNL MaxPoolGrad Updates to MaxPool for training with GPU Update oneDNN to version 1.8.1 Signed-off-by: George Nash <george.nash@intel.com> * Fixed issues found durring code review - error in code comment - using auto when the direct type would have been better - removed ternary operators that were returning bool values Signed-off-by: George Nash <george.nash@intel.com> Co-authored-by: Chethan Palangotu Keshava <chethan.palangotu.keshava@intel.com>
This commit is contained in:
George Nash
GitHub
parent
c8e2e3191b
commit
ba51774a1f
5 changed files with 315 additions and 195 deletions
6
cmake/external/dnnl.cmake
vendored
6
cmake/external/dnnl.cmake
vendored
|
|
@ -2,7 +2,7 @@ include (ExternalProject)
|
|||
|
||||
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)
|
||||
set(DNNL_TAG v1.8.1)
|
||||
|
||||
if(WIN32)
|
||||
set(DNNL_SHARED_LIB dnnl.dll)
|
||||
|
|
@ -15,9 +15,7 @@ else()
|
|||
endif()
|
||||
endif()
|
||||
|
||||
if(onnxruntime_USE_DNNL AND onnxruntime_ENABLE_TRAINING AND onnxruntime_DNNL_GPU_RUNTIME STREQUAL "ocl")
|
||||
message(FATAL_ERROR "--enable_training not supported with dnnl GPU runtime. Remove '--enable_training' or remove '--dnnl_gpu_runtime'.")
|
||||
elseif(onnxruntime_USE_DNNL AND onnxruntime_DNNL_GPU_RUNTIME STREQUAL "ocl" AND onnxruntime_DNNL_OPENCL_ROOT STREQUAL "")
|
||||
if(onnxruntime_USE_DNNL AND onnxruntime_DNNL_GPU_RUNTIME STREQUAL "ocl" AND onnxruntime_DNNL_OPENCL_ROOT STREQUAL "")
|
||||
message(FATAL_ERROR "--dnnl_opencl_root required")
|
||||
elseif(onnxruntime_USE_DNNL AND onnxruntime_DNNL_GPU_RUNTIME STREQUAL "" AND NOT (onnxruntime_DNNL_OPENCL_ROOT STREQUAL ""))
|
||||
message(FATAL_ERROR "--dnnl_gpu_runtime required")
|
||||
|
|
|
|||
|
|
@ -67,12 +67,11 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
|
||||
dnnl::engine cpu_engine;
|
||||
dnnl::engine engine_to_use;
|
||||
const auto iter = dnnl_engine.find(dnnl::engine::kind::cpu);
|
||||
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 = iter->second;
|
||||
engine_to_use = cpu_engine;
|
||||
}
|
||||
#if 0 // TODO update convgrad for gpu
|
||||
gpu_available_ = false;
|
||||
dnnl::engine gpu_engine;
|
||||
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
|
||||
|
|
@ -82,7 +81,6 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
engine_to_use = gpu_engine;
|
||||
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
|
||||
}
|
||||
#endif
|
||||
Ort::CustomOpApi ort{*api};
|
||||
stream_ = onnxruntime::make_unique<dnnl::stream>(dnnl::stream(engine_to_use));
|
||||
|
||||
|
|
@ -231,7 +229,7 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
TensorShape db_shape({wshape[0]});
|
||||
diff_bias_shape_ = db_shape;
|
||||
diff_bias_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
dnnl::memory::desc({wshape[0]}, DnnnType<T>(), dnnl::memory::format_tag::any));
|
||||
dnnl::memory::desc({wshape[0]}, DnnnType<T>(), dnnl::memory::format_tag::x));
|
||||
|
||||
dnnl::memory::dims filter_dims_mkl;
|
||||
if (group_mkl == 1) {
|
||||
|
|
@ -314,33 +312,50 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
*conv_bwd_weights_desc_, engine_to_use, *(conv_fwd_->GetPrimitiveDesc())));
|
||||
|
||||
diff_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(conv_bwd_weights_pd_.get()->diff_dst_desc(), engine_to_use, nullptr));
|
||||
dnnl::memory(conv_bwd_weights_pd_.get()->diff_dst_desc(), cpu_engine, nullptr));
|
||||
|
||||
src_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(conv_bwd_weights_pd_.get()->src_desc(), engine_to_use, nullptr));
|
||||
dnnl::memory(conv_bwd_weights_pd_.get()->src_desc(), cpu_engine, nullptr));
|
||||
|
||||
weights_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(conv_bwd_data_pd_.get()->weights_desc(), engine_to_use, nullptr));
|
||||
dnnl::memory(conv_bwd_data_pd_.get()->weights_desc(), cpu_engine, nullptr));
|
||||
|
||||
//diff_src_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
// dnnl::memory(conv_bwd_data_pd_.get()->diff_src_desc(), engine_to_use, nullptr));
|
||||
// dnnl::memory(conv_bwd_data_pd_.get()->diff_src_desc(), cpu_engine, nullptr));
|
||||
|
||||
if (gpu_available_) {
|
||||
diff_dst_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(*diff_dst_md_, gpu_engine);
|
||||
net.push_back(mkldnn::reorder(*diff_dst_mem_, *diff_dst_mem_gpu_));
|
||||
net_args.push_back({{MKLDNN_ARG_SRC, *diff_dst_mem_}, {MKLDNN_ARG_DST, *diff_dst_mem_gpu_}});
|
||||
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_}});
|
||||
weights_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(*weights_md_, gpu_engine);
|
||||
net.push_back(mkldnn::reorder(*weights_mem_, *weights_mem_gpu_));
|
||||
net_args.push_back({{MKLDNN_ARG_SRC, *weights_mem_}, {MKLDNN_ARG_DST, *weights_mem_gpu_}});
|
||||
}
|
||||
|
||||
primitive_src_desc_ = conv_bwd_data_pd_.get()->diff_dst_desc();
|
||||
primitive_dst_desc_ = conv_bwd_data_pd_.get()->diff_src_desc();
|
||||
|
||||
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_bwd_data_pd_.get()->diff_src_desc(), engine_to_use));
|
||||
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_bwd_data_pd_.get()->diff_src_desc(), cpu_engine));
|
||||
} else {
|
||||
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(conv_bwd_data_pd_.get()->diff_src_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_bwd_data_pd_.get()->diff_src_desc(), engine_to_use, nullptr));
|
||||
dnnl::memory(conv_bwd_data_pd_.get()->diff_src_desc(), cpu_engine));
|
||||
}
|
||||
} else {
|
||||
// last node of sub-graph. need to allocate memory for output_tensor
|
||||
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(conv_bwd_data_pd_.get()->diff_src_desc(), engine_to_use));
|
||||
dnnl::memory(conv_bwd_data_pd_.get()->diff_src_desc(), gpu_engine));
|
||||
}
|
||||
|
||||
diff_weights_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
|
|
@ -350,19 +365,45 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
dnnl::memory(conv_bwd_weights_pd_.get()->diff_bias_desc(), engine_to_use));
|
||||
|
||||
net.push_back(dnnl::convolution_backward_data(*conv_bwd_data_pd_));
|
||||
net_args.push_back({{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_WEIGHTS, *weights_mem_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_}});
|
||||
if (!gpu_available_) {
|
||||
net_args.push_back({{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_WEIGHTS, *weights_mem_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_}});
|
||||
} else {
|
||||
net_args.push_back({{DNNL_ARG_DIFF_DST, *diff_dst_mem_gpu_},
|
||||
{DNNL_ARG_WEIGHTS, *weights_mem_gpu_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_}});
|
||||
}
|
||||
|
||||
net.push_back(dnnl::convolution_backward_weights(*conv_bwd_weights_pd_));
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
|
||||
{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_DIFF_BIAS, *diff_bias_mem_},
|
||||
{DNNL_ARG_DIFF_WEIGHTS, *diff_weights_mem_}});
|
||||
if (!gpu_available_) {
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
|
||||
{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_DIFF_BIAS, *diff_bias_mem_},
|
||||
{DNNL_ARG_DIFF_WEIGHTS, *diff_weights_mem_}});
|
||||
} else {
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
|
||||
{DNNL_ARG_DIFF_DST, *diff_dst_mem_gpu_},
|
||||
{DNNL_ARG_DIFF_BIAS, *diff_bias_mem_},
|
||||
{DNNL_ARG_DIFF_WEIGHTS, *diff_weights_mem_}});
|
||||
}
|
||||
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
dnnl::memory::data_type t = DnnnType<T>();
|
||||
InitDstReorderOutput(engine_to_use, t, net, net_args);
|
||||
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
|
||||
// Allocate dst buffer if reorder is necessary
|
||||
if (gpu_available_) {
|
||||
// reorder to ONNXRuntime format
|
||||
diff_weights_reorder_mem_to_ = onnxruntime::make_unique<dnnl::memory>(dnnl::memory(*diff_weights_md_, cpu_engine));
|
||||
net.push_back(dnnl::reorder(*diff_weights_mem_, *diff_weights_reorder_mem_to_));
|
||||
net_args.push_back({{DNNL_ARG_FROM, *diff_weights_mem_},
|
||||
{DNNL_ARG_TO, *diff_weights_reorder_mem_to_}});
|
||||
|
||||
diff_bias_reorder_mem_to_ = onnxruntime::make_unique<dnnl::memory>(dnnl::memory(*diff_bias_md_, cpu_engine));
|
||||
net.push_back(dnnl::reorder(*diff_bias_mem_, *diff_bias_reorder_mem_to_));
|
||||
net_args.push_back({{DNNL_ARG_FROM, *diff_bias_mem_},
|
||||
{DNNL_ARG_TO, *diff_bias_reorder_mem_to_}});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -392,21 +433,33 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
OrtValue* dx_output = ort.KernelContext_GetOutput(context, 0, &dx_dims[0], static_cast<int>(dx_dims.size()));
|
||||
T* diff_src_data = ort.GetTensorMutableData<T>(dx_output);
|
||||
|
||||
if (primitive_dst_desc_ != ort_source_desc_) {
|
||||
reorder_dst_mem_to_->set_data_handle(diff_src_data);
|
||||
if (!gpu_available_) {
|
||||
if (primitive_dst_desc_ != ort_source_desc_) {
|
||||
reorder_dst_mem_to_->set_data_handle(diff_src_data);
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(diff_src_data);
|
||||
}
|
||||
} else {
|
||||
primitive_dst_mem_->set_data_handle(diff_src_data);
|
||||
reorder_dst_mem_to_->set_data_handle(diff_src_data);
|
||||
}
|
||||
|
||||
auto& dw_dims = diff_weights_shape_.GetDims();
|
||||
OrtValue* dw_output = ort.KernelContext_GetOutput(context, 1, &dw_dims[0], static_cast<int>(dw_dims.size()));
|
||||
T* dw_data = ort.GetTensorMutableData<T>(dw_output);
|
||||
diff_weights_mem_->set_data_handle(dw_data);
|
||||
if (!gpu_available_) {
|
||||
diff_weights_mem_->set_data_handle(dw_data);
|
||||
} else {
|
||||
diff_weights_reorder_mem_to_->set_data_handle(dw_data);
|
||||
}
|
||||
|
||||
auto& db_dims = diff_bias_shape_.GetDims();
|
||||
OrtValue* db_output = ort.KernelContext_GetOutput(context, 2, &db_dims[0], static_cast<int>(db_dims.size()));
|
||||
T* db_data = ort.GetTensorMutableData<T>(db_output);
|
||||
diff_bias_mem_->set_data_handle(db_data);
|
||||
if (!gpu_available_) {
|
||||
diff_bias_mem_->set_data_handle(db_data);
|
||||
} else {
|
||||
diff_bias_reorder_mem_to_->set_data_handle(db_data);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
|
@ -496,12 +549,15 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
|
||||
// input tensors
|
||||
std::unique_ptr<dnnl::memory> diff_dst_mem_;
|
||||
std::unique_ptr<dnnl::memory> diff_dst_mem_gpu_;
|
||||
std::unique_ptr<dnnl::memory::desc> diff_dst_md_;
|
||||
|
||||
std::unique_ptr<dnnl::memory> src_mem_;
|
||||
std::unique_ptr<dnnl::memory> src_mem_gpu_;
|
||||
std::unique_ptr<dnnl::memory::desc> src_md_;
|
||||
|
||||
std::unique_ptr<dnnl::memory> weights_mem_;
|
||||
std::unique_ptr<dnnl::memory> weights_mem_gpu_;
|
||||
std::unique_ptr<dnnl::memory::desc> weights_md_;
|
||||
//TensorShape bwd_src_shape_;
|
||||
|
||||
|
|
@ -513,14 +569,20 @@ class DnnlConvGrad : public DnnlKernel {
|
|||
std::shared_ptr<dnnl::memory> diff_weights_mem_;
|
||||
std::unique_ptr<dnnl::memory::desc> diff_weights_md_;
|
||||
TensorShape diff_weights_shape_;
|
||||
// memory used for reorders
|
||||
std::unique_ptr<dnnl::memory> diff_weights_reorder_mem_to_;
|
||||
|
||||
std::shared_ptr<dnnl::memory> diff_bias_mem_;
|
||||
std::unique_ptr<dnnl::memory::desc> diff_bias_md_;
|
||||
TensorShape diff_bias_shape_;
|
||||
// memory used for reorders
|
||||
std::unique_ptr<dnnl::memory> diff_bias_reorder_mem_to_;
|
||||
|
||||
IAllocatorUniquePtr<void> src_reorder_buffer_;
|
||||
IAllocatorUniquePtr<void> dst_reorder_buffer_;
|
||||
|
||||
bool gpu_available_;
|
||||
|
||||
private:
|
||||
Status ComputeKernelShape(const TensorShape& weight_shape, std::vector<int64_t>& kernel_shape) const {
|
||||
if (kernel_shape_specified_) {
|
||||
|
|
|
|||
|
|
@ -6,6 +6,33 @@
|
|||
#include "core/providers/dnnl/dnnl_execution_provider.h"
|
||||
#include "core/providers/dnnl/subgraph/dnnl_kernel.h"
|
||||
|
||||
|
||||
/*
|
||||
MaxPoolGrad: (According to OnnxRuntime discovered using code inspection and Onnx documentation)
|
||||
Inputs:
|
||||
0) dY - Gradient of output Y
|
||||
1) indices - indices
|
||||
Outputs:
|
||||
0) dX - Gradient of Input
|
||||
|
||||
+-----------------+
|
||||
(dY) diff_dst | |
|
||||
------------------->+ | (dX ) diff_src
|
||||
(indices) workspace | MaxPoolGrad +----------------->
|
||||
------------------->+ |
|
||||
| |
|
||||
+-----------------+
|
||||
|
||||
diff_dst = DNNL_ARG_DIFF_DST
|
||||
workspace = DNNL_ARG_WORKSPACE
|
||||
|
||||
diff_src = DNNL_ARG_DIFF_SRC
|
||||
|
||||
Attributes (auto_pad, dilations, group, kernel_shap, pads, and strides) should be the same as the forward pass Pool operator
|
||||
|
||||
The indices must come from the forward pool operator the indices input from OnnxRuntime will be ignored. For that reason the
|
||||
forward and backward operators must run using dnnl endpoint.
|
||||
*/
|
||||
namespace onnxruntime {
|
||||
namespace ort_dnnl {
|
||||
template <typename T>
|
||||
|
|
@ -30,12 +57,11 @@ class DnnlMaxPoolGrad : public DnnlKernel {
|
|||
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
|
||||
dnnl::engine cpu_engine;
|
||||
dnnl::engine engine_to_use;
|
||||
const auto iter = dnnl_engine.find(dnnl::engine::kind::cpu);
|
||||
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 = iter->second;
|
||||
engine_to_use = cpu_engine;
|
||||
}
|
||||
#if 0 // TODO update maxpoolgrad for gpu
|
||||
gpu_available_ = false;
|
||||
dnnl::engine gpu_engine;
|
||||
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
|
||||
|
|
@ -45,7 +71,7 @@ class DnnlMaxPoolGrad : public DnnlKernel {
|
|||
engine_to_use = gpu_engine;
|
||||
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
|
||||
}
|
||||
#endif
|
||||
|
||||
Ort::CustomOpApi ort{*api};
|
||||
int input_index = mklnode_ptr_->input_start_index < 0 ? 0 : mklnode_ptr_->input_start_index;
|
||||
|
||||
|
|
@ -54,7 +80,6 @@ class DnnlMaxPoolGrad : public DnnlKernel {
|
|||
const OrtValue* xgrad_input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
auto xgrad_tensor_info = ort.GetTensorTypeAndShape(xgrad_input_tensor);
|
||||
auto xgrad_tensor_shape = ort.GetTensorShape(xgrad_tensor_info);
|
||||
|
||||
ort.ReleaseTensorTypeAndShapeInfo(xgrad_tensor_info);
|
||||
auto xgradshape = xgrad_tensor_shape.data();
|
||||
auto xgraddim = xgrad_tensor_shape.size();
|
||||
|
|
@ -78,33 +103,12 @@ class DnnlMaxPoolGrad : public DnnlKernel {
|
|||
// TODO Sourcenode will set src of this node.
|
||||
x_shape_ = parents_[0].get()->primitive_dst_shape_;
|
||||
source_desc_ = parents_[0].get()->primitive_dst_desc_;
|
||||
dnnl::memory::dims src_dims_mkl(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
|
||||
|
||||
ort_source_format_ = parents_[0].get()->ort_source_format_;
|
||||
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 {
|
||||
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
dnnl::memory::desc(parents_[0].get()->primitive_dst_mem_->get_desc()));
|
||||
}
|
||||
}
|
||||
|
||||
//Get pointer to forward primitive descriptor created in maxpool forward
|
||||
fwd_primitive_desc_ = (pool_fwd_->GetPrimitiveDesc());
|
||||
|
||||
//The second input would be the workspace data populated during maxpool forward.
|
||||
//So just get the pointer to that data created during maxpool forward using the fwd desc
|
||||
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
dnnl::memory::desc(((fwd_primitive_desc_->workspace_desc()))));
|
||||
|
||||
src_mem_ = pool_fwd_->GetWorkspacePtr();
|
||||
|
||||
//Obtain output size and shape from the forward desc in maxpool.
|
||||
//This would be the input shape and size in the maxpool forward
|
||||
primitive_dst_shape_ = pool_fwd_->GetOutputShape();
|
||||
|
|
@ -144,82 +148,114 @@ class DnnlMaxPoolGrad : public DnnlKernel {
|
|||
dnnl::pooling_backward::desc(algo, *primitive_dst_md_, *diff_dst_md_,
|
||||
strides_mkl, kernel_mkl,
|
||||
padding_left_mkl, padding_right_mkl));
|
||||
|
||||
auto pool_fwd_pd = pool_fwd_->GetPrimitiveDesc();
|
||||
bwd_primitive_desc_ = onnxruntime::make_unique<dnnl::pooling_backward::primitive_desc>(
|
||||
dnnl::pooling_backward::primitive_desc(*bwd_desc_, engine_to_use, *(pool_fwd_->GetPrimitiveDesc())));
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
// Sub-graph's first node. Read input from input buffer
|
||||
diff_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_.get()->diff_dst_desc(), engine_to_use, nullptr));
|
||||
} else {
|
||||
// Sub-graph's inner node. set input to parent's output
|
||||
diff_dst_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
dnnl::pooling_backward::primitive_desc(*bwd_desc_, engine_to_use, *pool_fwd_pd));
|
||||
|
||||
primitive_src_desc_ = bwd_primitive_desc_.get()->diff_src_desc();
|
||||
primitive_dst_desc_ = bwd_primitive_desc_.get()->diff_dst_desc();
|
||||
|
||||
src_size_ = bwd_primitive_desc_.get()->diff_src_desc().get_size();
|
||||
dst_size_ = bwd_primitive_desc_.get()->diff_dst_desc().get_size();
|
||||
if (!gpu_available_) {
|
||||
// reorder source memory for best performance (AVX512);
|
||||
if (primitive_dst_desc_ != source_desc_) {
|
||||
dnnl::memory::desc pd(source_desc_);
|
||||
|
||||
// reorder source memory for best performance (AVX512);
|
||||
if (primitive_dst_desc_ != source_desc_) {
|
||||
dnnl::memory::dims src_dims(xgrad_shape_.GetDims().begin(), xgrad_shape_.GetDims().end());
|
||||
auto pd = dnnl::memory::desc(source_desc_);
|
||||
if (mklnode_ptr_->parent_nodes.empty())
|
||||
diff_dst_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(pd, cpu_engine, DNNL_MEMORY_NONE));
|
||||
else
|
||||
diff_dst_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty())
|
||||
diff_dst_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(pd, engine_to_use, nullptr));
|
||||
else
|
||||
diff_dst_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
|
||||
diff_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_->diff_dst_desc(), engine_to_use, nullptr));
|
||||
net.push_back(dnnl::reorder(*diff_dst_mem_from_, *diff_dst_mem_));
|
||||
net_args.push_back({{DNNL_ARG_FROM, *diff_dst_mem_from_},
|
||||
{DNNL_ARG_TO, *diff_dst_mem_}});
|
||||
} else {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
diff_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_->diff_dst_desc(), engine_to_use, nullptr));
|
||||
dnnl::memory(bwd_primitive_desc_->diff_dst_desc(), cpu_engine, DNNL_MEMORY_NONE));
|
||||
net.push_back(dnnl::reorder(*diff_dst_mem_from_, *diff_dst_mem_));
|
||||
net_args.push_back({{DNNL_ARG_FROM, *diff_dst_mem_from_},
|
||||
{DNNL_ARG_TO, *diff_dst_mem_}});
|
||||
} else {
|
||||
diff_dst_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
diff_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_->diff_dst_desc(), cpu_engine, DNNL_MEMORY_NONE));
|
||||
} else {
|
||||
diff_dst_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
}
|
||||
} else { //gpu_available_
|
||||
if (primitive_dst_desc_ != source_desc_) {
|
||||
//dnnl::memory::dims src_dims(xgrad_shape_.GetDims().begin(), xgrad_shape_.GetDims().end());
|
||||
dnnl::memory::desc pd(source_desc_);
|
||||
|
||||
if (mklnode_ptr_->parent_nodes.empty())
|
||||
diff_dst_mem_from_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(pd, cpu_engine, DNNL_MEMORY_NONE));
|
||||
else
|
||||
diff_dst_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
|
||||
diff_dst_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_->diff_dst_desc(), gpu_engine));
|
||||
net.push_back(dnnl::reorder(*diff_dst_mem_from_, *diff_dst_mem_gpu_));
|
||||
net_args.push_back({{DNNL_ARG_FROM, *diff_dst_mem_from_},
|
||||
{DNNL_ARG_TO, *diff_dst_mem_gpu_}});
|
||||
} else {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
// Sub-graph's first node. Read input from input buffer
|
||||
diff_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_->diff_dst_desc(), cpu_engine, DNNL_MEMORY_NONE));
|
||||
diff_dst_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_->diff_dst_desc(), gpu_engine));
|
||||
net.push_back(dnnl::reorder(*diff_dst_mem_, *diff_dst_mem_gpu_));
|
||||
net_args.push_back({{DNNL_ARG_SRC, *diff_dst_mem_},
|
||||
{DNNL_ARG_DST, *diff_dst_mem_gpu_}});
|
||||
} else {
|
||||
diff_dst_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
|
||||
primitive_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(bwd_primitive_desc_.get()->diff_src_desc(), engine_to_use));
|
||||
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(bwd_primitive_desc_.get()->diff_src_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(bwd_primitive_desc_.get()->diff_src_desc(), cpu_engine, DNNL_MEMORY_NONE));
|
||||
}
|
||||
} else {
|
||||
// Last node but re-order not needed. Allocate buffer to output of this node
|
||||
// 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(bwd_primitive_desc_.get()->diff_src_desc(), engine_to_use, nullptr));
|
||||
dnnl::memory(bwd_primitive_desc_.get()->diff_src_desc(), cpu_engine));
|
||||
}
|
||||
} 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(bwd_primitive_desc_.get()->diff_src_desc(), engine_to_use));
|
||||
dnnl::memory(bwd_primitive_desc_.get()->diff_src_desc(), gpu_engine));
|
||||
}
|
||||
|
||||
auto workspace_mem_ = pool_fwd_->GetWorkspacePtr();
|
||||
|
||||
pool_bwd_ = onnxruntime::make_unique<dnnl::pooling_backward>(
|
||||
dnnl::pooling_backward(*bwd_primitive_desc_));
|
||||
|
||||
net.push_back(*pool_bwd_);
|
||||
|
||||
net_args.push_back({{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_},
|
||||
{DNNL_ARG_WORKSPACE, *src_mem_}});
|
||||
if (!gpu_available_) {
|
||||
net_args.push_back({{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_},
|
||||
{DNNL_ARG_WORKSPACE, *workspace_mem_}});
|
||||
} else {
|
||||
net_args.push_back({{DNNL_ARG_DIFF_DST, *diff_dst_mem_gpu_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_},
|
||||
{DNNL_ARG_WORKSPACE, *workspace_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(engine_to_use, t, net, net_args);
|
||||
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -237,10 +273,11 @@ class DnnlMaxPoolGrad : public DnnlKernel {
|
|||
} else {
|
||||
diff_dst_mem_from_ = parents_[0].get()->primitive_dst_mem_;
|
||||
}
|
||||
|
||||
auto diff_dst_size = bwd_primitive_desc_.get()->diff_dst_desc().get_size();
|
||||
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, diff_dst_size);
|
||||
diff_dst_mem_->set_data_handle(src_reorder_buffer_.get());
|
||||
if (!gpu_available_) {
|
||||
auto diff_dst_size = bwd_primitive_desc_.get()->diff_dst_desc().get_size();
|
||||
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, diff_dst_size);
|
||||
diff_dst_mem_->set_data_handle(src_reorder_buffer_.get());
|
||||
}
|
||||
} else {
|
||||
if (mklnode_ptr_->parent_nodes.empty()) {
|
||||
const OrtValue* dx_input_tensor = ort.KernelContext_GetInput(context, input_index);
|
||||
|
|
@ -259,10 +296,14 @@ class DnnlMaxPoolGrad : 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_) {
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
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 {
|
||||
primitive_dst_mem_->set_data_handle(dst_data);
|
||||
reorder_dst_mem_to_->set_data_handle(dst_data);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
|
|
@ -334,38 +375,36 @@ class DnnlMaxPoolGrad : public DnnlKernel {
|
|||
}
|
||||
|
||||
private:
|
||||
size_t src_size_;
|
||||
size_t dst_size_;
|
||||
|
||||
std::shared_ptr<dnnl::memory> src_mem_;
|
||||
std::shared_ptr<dnnl::memory> diff_dst_mem_;
|
||||
std::shared_ptr<dnnl::memory> diff_dst_mem_gpu_;
|
||||
|
||||
std::unique_ptr<dnnl::memory::desc> src_md_;
|
||||
std::unique_ptr<dnnl::memory::desc> diff_dst_md_;
|
||||
|
||||
std::unique_ptr<dnnl::pooling_forward::desc> fwd_desc_;
|
||||
std::shared_ptr<dnnl::pooling_forward::primitive_desc> fwd_primitive_desc_;
|
||||
|
||||
std::shared_ptr<DnnlPool<T>> pool_fwd_;
|
||||
|
||||
std::unique_ptr<dnnl::pooling_backward::desc> bwd_desc_;
|
||||
std::unique_ptr<dnnl::pooling_backward::primitive_desc> bwd_primitive_desc_;
|
||||
std::unique_ptr<dnnl::primitive> pool_bwd_;
|
||||
|
||||
std::shared_ptr<dnnl::memory> src_mem_from_;
|
||||
std::unique_ptr<dnnl::memory> src_mem_to_;
|
||||
|
||||
std::shared_ptr<dnnl::memory> diff_dst_mem_from_;
|
||||
std::unique_ptr<dnnl::memory> diff_dst_mem_to_;
|
||||
|
||||
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;
|
||||
bool is_1D = src_dims_mkl.size() == 3 ? true : false;
|
||||
bool is_2D = src_dims_mkl.size() == 4;
|
||||
bool is_1D = src_dims_mkl.size() == 3;
|
||||
dnnl::memory::format_tag fmt = dnnl::memory::format_tag::any;
|
||||
if (CPUIDInfo::GetCPUIDInfo().HasAVX512f()) {
|
||||
if (gpu_available_) {
|
||||
if (is_1D)
|
||||
fmt = dnnl::memory::format_tag::ncw;
|
||||
else
|
||||
fmt = is_2D ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::ncdhw;
|
||||
} else if (CPUIDInfo::GetCPUIDInfo().HasAVX512f()) {
|
||||
if (is_1D)
|
||||
fmt = dnnl::memory::format_tag::nCw16c;
|
||||
else
|
||||
|
|
|
|||
|
|
@ -69,7 +69,6 @@ class DnnlPool : public DnnlKernel {
|
|||
// TODO Sourcenode will set src of this node.
|
||||
x_shape_ = parents_[0].get()->primitive_dst_shape_;
|
||||
source_desc_ = parents_[0].get()->primitive_dst_desc_;
|
||||
dnnl::memory::dims src_dims_mkl(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
|
||||
|
||||
ort_source_format_ = parents_[0].get()->ort_source_format_;
|
||||
ort_source_desc_ = parents_[0].get()->ort_source_desc_;
|
||||
|
|
@ -77,6 +76,7 @@ class DnnlPool : public DnnlKernel {
|
|||
|
||||
if (!gpu_available_) {
|
||||
if (source_desc_ == ort_source_desc_) {
|
||||
dnnl::memory::dims src_dims_mkl(x_shape_.GetDims().begin(), x_shape_.GetDims().end());
|
||||
// reorder for better performance
|
||||
dnnl::memory::format_tag fmt = GetAVXFormat(src_dims_mkl);
|
||||
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
|
|
@ -139,37 +139,16 @@ class DnnlPool : public DnnlKernel {
|
|||
strides_mkl, kernel_mkl,
|
||||
padding_left_mkl, padding_right_mkl));
|
||||
|
||||
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));
|
||||
//}
|
||||
fwd_primitive_desc_ = onnxruntime::make_unique<dnnl::pooling_forward::primitive_desc>(
|
||||
dnnl::pooling_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();
|
||||
|
||||
src_size_ = fwd_primitive_desc_.get()->src_desc().get_size();
|
||||
dst_size_ = fwd_primitive_desc_.get()->dst_desc().get_size();
|
||||
|
||||
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())
|
||||
|
|
@ -193,7 +172,6 @@ class DnnlPool : public DnnlKernel {
|
|||
}
|
||||
} 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()) {
|
||||
|
|
@ -248,22 +226,24 @@ class DnnlPool : public DnnlKernel {
|
|||
|
||||
#ifdef ENABLE_TRAINING
|
||||
workspace_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory(fwd_primitive_desc_.get()->workspace_desc(), cpu_engine));
|
||||
dnnl::memory(fwd_primitive_desc_.get()->workspace_desc(), engine_to_use));
|
||||
#endif // ENABLE_TRAINING
|
||||
|
||||
pool_fwd_ = onnxruntime::make_unique<dnnl::pooling_forward>(
|
||||
dnnl::pooling_forward(*fwd_primitive_desc_));
|
||||
|
||||
net.push_back(*pool_fwd_);
|
||||
if (!gpu_available_) {
|
||||
net.push_back(*pool_fwd_);
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
|
||||
#ifdef ENABLE_TRAINING
|
||||
{DNNL_ARG_WORKSPACE, *workspace_mem_},
|
||||
#endif // ENABLE_TRAINING
|
||||
{DNNL_ARG_WORKSPACE, *workspace_mem_},
|
||||
#endif //ENABLE_TRAINING
|
||||
{DNNL_ARG_DST, *primitive_dst_mem_}});
|
||||
} else { // gpu_available_
|
||||
net.push_back(*pool_fwd_);
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
|
||||
#ifdef ENABLE_TRAINING
|
||||
{DNNL_ARG_WORKSPACE, *workspace_mem_},
|
||||
#endif //ENABLE_TRAINING
|
||||
{DNNL_ARG_DST, *primitive_dst_mem_}});
|
||||
}
|
||||
if (mklnode_ptr_->output_index >= 0) {
|
||||
|
|
@ -409,14 +389,12 @@ class DnnlPool : public DnnlKernel {
|
|||
}
|
||||
|
||||
private:
|
||||
size_t src_size_;
|
||||
size_t dst_size_;
|
||||
|
||||
std::shared_ptr<dnnl::memory> src_mem_;
|
||||
std::shared_ptr<dnnl::memory> src_mem_gpu_;
|
||||
#ifdef ENABLE_TRAINING
|
||||
std::shared_ptr<dnnl::memory> workspace_mem_;
|
||||
#endif // ENABLE_TRAINING
|
||||
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_;
|
||||
|
|
@ -443,14 +421,21 @@ class DnnlPool : public DnnlKernel {
|
|||
bool is_1D = src_dims_mkl.size() == 3 ? true : false;
|
||||
#endif
|
||||
dnnl::memory::format_tag fmt = dnnl::memory::format_tag::any;
|
||||
if (!gpu_available_ && CPUIDInfo::GetCPUIDInfo().HasAVX512f()) {
|
||||
if (gpu_available_) {
|
||||
#ifdef ENABLE_TRAINING
|
||||
if (is_1D)
|
||||
fmt = dnnl::memory::format_tag::ncw;
|
||||
else
|
||||
#endif
|
||||
fmt = is_2D ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::ncdhw;
|
||||
} else if (CPUIDInfo::GetCPUIDInfo().HasAVX512f()) {
|
||||
#ifdef ENABLE_TRAINING
|
||||
if (is_1D)
|
||||
fmt = dnnl::memory::format_tag::nCw16c;
|
||||
else
|
||||
#endif
|
||||
fmt = is_2D ? dnnl::memory::format_tag::nChw16c : dnnl::memory::format_tag::nCdhw16c;
|
||||
} else if (!gpu_available_ && CPUIDInfo::GetCPUIDInfo().HasAVX2() && (src_dims_mkl[1] % 8 == 0)) {
|
||||
} else if (CPUIDInfo::GetCPUIDInfo().HasAVX2() && (src_dims_mkl[1] % 8 == 0)) {
|
||||
#ifdef ENABLE_TRAINING
|
||||
if (is_1D)
|
||||
fmt = dnnl::memory::format_tag::nCw8c;
|
||||
|
|
|
|||
|
|
@ -57,12 +57,11 @@ class DnnlReluGrad : public DnnlKernel {
|
|||
std::vector<std::unordered_map<int, dnnl::memory>>& net_args) override {
|
||||
dnnl::engine cpu_engine;
|
||||
dnnl::engine engine_to_use;
|
||||
const auto iter = dnnl_engine.find(dnnl::engine::kind::cpu);
|
||||
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 = iter->second;
|
||||
engine_to_use = cpu_engine;
|
||||
}
|
||||
#if 0 // TODO update relugrad for gpu
|
||||
gpu_available_ = false;
|
||||
dnnl::engine gpu_engine;
|
||||
iter = dnnl_engine.find(dnnl::engine::kind::gpu);
|
||||
|
|
@ -72,10 +71,11 @@ class DnnlReluGrad : public DnnlKernel {
|
|||
engine_to_use = gpu_engine;
|
||||
LOGS_DEFAULT(INFO) << "gpu engine found" << std::endl;
|
||||
}
|
||||
#endif
|
||||
if (!relu_fwd_){
|
||||
ORT_THROW("Unable to find forward pass Relu opt. "
|
||||
"Verify AddForwardDnnlKernel was called after DnnlReluGrad was created.");
|
||||
|
||||
if (!relu_fwd_) {
|
||||
ORT_THROW(
|
||||
"Unable to find forward pass Relu opt. "
|
||||
"Verify AddForwardDnnlKernel was called after DnnlReluGrad was created.");
|
||||
}
|
||||
|
||||
Ort::CustomOpApi ort{*api};
|
||||
|
|
@ -115,7 +115,7 @@ class DnnlReluGrad : public DnnlKernel {
|
|||
diff_dst_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
dnnl::memory::desc({xgrad_dims}, DnnnType<T>(), ort_source_format_));
|
||||
diff_dst_mem_ = onnxruntime::make_unique<dnnl::memory>(
|
||||
dnnl::memory({{xgrad_dims}, DnnnType<T>(), ort_source_format_}, engine_to_use, nullptr));
|
||||
dnnl::memory({{xgrad_dims}, DnnnType<T>(), ort_source_format_}, cpu_engine, nullptr));
|
||||
|
||||
// convert the onnx Tensor to dnnl::memory and dnnl::memory::desc
|
||||
const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_index + 1);
|
||||
|
|
@ -146,7 +146,17 @@ class DnnlReluGrad : public DnnlKernel {
|
|||
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
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_}, engine_to_use, nullptr));
|
||||
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_}});
|
||||
diff_dst_mem_gpu_ = onnxruntime::make_unique<dnnl::memory>(*diff_dst_md_, gpu_engine);
|
||||
net.push_back(mkldnn::reorder(*diff_dst_mem_, *diff_dst_mem_gpu_));
|
||||
net_args.push_back({{MKLDNN_ARG_SRC, *diff_dst_mem_},
|
||||
{MKLDNN_ARG_DST, *diff_dst_mem_gpu_}});
|
||||
}
|
||||
} else {
|
||||
ort_source_format_ = parents_[0].get()->ort_source_format_;
|
||||
ort_source_desc_ = parents_[0].get()->ort_source_desc_;
|
||||
|
|
@ -154,12 +164,18 @@ class DnnlReluGrad : public DnnlKernel {
|
|||
|
||||
diff_dst_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
dnnl::memory::desc(parents_[0].get()->primitive_dst_desc_));
|
||||
diff_dst_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
xgrad_shape = parents_[0].get()->primitive_dst_shape_;
|
||||
|
||||
src_md_ = onnxruntime::make_unique<dnnl::memory::desc>(
|
||||
dnnl::memory::desc(parents_[1].get()->primitive_dst_desc_));
|
||||
src_mem_ = parents_[1].get()->primitive_dst_mem_;
|
||||
|
||||
if (!gpu_available_) {
|
||||
diff_dst_mem_ = parents_[0].get()->primitive_dst_mem_;
|
||||
src_mem_ = parents_[1].get()->primitive_dst_mem_;
|
||||
} else {
|
||||
diff_dst_mem_gpu_ = parents_[0].get()->primitive_dst_mem_;
|
||||
src_mem_gpu_ = parents_[1].get()->primitive_dst_mem_;
|
||||
}
|
||||
|
||||
xgrad_shape = parents_[0].get()->primitive_dst_shape_;
|
||||
x_shape = parents_[1].get()->primitive_dst_shape_;
|
||||
}
|
||||
|
||||
|
|
@ -169,45 +185,55 @@ class DnnlReluGrad : public DnnlKernel {
|
|||
dnnl::algorithm algo = dnnl::algorithm::eltwise_relu;
|
||||
|
||||
relu_bwd_desc_ = onnxruntime::make_unique<dnnl::eltwise_backward::desc>(
|
||||
dnnl::eltwise_backward::desc(algo, *diff_dst_md_,*src_md_, 0.0, 0.0));
|
||||
dnnl::eltwise_backward::desc(algo, *diff_dst_md_, *src_md_, 0.0, 0.0));
|
||||
relu_bwd_pd_ = onnxruntime::make_unique<dnnl::eltwise_backward::primitive_desc>(
|
||||
dnnl::eltwise_backward::primitive_desc(*relu_bwd_desc_, engine_to_use, *(relu_fwd_->GetPrimitiveDesc())));
|
||||
|
||||
primitive_src_desc_ = relu_bwd_pd_.get()->src_desc();
|
||||
primitive_dst_desc_ = relu_bwd_pd_.get()->diff_src_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_bwd_pd_.get()->diff_src_desc(), engine_to_use));
|
||||
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_bwd_pd_.get()->diff_src_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_bwd_pd_.get()->diff_src_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_bwd_pd_.get()->diff_src_desc(), engine_to_use, 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_bwd_pd_.get()->diff_src_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_bwd_pd_.get()->diff_src_desc(), engine_to_use));
|
||||
primitive_dst_mem_ = std::make_shared<dnnl::memory>(dnnl::memory(relu_bwd_pd_.get()->diff_src_desc(), gpu_engine));
|
||||
}
|
||||
|
||||
net.push_back(dnnl::eltwise_backward(*relu_bwd_pd_));
|
||||
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
|
||||
{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_}});
|
||||
if (!gpu_available_) {
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_},
|
||||
{DNNL_ARG_DIFF_DST, *diff_dst_mem_},
|
||||
{DNNL_ARG_DIFF_SRC, *primitive_dst_mem_}});
|
||||
} else {
|
||||
net_args.push_back({{DNNL_ARG_SRC, *src_mem_gpu_},
|
||||
{DNNL_ARG_DIFF_DST, *diff_dst_mem_gpu_},
|
||||
{DNNL_ARG_DIFF_SRC, *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(engine_to_use, t, net, net_args);
|
||||
InitDstReorderOutput(cpu_engine, t, net, net_args, gpu_available_);
|
||||
}
|
||||
}
|
||||
|
||||
virtual Status Bind(const OrtCustomOpApi* api,
|
||||
OrtKernelContext* context) {
|
||||
OrtKernelContext* context) {
|
||||
Ort::CustomOpApi ort{*api};
|
||||
|
||||
ORT_RETURN_IF_ERROR(primitive_created_status_);
|
||||
|
|
@ -229,25 +255,35 @@ class DnnlReluGrad : 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();
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
std::shared_ptr<DnnlRelu<T>> relu_fwd_;
|
||||
std::shared_ptr<dnnl::memory> diff_dst_mem_;
|
||||
std::shared_ptr<dnnl::memory> diff_dst_mem_gpu_;
|
||||
|
||||
std::unique_ptr<dnnl::memory::desc> diff_dst_md_;
|
||||
|
||||
std::shared_ptr<dnnl::memory> src_mem_;
|
||||
std::shared_ptr<dnnl::memory> src_mem_gpu_;
|
||||
|
||||
std::unique_ptr<dnnl::memory::desc> src_md_;
|
||||
|
||||
std::unique_ptr<dnnl::eltwise_backward::desc> relu_bwd_desc_;
|
||||
std::unique_ptr<dnnl::eltwise_backward::primitive_desc> relu_bwd_pd_;
|
||||
|
||||
bool gpu_available_;
|
||||
}; // namespace ort_dnnl
|
||||
} // namespace ort_dnnl
|
||||
} // namespace onnxruntime
|
||||
Loading…
Reference in a new issue