[tmva][sofie] Extend ONNX operator and type support, fix two codegen bugs, speed up Reduce and TopK - #23331
Merged
guitargeek merged 6 commits intoSep 10, 2026
Conversation
…olic TopK K Adds the ONNX operators and shape handling needed to convert a GNN model with a dynamic input dimension: - ReduceMax and ReduceMin, on all three reduction paths - Elu on tensors with a parametric dimension (Dim shapes, as in Relu) - TopK with K taken from a shape tensor, not only an initializer - Min/Max/Sum/Mean fold to a shape tensor when all inputs are rank <= 1 INT64 and known at initialization, as BasicBinary already did Assisted-by: ClaudeCode:claude-opus-5
The parser instantiated the n-ary operators for float only and threw for anything else, so an INT64 Min could not be converted. Adds DOUBLE, INT32 and INT64, and makes the Mean trait generic, since it was specialised for float and the integer instantiations would not have compiled. Error messages now name the operator instead of always saying "Max". Assisted-by: ClaudeCode:claude-opus-5
… sorted=0 Reduce emitted the strides unparenthesised, so for a reduction over an interior axis "i / 126*std::min(a,b) % (...)" parses as "((i / 126) * min) % (...)". The index arithmetic is then wrong and the output is garbage. This affects ReduceMean and ReduceSum too, whenever the shape expressions are not single tokens, i.e. for dynamic shapes. TopK with sorted=0 called std::partial_sort without a comparator, so the default ascending order on std::pair applied and largest=1 returned the K smallest elements. Assisted-by: ClaudeCode:claude-opus-5
All three changes leave the generated output bit-identical; they were checked byte-for-byte against the previous generated code. - Reduce recovered the indices with a division and a modulo per element, which for dynamic shapes are real integer divisions. Replaced by one loop per axis in memory order, so the inner loop is division-free and vectorises. Reduce went from 66% to 18% of a GNN model's runtime. - TopK selects with nth_element + sort over the K selected, O(n) + O(K log K), instead of partial_sort's O(n log K). - TopK packs (value, index) of a float tensor into one uint64 with an order-preserving key: one instruction per comparison and 8 bytes per element instead of 16. Other types keep the pairs. Assisted-by: ClaudeCode:claude-opus-5
Eight models in generate_input_models.py with the matching gtest cases: - ReduceMax, ReduceMin - EluDynShape, for Elu on a parametric dimension - TopKWithDynShapeK, where K = min(N, 4) reaches TopK as a shape tensor - MinInt64, MaxInt64, for the n-ary operators on a non-float type - ReduceMean_kMiddle_DynShape, a reduction over an interior axis whose strides are expressions, i.e. the case that was generated wrong - TopKLargestUnsorted, largest=1 with sorted=0 Assisted-by: ClaudeCode:claude-opus-5
Contributor
Author
|
@guitargeek I opened the PR as requested. :) |
guitargeek
added
in:SOFIE
bug
AI-Assisted
experiment
guitargeek
approved these changes
Sep 10, 2026
guitargeek
left a comment
guitargeek
left a comment
Contributor
There was a problem hiding this comment.
LGTM! Thanks a lot for trying out SOFIE and fixing it up for your usecase!
Would it help you if we try to backport these fixes to 6.40, or you don't care about it working in the next 6.40 patch release?
Contributor
Author
|
I am currently just experimenting few things, so having this in the master branch (and thus in the next nightly builds) would be sufficient! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This Pull request:
Extend ONNX operator and type support, fix two codegen bugs, speed up Reduce and TopK
Changes or fixes:
Adds the ONNX support needed to convert a GNN track finder with a dynamic input dimension, fixes two code-generation bugs found along the way, and speeds up the two operators that dominate its runtime. Full description, measurements and a public reproducer are in the linked issue.
ReduceMax/ReduceMin;Eluon a parametric dimension;TopKwithKfrom a shape tensor rather than an initializer;Min/Max/Sum/Meanfolding to a shape tensor, asBasicBinaryalready did.floatonly, so an INT64Mincould not be converted. AddsDOUBLE,INT32,INT64.Reduceemitted unparenthesised strides, so a reduction over an interior axis generated wrong index arithmetic. This affects the already-supportedReduceMeanandReduceSumwhenever shapes are dynamic.TopKwithsorted=0ignoredlargestand returned the K smallest.Reduceover an interior axis no longer does two integer divisions per element (66% → 18% of my model's runtime);TopKselects withnth_elementinstead ofpartial_sortand packs(value, index)into oneuint64for float tensors.All the modifications were done by Claude (Opus 5). I tested everything against my original ONNX model and did not observe difference in terms of numerical performance. More details are in the linked issue or here.
Checklist:
This PR fixes #23308