From 267fb898e33c263c6337168556546aa5d6bc3632 Mon Sep 17 00:00:00 2001 From: chethanpk <63478277+chethanpk@users.noreply.github.com> Date: Tue, 21 Sep 2021 23:30:58 -0700 Subject: [PATCH] Added code to support Softmaxgrad for DNNL EP (#9022) * Added code to support Softmaxgrad Signed-off-by: Chethan Palangotu Keshava * Bringing back the opset filters for softmax that I had removed. This will fix the test failures from onnnx repo. Signed-off-by: Chethan Palangotu Keshava --- .../providers/dnnl/dnnl_node_capability.cc | 1 + .../core/providers/dnnl/dnnl_op_manager.cc | 1 + .../providers/dnnl/subgraph/dnnl_softmax.cc | 2 +- .../dnnl/subgraph/dnnl_softmaxgrad.cc | 55 +++++++++++++++++++ .../dnnl/subgraph/dnnl_softmaxgrad.h | 28 ++++++++++ .../dnnl/subgraph/dnnl_subgraph_primitive.cc | 3 + .../test/providers/cpu/math/softmax_test.cc | 4 +- 7 files changed, 91 insertions(+), 3 deletions(-) create mode 100644 onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.cc create mode 100644 onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.h diff --git a/onnxruntime/core/providers/dnnl/dnnl_node_capability.cc b/onnxruntime/core/providers/dnnl/dnnl_node_capability.cc index 79e440c954..ac1bd52e66 100644 --- a/onnxruntime/core/providers/dnnl/dnnl_node_capability.cc +++ b/onnxruntime/core/providers/dnnl/dnnl_node_capability.cc @@ -201,6 +201,7 @@ bool DnnlSoftmaxNodeCapability::IsAttributeSupported(const Node* node) const { return true; } + // DnnlMatMulNodeCapability class //------------------------------------- bool DnnlMatMulNodeCapability::Supported(const Node* node) const { diff --git a/onnxruntime/core/providers/dnnl/dnnl_op_manager.cc b/onnxruntime/core/providers/dnnl/dnnl_op_manager.cc index d8b56ba789..26a524050e 100644 --- a/onnxruntime/core/providers/dnnl/dnnl_op_manager.cc +++ b/onnxruntime/core/providers/dnnl/dnnl_op_manager.cc @@ -38,6 +38,7 @@ DnnlOpManager::DnnlOpManager() { dnnl_ops_map_.emplace(std::make_pair("ConvGrad", std::unique_ptr(new DnnlDefaultNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("MaxPoolGrad", std::unique_ptr(new DnnlPoolNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("ReluGrad", std::unique_ptr(new DnnlDefaultNodeCapability()))); + dnnl_ops_map_.emplace(std::make_pair("SoftmaxGrad", std::unique_ptr(new DnnlSoftmaxNodeCapability()))); #endif // ENABLE_TRAINING } diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmax.cc b/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmax.cc index ff96ac438e..3e5d1bf5de 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmax.cc +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmax.cc @@ -42,7 +42,7 @@ void DnnlSoftmax::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) { int64_t DnnlSoftmax::ReadAxis(DnnlNode& node) { auto attr = node.Attributes().find("axis"); - int64_t axis = 1; //Default value according to ONNX spec + int64_t axis = -1; //Default value according to ONNX spec 13 but works with lower opset too if (attr != node.Attributes().end() && attr->second().type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_INT) { axis = attr->second().i(); diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.cc b/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.cc new file mode 100644 index 0000000000..8c0adc9f2c --- /dev/null +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.cc @@ -0,0 +1,55 @@ +#include "dnnl_softmaxgrad.h" +#include "dnnl_subgraph.h" +#include "dnnl_subgraph_primitive.h" + +namespace onnxruntime { +namespace ort_dnnl { + +DnnlSoftmaxGrad::DnnlSoftmaxGrad() {} + +void DnnlSoftmaxGrad::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) { + auto eng = sp.GetEngine(); + + //get what's available as input + auto src_mem = sp.GetMemory(node.Input(IN_X).Name()); + auto diff_dst_mem = sp.GetMemory(node.Input(IN_dY).Name()); + + //reorder if needed (gpu) + auto softmax_bwd_src_mem = sp.GetMemoryAndReshape(node.Input(IN_X), src_mem.get_desc(), eng); + auto softmax_bwd_diff_dst_mem = sp.GetMemoryAndReshape(node.Input(IN_dY), diff_dst_mem.get_desc(), eng); + + auto axis = ReadAxis(node); + + if (axis < 0) + axis = src_mem.get_desc().dims().size() + axis; + + //create hints on the fly + auto hints_d = dnnl::softmax_forward::desc(dnnl::prop_kind::forward_training, softmax_bwd_src_mem.get_desc(), (int) axis); + auto hints_pd = dnnl::softmax_forward::primitive_desc(hints_d, eng); + + auto softmax_bwd_d = dnnl::softmax_backward::desc(softmax_bwd_diff_dst_mem.get_desc(), softmax_bwd_src_mem.get_desc(), (int) axis); + + auto softmax_bwd_pd = dnnl::softmax_backward::primitive_desc(softmax_bwd_d, eng, hints_pd); + + auto softmax_bwd_diff_src_mem = dnnl::memory(softmax_bwd_pd.diff_src_desc(), eng); + + auto softmax_bwd = dnnl::softmax_backward(softmax_bwd_pd); + + sp.AddPrimitive(softmax_bwd, {{DNNL_ARG_DST, softmax_bwd_src_mem}, + {DNNL_ARG_DIFF_DST, softmax_bwd_diff_dst_mem}, + {DNNL_ARG_DIFF_SRC, softmax_bwd_diff_src_mem}}); + + sp.SetMemory(node.Output(OUT_dX), softmax_bwd_diff_src_mem); +} + +int64_t DnnlSoftmaxGrad::ReadAxis(DnnlNode& node) { + auto attr = node.Attributes().find("axis"); + int64_t axis = -1; //Default value according to ONNX spec 13 but works with lower opsets too + if (attr != node.Attributes().end() && + attr->second().type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_INT) { + axis = attr->second().i(); + } + return axis; +} +} // namespace ort_dnnl +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.h new file mode 100644 index 0000000000..649c910973 --- /dev/null +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_softmaxgrad.h @@ -0,0 +1,28 @@ +// Copyright(C) 2021 Intel Corporation +// Licensed under the MIT License + +#pragma once +#include "dnnl_subgraph.h" +#include "dnnl_subgraph_primitive.h" + +namespace onnxruntime { +namespace ort_dnnl { + +class DnnlSoftmaxGrad { + public: + enum InputTensors : int { + IN_dY = 0, + IN_X = 1 + }; + + enum OutputTensors : int { + OUT_dX = 0 + }; + + DnnlSoftmaxGrad(); + void CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node); + int64_t ReadAxis(DnnlNode& node); +}; + +} // namespace ort_dnnl +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc b/onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc index 6432ef9632..2d8385b595 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc @@ -14,6 +14,7 @@ #include "dnnl_pool.h" #include "dnnl_reducemean.h" #include "dnnl_softmax.h" +#include "dnnl_softmaxgrad.h" #include "dnnl_sum.h" #if defined(ENABLE_TRAINING) @@ -73,6 +74,8 @@ void DnnlSubgraphPrimitive::AddKernels() { DnnlConvGrad().CreatePrimitive(*this, node); } else if (node.OpType() == "ReluGrad") { DnnlReluGrad().CreatePrimitive(*this, node); + } else if (node.OpType() == "SoftmaxGrad") { + DnnlSoftmaxGrad().CreatePrimitive(*this, node); #endif } else { throw std::invalid_argument("Kernel not found"); diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 197dead786..d4c7d0fcf1 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -98,7 +98,7 @@ TEST(SoftmaxOperator, ThreeDimsAxis0) { 0.017545262f, 0.0135920765f, 0.027506188f, 0.010684152f, 0.0049549243f, 0.01401341f, 0.011721271f, 0.027815264f, 0.021463264f, 0.014014485f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 0, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // Axis=0 is not supported by TensorRT + RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 0, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); // Axis=0 is not supported by TensorRT } TEST(SoftmaxOperator, ThreeDimsAxis1) { @@ -124,7 +124,7 @@ TEST(SoftmaxOperator, ThreeDimsAxis1) { 0.050680935f, 0.03926183f, 0.079453886f, 0.030862054f, 0.014312706f, 0.040478885f, 0.033857856f, 0.080346674f, 0.06199841f, 0.040481992f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 1, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); + RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 1, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); } TEST(SoftmaxOperator, ThreeDimsAxis1_opset13) {