diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc index 3525f7b009..d2240b5d50 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc @@ -33,6 +33,7 @@ void Selectors::RegisterSelector(const OpVersionsAndSelector::OpVersionsMap& ops // output Q have the same scale and zero_point. static const OpVersionsAndSelector::OpVersionsMap GetMiscOpVersionsMap() { return {{"Gather", {}}, + {"GatherElements", {}}, {"Reshape", {}}, {"Expand", {}}, {"Flatten", {}}, diff --git a/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc b/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc index 8c34a7a60e..dd5c6a5a79 100644 --- a/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc +++ b/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc @@ -110,6 +110,7 @@ OpBuilderRegistrations::OpBuilderRegistrations() { { CreateGatherOpBuilder("Gather", *this); + CreateGatherOpBuilder("GatherElements", *this); } { diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h index 6886845ff3..8a953616e2 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h @@ -122,6 +122,7 @@ class BaseOpBuilder : public IOpBuilder { {"Exp", QNN_OP_ELEMENT_WISE_EXP}, {"Floor", QNN_OP_ELEMENT_WISE_FLOOR}, {"Gather", QNN_OP_GATHER}, + {"GatherElements", QNN_OP_GATHER_ELEMENTS}, {"Greater", QNN_OP_ELEMENT_WISE_GREATER}, {"GreaterOrEqual", QNN_OP_ELEMENT_WISE_GREATER_EQUAL}, {"Less", QNN_OP_ELEMENT_WISE_LESS}, diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc index 1b759af256..3737fcb54f 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include "core/providers/common.h" #include "core/providers/shared/utils/utils.h" #include "core/providers/qnn/builder/qnn_model_wrapper.h" @@ -13,11 +14,16 @@ namespace onnxruntime { namespace qnn { +// Handles Gather and GatherElements class GatherOpBuilder : public BaseOpBuilder { public: GatherOpBuilder() : BaseOpBuilder("GatherOpBuilder") {} ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(GatherOpBuilder); + Status IsOpSupported(QnnModelWrapper& qnn_model_wrapper, + const NodeUnit& node_unit, + const logging::Logger& logger) const override ORT_MUST_USE_RESULT; + protected: Status ProcessInputs(QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit, @@ -32,79 +38,154 @@ class GatherOpBuilder : public BaseOpBuilder { bool do_op_validation) const override ORT_MUST_USE_RESULT; }; -Status GatherOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, +Status GatherOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit, - const logging::Logger& logger, - std::vector& input_names, - bool do_op_validation) const { - const auto& inputs = node_unit.Inputs(); - ORT_RETURN_IF(inputs.size() != 2, "Gather should has 2 inputs at least!"); - ORT_RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, inputs[0], logger, input_names)); + const logging::Logger& logger) const { + // On QNN CPU backend, the QNN validator does not properly reject unsupported input shapes. + // This causes a Qnn graph execution error. So, reject those configs here. + // We should consider not using QNN CPU backend for onnxruntime unit tests. + const std::string& op_type = node_unit.OpType(); + if (qnn_model_wrapper.GetQnnBackendType() == QnnBackendType::CPU && op_type == "GatherElements") { + const auto& input0 = node_unit.Inputs()[0]; + std::vector input0_shape; + ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(input0.node_arg, input0_shape), + "Cannot get input[0] shape for ", op_type, " node ", node_unit.Name()); - // Process indices - const auto& input_name = inputs[1].node_arg.Name(); + const size_t input0_rank = input0_shape.size(); + ORT_RETURN_IF_NOT(input0_rank > 1 && input0_rank <= 4, + "QNN CPU backend does not support ", op_type, " with input[0] of rank ", input0_rank); + } + + return BaseOpBuilder::IsOpSupported(qnn_model_wrapper, node_unit, logger); +} + +// Makes negative indices positive and converts int64 indices to another integer type (typically int32 or uint32). +// The input and output are both represented as byte arrays. +template +static bool FixStaticIndices(const std::vector& onnx_bytes, + int64_t input0_axis_dim, + /*out*/ std::vector& qnn_bytes) { + const size_t num_elems = onnx_bytes.size() / sizeof(SrcType); + gsl::span onnx_indices{reinterpret_cast(onnx_bytes.data()), num_elems}; + + qnn_bytes.resize(num_elems * sizeof(DstType)); + gsl::span qnn_indices{reinterpret_cast(qnn_bytes.data()), num_elems}; + + for (size_t i = 0; i < num_elems; i++) { + SrcType onnx_index = onnx_indices[i]; + + // Try to make a negative index positive by adding rank. + if (onnx_index < 0) { + onnx_index += static_cast(input0_axis_dim); + } + + if (onnx_index < 0 || static_cast(onnx_index) >= input0_axis_dim) { + return false; // QNN does not support out-of-bounds indices. + } + + qnn_indices[i] = static_cast(onnx_index); + } + + return true; +} + +// Gets the size of input0 on the axis dimension. +static Status GetInpu0AxisDimValue(const QnnModelWrapper& qnn_model_wrapper, + const NodeUnit& node_unit, + int64_t default_axis_value, + /*out*/ int64_t& axis_dim_value) { + const auto& input0 = node_unit.Inputs()[0]; + std::vector input0_shape; + ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(input0.node_arg, input0_shape), + "Cannot get shape for ", node_unit.OpType(), " input[0] ", input0.node_arg.Name()); + + int64_t rank = static_cast(input0_shape.size()); + NodeAttrHelper node_helper(node_unit); + int64_t onnx_axis = node_helper.Get("axis", default_axis_value); + if (onnx_axis < 0) { + onnx_axis += rank; + } + ORT_RETURN_IF_NOT((onnx_axis >= 0 && onnx_axis < static_cast(input0_shape.size())), + "QNN requires axis range [0, rank-1] for ", node_unit.OpType()); + + axis_dim_value = static_cast(input0_shape[onnx_axis]); + + return Status::OK(); +} + +// Processes the indices input to Gather operators. +// +// In general, QNN only supports int32/uint32 indices. QNN EP has to add Cast for dynamic int64 indices or +// convert static int64 indices to int32. +// +// The HTP backend only supports dynamic int64 indices if they are a graph input. +static Status ProcessIndicesInput(QnnModelWrapper& qnn_model_wrapper, + const NodeUnitIODef& indices_input, + int64_t input0_axis_dim, + const logging::Logger& logger, + std::vector& input_names, + bool do_op_validation) { + const auto& input_name = indices_input.node_arg.Name(); if (qnn_model_wrapper.IsQnnTensorWrapperExist(input_name)) { LOGS(logger, VERBOSE) << "Tensor already added, skip it: " << input_name; input_names.push_back(input_name); return Status::OK(); } - std::string indices_input_name(input_name); - Qnn_DataType_t qnn_data_type = QNN_DATATYPE_INT_32; - const auto* type_proto = inputs[1].node_arg.TypeAsProto(); - ORT_RETURN_IF_ERROR(utils::GetQnnDataType(false, type_proto, qnn_data_type)); + TensorInfo indices_info = {}; + ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(indices_input, indices_info)); - std::vector unpacked_tensor; - std::vector gather_indices; - bool is_initializer_input = qnn_model_wrapper.IsInitializerInput(input_name); + const bool is_npu_backend = IsNpuBackend(qnn_model_wrapper.GetQnnBackendType()); + const bool is_graph_input = qnn_model_wrapper.IsGraphInput(input_name); + ORT_RETURN_IF(is_npu_backend && + (indices_info.qnn_data_type == QNN_DATATYPE_INT_64) && + !(indices_info.is_initializer || is_graph_input), + "HTP backend doesn't support a Gather* op with a dynamic int64 input activation ", + "unless it is a graph input."); - // Gather input 0 is quantized tensor, input 1 (indices) is int64, this is not supported by QNN - bool is_quantized_tensor = inputs[0].quant_param.has_value(); - ORT_RETURN_IF(is_quantized_tensor && qnn_data_type == QNN_DATATYPE_INT_64 && !is_initializer_input, - "HTP backend doesn't support any int64 data type."); + std::vector qnn_indices_bytes; - if (is_initializer_input) { - const auto& input_tensor = qnn_model_wrapper.GetInitializerTensors().at(input_name); - ORT_RETURN_IF_ERROR(qnn_model_wrapper.UnpackInitializerData(*input_tensor, unpacked_tensor)); - if (qnn_data_type == QNN_DATATYPE_INT_64) { - // Convert initializer from int64 to int32 - size_t size = unpacked_tensor.size() / sizeof(int64_t); - const int64_t* gather_indices_int64 = reinterpret_cast(unpacked_tensor.data()); - gather_indices.resize(size * sizeof(int32_t)); - int32_t* gather_indices_int32 = reinterpret_cast(gather_indices.data()); - std::transform(gather_indices_int64, gather_indices_int64 + size, gather_indices_int32, - [](int64_t item) { return SafeInt(item); }); + // Get raw bytes for static indices. + // If indices are int64, convert them to int32 and update indices_info.qnn_data_type. + if (indices_info.is_initializer) { + std::vector onnx_indices_bytes; + ORT_RETURN_IF_ERROR(qnn_model_wrapper.UnpackInitializerData(*indices_info.initializer_tensor, onnx_indices_bytes)); + + if (indices_info.qnn_data_type == QNN_DATATYPE_INT_64) { + ORT_RETURN_IF_NOT((FixStaticIndices(onnx_indices_bytes, input0_axis_dim, qnn_indices_bytes)), + "QNN does not support negative index values for Gather* ops"); + indices_info.qnn_data_type = QNN_DATATYPE_INT_32; + } else if (indices_info.qnn_data_type == QNN_DATATYPE_INT_32) { + ORT_RETURN_IF_NOT((FixStaticIndices(onnx_indices_bytes, input0_axis_dim, qnn_indices_bytes)), + "QNN does not support negative index values for Gather* ops"); } else { - gather_indices = std::move(unpacked_tensor); + qnn_indices_bytes = std::move(onnx_indices_bytes); } - qnn_data_type = QNN_DATATYPE_INT_32; } Qnn_TensorType_t tensor_type = qnn_model_wrapper.GetTensorType(input_name); - std::vector input_shape; - ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(inputs[1].node_arg, input_shape), "Cannot get shape"); - std::vector cast_output_shape(input_shape); - QnnTensorWrapper input_tensorwrapper(input_name, tensor_type, qnn_data_type, QnnQuantParamsWrapper(), - std::move(input_shape), std::move(gather_indices)); + std::vector cast_output_shape(indices_info.shape); + QnnTensorWrapper input_tensorwrapper(input_name, tensor_type, indices_info.qnn_data_type, QnnQuantParamsWrapper(), + std::move(indices_info.shape), std::move(qnn_indices_bytes)); ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(input_tensorwrapper)), "Failed to add tensor."); - if (!is_initializer_input && qnn_data_type == QNN_DATATYPE_INT_64) { - // Insert cast node int64 -> int32 - if (qnn_data_type == QNN_DATATYPE_INT_64) { - // Add Cast node for indices - indices_input_name = input_name + "_ort_qnn_ep_cast"; - QnnTensorWrapper cast_output(indices_input_name, QNN_TENSOR_TYPE_NATIVE, QNN_DATATYPE_INT_32, - QnnQuantParamsWrapper(), std::move(cast_output_shape)); - ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_output)), "Failed to add tensor."); - ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(indices_input_name, - QNN_OP_PACKAGE_NAME_QTI_AISW, - "Cast", - {input_name}, - {indices_input_name}, - {}, - do_op_validation), - "Failed to add node."); - } + // Insert QNN Cast op to convert dynamic indices from int64 to int32. + std::string indices_input_name(input_name); + if (indices_info.qnn_data_type == QNN_DATATYPE_INT_64) { + assert(!indices_info.is_initializer); + + indices_input_name = input_name + "_ort_qnn_ep_cast"; + QnnTensorWrapper cast_output(indices_input_name, QNN_TENSOR_TYPE_NATIVE, QNN_DATATYPE_INT_32, + QnnQuantParamsWrapper(), std::move(cast_output_shape)); + ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_output)), "Failed to add tensor."); + ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(indices_input_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + "Cast", + {input_name}, + {indices_input_name}, + {}, + do_op_validation), + "Failed to add node."); } input_names.push_back(indices_input_name); @@ -112,20 +193,44 @@ Status GatherOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, return Status::OK(); } +Status GatherOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, + const NodeUnit& node_unit, + const logging::Logger& logger, + std::vector& input_names, + bool do_op_validation) const { + const auto& inputs = node_unit.Inputs(); + ORT_RETURN_IF(inputs.size() != 2, "QNN EP: ", node_unit.OpType(), " operator must have two inputs"); + ORT_RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, inputs[0], logger, input_names)); + + int64_t input0_axis_dim = 0; + ORT_RETURN_IF_ERROR(GetInpu0AxisDimValue(qnn_model_wrapper, node_unit, /*default_axis*/ 0, input0_axis_dim)); + + return ProcessIndicesInput(qnn_model_wrapper, inputs[1], input0_axis_dim, logger, input_names, do_op_validation); +} + Status GatherOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit, std::vector&& input_names, const logging::Logger& logger, bool do_op_validation) const { - ORT_UNUSED_PARAMETER(logger); + const bool is_gather_elems = node_unit.OpType() == "GatherElements"; + + // Create QNN 'axis' parameter. std::vector param_tensor_names; int32_t axis_value = 0; Qnn_Scalar_t axis_qnn_scalar = QNN_SCALAR_INIT; ORT_RETURN_IF_ERROR(ProcessAxisAttribute(qnn_model_wrapper, node_unit, axis_qnn_scalar, axis_value)); - QnnParamWrapper axis_param(node_unit.Index(), node_unit.Name(), QNN_OP_GATHER_PARAM_AXIS, axis_qnn_scalar); + QnnParamWrapper axis_param(node_unit.Index(), node_unit.Name(), + (is_gather_elems ? QNN_OP_GATHER_ELEMENTS_PARAM_AXIS : QNN_OP_GATHER_PARAM_AXIS), + axis_qnn_scalar); param_tensor_names.push_back(axis_param.GetParamTensorName()); qnn_model_wrapper.AddParamWrapper(std::move(axis_param)); + if (is_gather_elems) { + return ProcessOutputs(qnn_model_wrapper, node_unit, std::move(input_names), std::move(param_tensor_names), + logger, do_op_validation, GetQnnOpType(node_unit.OpType())); + } + // if indicies is scalar shape, then need to add Reshape node const auto& input_tensor_wrapper = qnn_model_wrapper.GetQnnTensorWrapper(input_names[0]); const auto& indices_input_tensor_wrapper = qnn_model_wrapper.GetQnnTensorWrapper(input_names[1]); diff --git a/onnxruntime/python/tools/quantization/operators/gather.py b/onnxruntime/python/tools/quantization/operators/gather.py index e390e874a2..e6314407cd 100644 --- a/onnxruntime/python/tools/quantization/operators/gather.py +++ b/onnxruntime/python/tools/quantization/operators/gather.py @@ -55,7 +55,7 @@ class QDQGather(QDQOperatorBase): def quantize(self): node = self.node - assert node.op_type == "Gather" + assert node.op_type == "Gather" or node.op_type == "GatherElements" if self.quantizer.is_valid_quantize_weight(node.input[0]) or self.quantizer.force_quantize_no_input_check: self.quantizer.quantize_activation_tensor(node.input[0]) diff --git a/onnxruntime/python/tools/quantization/registry.py b/onnxruntime/python/tools/quantization/registry.py index caac829126..160b056e1d 100644 --- a/onnxruntime/python/tools/quantization/registry.py +++ b/onnxruntime/python/tools/quantization/registry.py @@ -79,6 +79,7 @@ QDQRegistry = { "MatMul": QDQMatMul, "Split": QDQSplit, "Gather": QDQGather, + "GatherElements": QDQGather, "Where": QDQWhere, "InstanceNormalization": QDQNormalization, "LayerNormalization": QDQNormalization, diff --git a/onnxruntime/test/onnx/main.cc b/onnxruntime/test/onnx/main.cc index 343027b42c..0397bba904 100644 --- a/onnxruntime/test/onnx/main.cc +++ b/onnxruntime/test/onnx/main.cc @@ -827,7 +827,8 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)"); ORT_TSTR("sce_NCd1d2d3_sum_weight_high_ii_expanded"), ORT_TSTR("sce_none_weights_log_prob_expanded"), ORT_TSTR("sce_none_weights_expanded"), - ORT_TSTR("convtranspose_3d")}; + ORT_TSTR("convtranspose_3d"), + ORT_TSTR("gather_elements_negative_indices")}; std::unordered_set> all_disabled_tests(std::begin(immutable_broken_tests), std::end(immutable_broken_tests)); diff --git a/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc b/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc index b4bd3fca7b..5b2d00bb95 100644 --- a/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc @@ -42,9 +42,15 @@ void GetData(const std::vector& input_dims, const std::vector& output_data.resize(output_size); std::srand(static_cast(std::time(0))); for (size_t i = 0; i < indices_size; ++i) { +#if defined(USE_QNN) + // Negative index not possible. + indices_data[i] = + static_cast(static_cast(std::rand()) % input_dims[axis]); +#else // Negative index possible. indices_data[i] = static_cast((static_cast(std::rand()) % (input_dims[axis] * 2)) - input_dims[axis]); +#endif } for (size_t i = 0; i < output_size; ++i) { int64_t input_offset = 0; @@ -382,9 +388,10 @@ TEST(GatherElementsOpTest, IndicesOutOfBounds) { // skip cuda as the cuda kernel won't throw the error message // skip openvino which will not throw error message but will ensure no out-of-bound access // skip TensorRT because it doesn't support out of bounds indices + // skip QNN because it doesn't support out of bounds indices test.Run(OpTester::ExpectResult::kExpectFailure, "", {kCudaExecutionProvider, kCudaNHWCExecutionProvider, kRocmExecutionProvider, kOpenVINOExecutionProvider, - kTensorrtExecutionProvider, kDmlExecutionProvider}); + kTensorrtExecutionProvider, kDmlExecutionProvider, kQnnExecutionProvider}); } TEST(GatherElementsOpTest, BigIndices) { diff --git a/onnxruntime/test/providers/qnn/gather_elems_op_test.cc b/onnxruntime/test/providers/qnn/gather_elems_op_test.cc new file mode 100644 index 0000000000..81c0887306 --- /dev/null +++ b/onnxruntime/test/providers/qnn/gather_elems_op_test.cc @@ -0,0 +1,319 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include + +#include "core/graph/node_attr_utils.h" +#include "test/optimizer/qdq_test_utils.h" +#include "test/providers/qnn/qnn_test_utils.h" + +#include "onnx/onnx_pb.h" + +#include "gtest/gtest.h" + +namespace onnxruntime { +namespace test { + +// Creates a graph with a single Q/DQ GatherElements operator. Used for testing HTP backend. +template +static GetTestQDQModelFn BuildQDQGatherElemsTestCase(const TestInputDef& input_def, + const TestInputDef& indices_def, + const std::vector& attrs, + bool use_contrib_qdq = false) { + return [input_def, indices_def, attrs, use_contrib_qdq](ModelTestBuilder& builder, + std::vector>& output_qparams) { + // input -> Q -> DQ -> + NodeArg* input = MakeTestInput(builder, input_def); + QuantParams input_qparams = GetTestInputQuantParams(input_def); + NodeArg* input_qdq = AddQDQNodePair(builder, input, input_qparams.scale, input_qparams.zero_point, + use_contrib_qdq); + + // indices input + NodeArg* indices_input = MakeTestInput(builder, indices_def); + + // GatherElements op + NodeArg* gather_output = builder.MakeIntermediate(); + Node& gather_node = builder.AddNode("GatherElements", {input_qdq, indices_input}, {gather_output}); + + for (const auto& attr : attrs) { + gather_node.AddAttributeProto(attr); + } + + // op_output -> Q -> DQ -> output + // NOTE: Input and output quantization parameters must be equal for GatherElements. + output_qparams[0] = input_qparams; // Overwrite! + AddQDQNodePairWithOutputAsGraphOutput(builder, gather_output, input_qparams.scale, + input_qparams.zero_point, use_contrib_qdq); + }; +} + +// Runs an GatherElements model on the QNN CPU backend. Checks the graph node assignment, and that inference +// outputs for QNN EP and CPU EP match. +template +static void RunCPUGatherElemsOpTest(const TestInputDef& input_def, + const TestInputDef& indices_def, + const std::vector& attrs, + ExpectedEPNodeAssignment expected_ep_assignment, + int opset = 13) { + ProviderOptions provider_options; + float fp32_abs_err = 1e-5f; // default tolerance + +#if defined(_WIN32) + provider_options["backend_path"] = "QnnCpu.dll"; +#else + provider_options["backend_path"] = "libQnnCpu.so"; +#endif + + RunQnnModelTest(BuildOpTestCase("GatherElements", {input_def}, {indices_def}, attrs), + provider_options, + opset, + expected_ep_assignment, + fp32_abs_err); +} + +// Runs a QDQ GatherElements model on the QNN HTP backend. Checks the graph node assignment, and that inference +// outputs for QNN EP and CPU EP match with expected accuracy. +template +static void RunHTPQDQGatherElemsOpTest(const TestInputDef& input_def, + const TestInputDef& indices_def, + const std::vector& attrs, + ExpectedEPNodeAssignment expected_ep_assignment, + int opset = 13, + bool use_contrib_qdq = false) { + ProviderOptions provider_options; + +#if defined(_WIN32) + provider_options["backend_path"] = "QnnHtp.dll"; +#else + provider_options["backend_path"] = "libQnnHtp.so"; +#endif + + auto f32_model_builder = BuildOpTestCase("GatherElements", {input_def}, {indices_def}, attrs); + auto qdq_model_builder = BuildQDQGatherElemsTestCase(input_def, indices_def, attrs, + use_contrib_qdq); + + TestQDQModelAccuracy(f32_model_builder, + qdq_model_builder, + provider_options, + opset, + expected_ep_assignment); +} + +// Runs a non-quantized GatherElements model on the QNN HTP backend. Checks the graph node assignment, +// and that inference outputs for QNN EP and CPU EP match. +template +static void RunHTPGatherElemsOpTest(const TestInputDef& input_def, + const TestInputDef& indices_def, + const std::vector& attrs, + ExpectedEPNodeAssignment expected_ep_assignment, + int opset = 13) { + ProviderOptions provider_options; + float fp32_abs_err = 1e-5f; // default tolerance + +#if defined(_WIN32) + provider_options["backend_path"] = "QnnHtp.dll"; +#else + provider_options["backend_path"] = "libQnnHtp.so"; +#endif + + RunQnnModelTest(BuildOpTestCase("GatherElements", {input_def}, {indices_def}, attrs), + provider_options, + opset, + expected_ep_assignment, + fp32_abs_err); +} + +// +// CPU tests: +// + +// Test GatherElements op on CPU backend: +// positive, dynamic, int64 indices. +TEST_F(QnnCPUBackendTests, GatherElems_DataF32_IndicesInt64) { + RunCPUGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, false, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test GatherElements op on CPU backend: +// positive, dynamic, int32 indices. +TEST_F(QnnCPUBackendTests, GatherElems_DataF32_IndicesInt32) { + RunCPUGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, false, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test GatherElements op on CPU backend: +// positive, static, int64 indices. +TEST_F(QnnCPUBackendTests, GatherElems_DataF32_StaticIndicesInt64) { + RunCPUGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, true, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test GatherElements op on CPU backend: +// positive, static, int32 indices. +TEST_F(QnnCPUBackendTests, GatherElems_DataF32_StaticIndicesInt32) { + RunCPUGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, true, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +#if defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__) +// +// HTP tests: +// + +// Test non-quantized GatherElements op on HTP backend: +// Input[0] is int32. Indices are int32, positive, and dynamic. Both inputs are rank 1. +TEST_F(QnnHTPBackendTests, GatherElems_DataInt32_IndicesInt32_Rank1) { + RunHTPGatherElemsOpTest( + TestInputDef({3}, false, {1, 2, 3}), + TestInputDef({2}, false, {1, 2}), + {utils::MakeAttribute("axis", static_cast(0))}, + ExpectedEPNodeAssignment::All); +} + +// Test uint8 QDQ GatherElements op on HTP backend: +// positive, dynamic, int32 indices. +TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_IndicesInt32) { + RunHTPQDQGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, false, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test uint8 QDQ GatherElements op on HTP backend: +// positive, static, int32 indices. +TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticIndicesInt32) { + RunHTPQDQGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, true, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test GatherElements op on HTP backend: +// positive, static, int64 indices. +// HTP does not support int64, but QNN EP converts static int64 indices into int32. +TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticIndicesInt64) { + RunHTPQDQGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, true, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test GatherElements op on HTP backend: +// negative, static, int32 indices. +TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticNegIndicesInt32) { + RunHTPQDQGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, true, {1, 2, -3, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test GatherElements op on HTP backend: +// negative, static, int64 indices. +TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticNegIndicesInt64) { + RunHTPQDQGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, true, {1, -1, -3, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +// Test QDQ GatherElements op on HTP backend: +// Input[0] is uint16_t, and rank 4. Indices are int64_t, rank 4, negative, and static. +// Axis is negative (points to last dim). +TEST_F(QnnHTPBackendTests, GatherElems_DataUint16_StaticNegIndicesInt64) { + const std::vector input_shape = {1, 2, 3, 3}; + const std::vector input_data = GetSequentialFloatData(input_shape, -8.0f, 1.0f); + RunHTPQDQGatherElemsOpTest( + TestInputDef(input_shape, false, input_data), + TestInputDef({1, 1, 2, 2}, true, {0, -1, -3, 2}), + {utils::MakeAttribute("axis", static_cast(-1))}, + ExpectedEPNodeAssignment::All, + /*opset*/ 21); +} + +// Test QDQ GatherElements op on HTP backend with large number of indices: +TEST_F(QnnHTPBackendTests, GatherElems_DataUint16_StaticNegIndicesInt64_Large) { + const std::vector input_shape = {12, 1024, 512}; + std::vector input_data(12 * 1024 * 512); + for (size_t i = 0; i < input_data.size(); i++) { + input_data[i] = static_cast((static_cast(i) % 8)); + } + + RunHTPQDQGatherElemsOpTest( + TestInputDef(input_shape, false, input_data), + TestInputDef({12, 1024, 1024}, true, -512, 511), + {utils::MakeAttribute("axis", static_cast(-1))}, + ExpectedEPNodeAssignment::All, + /*opset*/ 21); +} + +// Test QDQ GatherElements op on HTP backend with large number of indices. +// TODO: Investigate inaccuracy. +// Negative input[0] values seem to cause inaccuracies with 2M+ indices. +// Inaccuracy detected for output 'output_0', element 131072 +// output_range=300, tolerance=0.40000000596046448%. +// Expected val (f32@CPU_EP): 121 +// qdq@QNN_EP val: -97.999542236328125 (err: 218.99954223632812, err/output_range: 72.999847412109375%) +// qdq@CPU_EP val: 120.99794006347656 (err: 0.0020599365234375, err/output_range: 0.0006866455078125%) +// abs(qdq@QNN_EP - qdq@CPU_EP) / output_range = 72.999160766601562% +TEST_F(QnnHTPBackendTests, DISABLED_GatherElems_DataUint16_StaticNegIndicesInt64_Large2) { + // Input data with sequential values from -98.0f to 202.0f + const std::vector input_shape = {12, 1024, 512}; + std::vector input_data(12 * 1024 * 512); + for (size_t i = 0; i < input_data.size(); i++) { + int32_t int_val = -98 + (static_cast(i) % 301); + input_data[i] = static_cast(int_val); + } + + // Indices with values between -512 to 511. + const std::vector indices_shape = {12, 1024, 1024}; + std::vector indices(12 * 1024 * 1024); + for (size_t i = 0; i < indices.size(); i++) { + indices[i] = static_cast(-512 + (static_cast(i) % 1024)); + } + + RunHTPQDQGatherElemsOpTest( + TestInputDef(input_shape, false, input_data), + TestInputDef(indices_shape, true, indices), + {utils::MakeAttribute("axis", static_cast(-1))}, + ExpectedEPNodeAssignment::All, + /*opset*/ 21); +} + +// Test GatherElements op on HTP backend: +// Tests that dynamic int64 indices are supported on HTP backend if the indices are a graph input. +// QNN SDK 2.23 added support for Cast from int64 to int32. +TEST_F(QnnHTPBackendTests, GatherElems_DynamicInt64IndicesSupportedAsGraphInput) { + RunHTPQDQGatherElemsOpTest( + TestInputDef({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}), + TestInputDef({2, 3}, false, {1, 2, 0, 2, 0, 0}), + {utils::MakeAttribute("axis", static_cast(1))}, + ExpectedEPNodeAssignment::All); +} + +#endif // defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__) + +} // namespace test +} // namespace onnxruntime + +#endif // !defined(ORT_MINIMAL_BUILD) diff --git a/onnxruntime/test/providers/qnn/gather_op_htp_test.cc b/onnxruntime/test/providers/qnn/gather_op_htp_test.cc index 48cd5ad995..e3f09e9259 100644 --- a/onnxruntime/test/providers/qnn/gather_op_htp_test.cc +++ b/onnxruntime/test/providers/qnn/gather_op_htp_test.cc @@ -97,13 +97,14 @@ TEST_F(QnnHTPBackendTests, GatherOp_U16_IndicesStaticInt64_Axis0) { true); // Use 'com.microsoft' Q/DQ ops } -// Tests that dynamic int64 indices are not supported on HTP backend. +// Tests that dynamic int64 indices are supported on HTP backend if the indices are a graph input. +// QNN SDK 2.23 added support for Cast from int64 to int32. TEST_F(QnnHTPBackendTests, GatherOp_IndicesDynamicInt64_Axis0) { RunQDQGatherOpTest(TestInputDef({3, 2}, false, {1.0f, 1.2f, 2.3f, 3.4f, 4.5f, 5.7f}), TestInputDef({2, 2}, false, {0, 1, 1, 2}), {utils::MakeAttribute("axis", static_cast(0))}, 13, - ExpectedEPNodeAssignment::None); + ExpectedEPNodeAssignment::All); } // Test creates a DQ -> Gather -> Q -> DQ graph, and checks that all diff --git a/onnxruntime/test/python/quantization/test_op_gather_elems.py b/onnxruntime/test/python/quantization/test_op_gather_elems.py new file mode 100644 index 0000000000..04da34ef2c --- /dev/null +++ b/onnxruntime/test/python/quantization/test_op_gather_elems.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import os +import tempfile +import unittest + +import numpy as np +import onnx +from op_test_utils import TestDataFeeds, check_op_type_count + +from onnxruntime.quantization import QuantFormat, QuantType, quantize_static + + +class TestQDQGatherElements(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls._tmp_model_dir = tempfile.TemporaryDirectory(prefix="ort.qdq.gatherelems_") + + # Note: swap with the commented line if you want to see the models in local test dir. + cls._tmp_dir_path = cls._tmp_model_dir.name + # cls._tmp_dir_path = "." + + @classmethod + def tearDownClass(cls): + cls._tmp_model_dir.cleanup() + + def build_test_model_static_indices( + self, + inp_shape: list[int], + indices_data: np.ndarray, + axis: int = 0, + ): + input_0 = onnx.helper.make_tensor_value_info("input_0", onnx.TensorProto.FLOAT, inp_shape) + output_0 = onnx.helper.make_tensor_value_info("output_0", onnx.TensorProto.FLOAT, None) + indices = onnx.numpy_helper.from_array(indices_data, "indices") + + gatherelems_node = onnx.helper.make_node( + "GatherElements", ["input_0", "indices"], ["output_0"], axis=axis, name="GatherElems0" + ) + graph = onnx.helper.make_graph( + [gatherelems_node], + "GatherElemsf32", + [input_0], + [output_0], + initializer=[indices], + ) + opset_imports = [onnx.helper.make_opsetid("", 21)] + model = onnx.helper.make_model(graph, opset_imports=opset_imports) + + return onnx.shape_inference.infer_shapes(model) + + def test_qdq_gatherelems_static_indices(self): + """ + Test quantization of GatherElements with static indices. + """ + float_model_path = os.path.join(self._tmp_dir_path, "gather_elems.f32.onnx") + qdq_model_path = os.path.join(self._tmp_dir_path, "gather_elems.qdq.onnx") + + inp_shape = [3, 3] + indices_data = np.array([[1, 2, 0], [2, 0, 0]], dtype=np.int64) + float_model = self.build_test_model_static_indices(inp_shape, indices_data, axis=0) + + onnx.checker.check_model(float_model, True) + onnx.save_model(float_model, float_model_path) + + input_data1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32) + input_data2 = np.array([[8, 2, 9], [4, 5, 6], [7, 8, 1]], dtype=np.float32) + data_reader = TestDataFeeds([{"input_0": input_data1}, {"input_0": input_data2}]) + + quantize_static( + float_model_path, + qdq_model_path, + data_reader, + quant_format=QuantFormat.QDQ, + activation_type=QuantType.QUInt16, + weight_type=QuantType.QUInt8, + op_types_to_quantize=[node.op_type for node in float_model.graph.node], + extra_options={ + "ForceQuantizeNoInputCheck": True, + }, + ) + + qdq_node_counts = {"QuantizeLinear": 2, "DequantizeLinear": 2} + check_op_type_count(self, qdq_model_path, **qdq_node_counts) + + qdq_model = onnx.load_model(qdq_model_path) + onnx.checker.check_model(qdq_model, True) + + initializers = {init.name: init for init in qdq_model.graph.initializer} + + zp_name = "input_0_zero_point" + scale_name = "input_0_scale" + self.assertIn(zp_name, initializers) + self.assertIn(scale_name, initializers) + + # Check that all Q/DQ nodes use the same scale and zero-point + for node in qdq_model.graph.node: + if node.op_type == "QuantizeLinear" or node.op_type == "DequantizeLinear": + self.assertEqual(node.input[1], scale_name) + self.assertEqual(node.input[2], zp_name) + + +if __name__ == "__main__": + unittest.main()