Added code to support Softmaxgrad for DNNL EP (#9022)

* Added code to support Softmaxgrad

Signed-off-by: Chethan Palangotu Keshava <chethan.palangotu.keshava@intel.com>

* 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 <chethan.palangotu.keshava@intel.com>

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 {

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<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("MaxPoolGrad", std::unique_ptr<DnnlNodeCapability>(new DnnlPoolNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("ReluGrad", std::unique_ptr<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
+  dnnl_ops_map_.emplace(std::make_pair("SoftmaxGrad", std::unique_ptr<DnnlNodeCapability>(new DnnlSoftmaxNodeCapability())));
 #endif  // ENABLE_TRAINING
 }
 

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();

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

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

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");

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) {