Skip to content

Commit

Permalink
Fix CUDA reduction ops handling of optional axes input (#22149)
Browse files Browse the repository at this point in the history
### 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
  • Loading branch information
skottmckay authored Sep 20, 2024
1 parent f3cbe76 commit da3bd45
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -89,20 +89,29 @@ 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)) {
LOGS(logger, VERBOSE) << "Axes of reduction must be a constant initializer";
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;
Expand Down
5 changes: 2 additions & 3 deletions onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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>();
Expand Down
29 changes: 29 additions & 0 deletions onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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});
Expand Down Expand Up @@ -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

0 comments on commit da3bd45

Please sign in to comment.