mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
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
This commit is contained in:
parent
f3cbe76059
commit
da3bd45cdd
3 changed files with 45 additions and 8 deletions
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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>();
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue