diff --git a/tmva/sofie/inc/TMVA/ROperator_Gather.hxx b/tmva/sofie/inc/TMVA/ROperator_Gather.hxx index fafe0dd63ae92..605cfb6993db9 100644 --- a/tmva/sofie/inc/TMVA/ROperator_Gather.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_Gather.hxx @@ -186,16 +186,6 @@ public: auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY); auto stridesIndices = UTILITY::ComputeStrideFromShape(fShapeIndices); - // case fIndices is not known we need to correct for negative axis indices at run-time - if (fIndices.empty()) { - auto indicesLength = ConvertDimShapeToLength(fShapeIndices); - out << SP << "// correct in case of negative gather indices\n"; - out << SP << "for (size_t i = 0; i < " << indicesLength << "; i++){\n"; - out << SP << SP << "if (tensor_" << fNIndices << "[i] < 0)\n"; - out << SP << SP << SP << "tensor_" << fNIndices << "[i] += " << fShapeX[fAttrAxis] << ";\n"; - out << SP << "}\n"; - } - // Fill the output Y[j_0, j_1, ..., j_{axis - 1}, i_0, i_1, ..., i_{q - 1}, j_{axis + 1}, ..., j_{r - 1}] // [0 ... axis) [axis ... axis + q) [axis + q ... q + r - 1) // iterate in [0 ... axis) [0 ... q) [axis ... r - 1) @@ -260,8 +250,16 @@ public: out << ";\n"; // K + // when the indices are not a known constant, correct at the read site for + // possible negative values (the indices tensor may be a const model input) for (size_t k = 0; k < q + r; k++) out << SP; - out << "size_t k = static_cast(" << "tensor_" << fNIndices << "[i_index]" << ");\n"; + if (fIndices.empty()) { + out << "int64_t k_i = static_cast(tensor_" << fNIndices << "[i_index]);\n"; + for (size_t k = 0; k < q + r; k++) out << SP; + out << "size_t k = static_cast(k_i < 0 ? k_i + " << fShapeX[fAttrAxis] << " : k_i);\n"; + } else { + out << "size_t k = static_cast(" << "tensor_" << fNIndices << "[i_index]" << ");\n"; + } // Input for (size_t k = 0; k < q + r; k++) out << SP; out << "size_t x_index = k"; diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx index 9867cc2b8c215..7721e039273ea 100644 --- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx +++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx @@ -1228,6 +1228,18 @@ TEST(ONNX, GatherNegativeIndices) expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); } +// The gather indices are a graph input here, so the generated code receives +// them as an "int64_t const*": correcting the negative ones must not write +// back into the indices tensor, otherwise the model does not even compile. +TEST(ONNX, GatherRuntimeNegativeIndices) +{ + SofieReference ref = readReference("GatherRuntimeNegativeIndices"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "GatherRuntimeNegativeIndices", ref.f32("input0"), ref.i64("input1")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + TEST(ONNX, Slice) { SofieReference ref = readReference("Slice"); diff --git a/tmva/sofie/test/generate_input_models.py b/tmva/sofie/test/generate_input_models.py index 0946fb3985e92..271bd910771b7 100644 --- a/tmva/sofie/test/generate_input_models.py +++ b/tmva/sofie/test/generate_input_models.py @@ -1846,6 +1846,28 @@ def make_GatherNegativeIndices(): return _model(graph, opset=17, ir_version=8) +def make_GatherRuntimeNegativeIndices(): + """Ops: Gather""" + # The indices are a graph input and not an initializer, so their values are + # only known at run time and the generated code gets them as a pointer to + # const: it must correct the negative ones without writing to the tensor. + nodes = [ + helper.make_node('Gather', ['X', 'I'], ['Y'], axis=0), + ] + graph = helper.make_graph( + nodes, + 'Gather', + inputs=[ + _vi('X', FLOAT, [5, 2]), + _vi('I', INT64, [3]), + ], + outputs=[ + _vi('Y', FLOAT, [3, 2]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + def make_Gelu(): """Ops: Gelu""" nodes = [ @@ -5298,6 +5320,7 @@ def make_Where(): 'GatherND_2': make_GatherND_2, 'GatherND_3': make_GatherND_3, 'GatherNegativeIndices': make_GatherNegativeIndices, + 'GatherRuntimeNegativeIndices': make_GatherRuntimeNegativeIndices, 'Gelu': make_Gelu, 'Gemm_ConstantFolding': make_Gemm_ConstantFolding, 'Gemm_ConstantFolding_Shared': make_Gemm_ConstantFolding_Shared, @@ -5517,6 +5540,7 @@ def rand_f32(seed, shape): 'GatherAxis2': [f32(np.arange(0.0, 120.0), (5, 4, 3, 2))], 'GatherAxis3': [f32(np.arange(0.0, 120.0), (5, 4, 3, 2))], 'GatherNegativeIndices': [f32(np.arange(0.0, 10.0), (10,))], + 'GatherRuntimeNegativeIndices': [f32(np.arange(0.0, 10.0), (5, 2)), i64([-1, 2, -5])], 'Gelu': [f32([1.0, -2.0, 3.0, 0.5, -1.0, 2.0], (6,))], # Note: the second operand must produce a mix of true and false results, # otherwise a constant implementation would pass the test.