[tmva][sofie] Do not mutate the indices tensor in Gather codegen

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<size_t>(" << "tensor_" << fNIndices << "[i_index]" << ");\n";
+      if (fIndices.empty()) {
+         out << "int64_t k_i = static_cast<int64_t>(tensor_" << fNIndices << "[i_index]);\n";
+         for (size_t k = 0; k < q + r; k++) out << SP;
+         out << "size_t k = static_cast<size_t>(k_i < 0 ? k_i + " << fShapeX[fAttrAxis] << " : k_i);\n";
+      } else {
+         out << "size_t k = static_cast<size_t>(" << "tensor_" << fNIndices << "[i_index]" << ");\n";
+      }
       // Input
       for (size_t k = 0; k < q + r; k++) out << SP;
       out << "size_t x_index = k";

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<float>, "GatherRuntimeNegativeIndices", ref.f32("input0"), ref.i64("input1"));
+
+   expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
+}
+
 TEST(ONNX, Slice)
 {
    SofieReference ref = readReference("Slice");

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.