onnxruntime
7a795ed7 - Improve SparseTensors public API input validation as well as sparse utilities (#28227)

Commit
6 days ago
Improve SparseTensors public API input validation as well as sparse utilities (#28227) This pull request significantly improves the safety and robustness of sparse tensor handling in ONNX Runtime. The main focus is on adding thorough bounds checking and using safe integer arithmetic to prevent overflows and invalid memory accesses when working with sparse tensor indices. Additionally, the Python bindings for sparse tensors are refactored to ensure correct object lifetimes and memory management when exposing data to NumPy. **Sparse Tensor Index Validation and Safety** * Added comprehensive bounds checks for COO and CSR sparse tensor indices in both the C API (`onnxruntime_c_api.cc`) and core conversion utilities, ensuring indices are within valid ranges and, for CSR, that outer indices are non-decreasing and within bounds. [[1]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R449-R485) [[2]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R521-R547) [[3]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R659-R696) [[4]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R721-R747) [[5]](diffhunk://#diff-620fd022510c5134fc9bd3c8d01bc5772cc78a82043b0da5e44cf2482038dc37L267-R273) [[6]](diffhunk://#diff-620fd022510c5134fc9bd3c8d01bc5772cc78a82043b0da5e44cf2482038dc37L359-R376) * Replaced direct arithmetic with `SafeInt` for all index and size calculations to prevent integer overflows, especially when converting between types or computing dense tensor offsets. [[1]](diffhunk://#diff-620fd022510c5134fc9bd3c8d01bc5772cc78a82043b0da5e44cf2482038dc37L267-R273) [[2]](diffhunk://#diff-d31e9fbe0f5334fcd949833e035f2b25d5ae810dcd505c545f6b372b546b1406L2077-R2077) [[3]](diffhunk://#diff-d31e9fbe0f5334fcd949833e035f2b25d5ae810dcd505c545f6b372b546b1406L2091-R2091) [[4]](diffhunk://#diff-d31e9fbe0f5334fcd949833e035f2b25d5ae810dcd505c545f6b372b546b1406L2110-R2110) [[5]](diffhunk://#diff-d31e9fbe0f5334fcd949833e035f2b25d5ae810dcd505c545f6b372b546b1406L2291-R2298) * Improved error messages for invalid indices, making debugging easier by providing more context about the specific error. [[1]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R449-R485) [[2]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R521-R547) [[3]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R659-R696) [[4]](diffhunk://#diff-cff364b6b1ab4ef507d87a661a97b873405f569797fcaf91af29491f223555a8R721-R747) [[5]](diffhunk://#diff-620fd022510c5134fc9bd3c8d01bc5772cc78a82043b0da5e44cf2482038dc37L267-R273) [[6]](diffhunk://#diff-620fd022510c5134fc9bd3c8d01bc5772cc78a82043b0da5e44cf2482038dc37L359-R376) **Python Bindings Improvements** * Refactored the pybind11 bindings for sparse tensor views so that NumPy arrays referencing sparse tensor memory correctly keep the parent Python object alive, preventing potential memory issues when the sparse tensor is on the GPU or managed by Python. [[1]](diffhunk://#diff-3c1b21fe3d5903c277b4d3888f5a4c57ff8f8f6f593183a3f4865825c5ab8e0cL98-R120) [[2]](diffhunk://#diff-3c1b21fe3d5903c277b4d3888f5a4c57ff8f8f6f593183a3f4865825c5ab8e0cL299-R304) [[3]](diffhunk://#diff-3c1b21fe3d5903c277b4d3888f5a4c57ff8f8f6f593183a3f4865825c5ab8e0cL314-R319) **General Code Quality** * Added missing header include for `safeint.h` to ensure `SafeInt` is available where needed. * Minor cleanups and improved assertions to clarify intent and ensure correctness. These changes collectively make sparse tensor support in ONNX Runtime safer, more reliable, and easier to use from both C++ and Python.
Author
Parents
Loading