From e4b49be06f4b1ac66c24388f20fbb3b405cf6d8d Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Tue, 25 Aug 2026 07:46:14 +0000 Subject: [PATCH] [tmva][sofie] Do not mutate the indices tensor in Gather codegen MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the gather indices are not a known constant, the generated code corrected possible negative indices by writing them back into the indices tensor: if (tensor_receivers[i] < 0) tensor_receivers[i] += 5; If the indices tensor is a model input, the generated pointer is "int64_t const*" and the emitted code does not compile. Runtime gather indices that are graph inputs are the norm for graph neural networks, where the edge sender/receiver lists are fed at inference time. Emit the correction at the read site instead, leaving the input tensor untouched. Covered by the new GatherRuntimeNegativeIndices test model, whose gather indices are a graph input instead of an initializer. The existing Gather tests all pass the indices as initializers, so the generated code took the constant-index branch and never exercised the correction. 🤖 Done with the help of AI --- tmva/sofie/inc/TMVA/ROperator_Gather.hxx | 20 ++++++++-------- tmva/sofie/test/TestCustomModelsFromONNX.cxx | 12 ++++++++++ tmva/sofie/test/generate_input_models.py | 24 ++++++++++++++++++++ 3 files changed, 45 insertions(+), 11 deletions(-) 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.