robertknight / rten

ONNX neural network inference engine
123 stars 9 forks source link

Share implementations for operators based on data type width #244

Open robertknight opened 4 months ago

robertknight commented 4 months ago

Instantiating copies of all operators for all supported data types increases the binary size and compile time of the library. For example I'm currently prototyping adding f16 support, and a naive implementation increased the size of the rten CLI tool by 300 KB / 11%.

For operators which merely move or copy data, such as Transpose, we only need different code based on the size of elements. A single instantiation could handle i32, u32 and f32 for example.

Operators which could use this optimization include:

robertknight commented 4 months ago

Maybe operations which care about bitwise equality with zero. For floats this is complicated by +0.0 vs -0.0.

I suppose we could implement eg. NonZero for a given bit width with shared code using a function which receives two different zero values as arguments. When used on types which only have a single zero value, these two values would be the same.