From da3bd45cdd50782a22fa919a343101683468002d Mon Sep 17 00:00:00 2001
From: Scott McKay <skottmckay@gmail.com>
Date: Fri, 20 Sep 2024 13:44:47 +1000
Subject: [PATCH] Fix CUDA reduction ops handling of optional axes input
 (#22149)

### Description
<!-- Describe your changes. -->
The optional `axes` input may exist with an empty name and be a nullptr.

Update the CUDA implementation to handle this.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

#22035
---
 .../builders/impl/reduction_op_builder.cc     | 19 ++++++++----
 .../providers/cuda/reduction/reduction_ops.cc |  5 ++--
 .../cpu/reduction/reduction_ops_test.cc       | 29 +++++++++++++++++++
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc
index 32378b1f654d8..5651b9cc5793e 100644
--- a/onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc
+++ b/onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc
@@ -89,7 +89,15 @@ bool ReductionOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInpu
                                            const logging::Logger& logger) const {
   const auto& input_defs = node.InputDefs();
 
+  NodeAttrHelper helper(node);
+
+  // noop_with_empty_axes defaults to false and is only available in newer opsets where 'axes' is an optional input
+  // so we don't need to check the 'axes' attribute used in older opsets here.
+  const bool noop_with_empty_axes = helper.Get("noop_with_empty_axes", 0) != 0;
+  bool empty_axes = true;
+
   if (input_defs.size() > 1 && input_defs[1]->Exists()) {
+    // 'axes' is optional input in new opsets
     const auto& axes_name = input_defs[1]->Name();
     const auto& initializers = input_params.graph_viewer.GetAllInitializedTensors();
     if (!Contains(initializers, axes_name)) {
@@ -97,12 +105,13 @@ bool ReductionOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInpu
       return false;
     }
 
-    NodeAttrHelper helper(node);
+    empty_axes = initializers.at(axes_name)->int64_data_size() == 0;
+  }
 
-    if (initializers.at(axes_name)->int64_data_size() == 0 && helper.Get("noop_with_empty_axes", 0) != 0) {
-      LOGS(logger, VERBOSE) << "CoreML doesn't support noop on empty axes for reduction layers" << std::endl;
-      return false;
-    }
+  if (empty_axes && noop_with_empty_axes) {
+    // TODO: When we add ML Program support we should enable this as it makes the node an Identity op
+    LOGS(logger, VERBOSE) << "CoreML doesn't support noop on empty axes for reduction layers" << std::endl;
+    return false;
   }
 
   return true;
diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
index c921339ee6f33..860bea67dc719 100644
--- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
+++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
@@ -645,10 +645,9 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
   TensorShapeVector axes;
 
   size_t num_inputs = ctx->InputCount();
-  if (num_inputs == 2) {
+  const Tensor* axes_tensor = num_inputs == 2 ? ctx->Input<Tensor>(1) : nullptr;  // optional input. may be nullptr.
+  if (axes_tensor != nullptr) {
     // override the attribute value with the input value for reduction_axes
-    const Tensor* axes_tensor = ctx->Input<Tensor>(1);
-    ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null");
     ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor.");
     auto nDims = static_cast<size_t>(axes_tensor->Shape()[0]);
     const auto* data = axes_tensor->Data<int64_t>();
diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
index ab24337046b99..0697187a777d6 100644
--- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
+++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
@@ -2780,6 +2780,22 @@ TEST(ReductionOpTest, ReduceSum_empty_axes_input_initializer_opset_13) {
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider});
 }
 
+TEST(ReductionOpTest, ReduceSum_missing_axes_input_noop_opset_13) {
+  OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain);
+  test.AddAttribute("keepdims", (int64_t)0);
+  test.AddAttribute("noop_with_empty_axes", (int64_t)1);  // missing axes input and noop. should be identity
+  test.AddInput<float>("data", {1, 2, 2},
+                       {1.0f, 2.0f,
+                        3.0f, 4.0f});
+  test.AddOptionalInputEdge<int64_t>();  // missing optional axes input
+  test.AddOutput<float>("reduced", {1, 2, 2},
+                        {1.0f, 2.0f,
+                         3.0f, 4.0f});
+
+  // TODO: OpenVINO doesn't support "axes" input in opset 13
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider});
+}
+
 TEST(ReductionOpTest, ReduceSum0DTensor) {
   OpTester test("ReduceSum");
   test.AddInput<float>("data", {}, {2});
@@ -5940,5 +5956,18 @@ TEST(ReductionOpTest, empty_set_ReduceSumSquare) {
 TEST(ReductionOpTest, empty_set_ReduceSumSquare_13) {
   test_empty_set("ReduceSumSquare", 13, false, 0.0f);
 }
+
+TEST(ReductionOpTest, MissingOptionalAxes) {
+  OpTester test("ReduceMax", 18);
+  test.AddInput<float>("data", {2, 2},
+                       {1.0f, 4.0f,
+                        3.0f, 2.0f});
+  test.AddOptionalInputEdge<int64_t>();
+  test.AddOutput<float>("reduced", {1, 1}, {4.0f});
+
+  // OpenVINO doesn't support "axes" input
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider});
+}
+
 }  // namespace test
 }  // namespace onnxruntime