search

ahrefs / ocannl

OCANNL: OCaml Compiles Algorithms for Neural Networks Learning
BSD 2-Clause "Simplified" License
67 stars 2 forks source link

Enable inlining for virtual nodes with non-linear index symbols (diagonal and partially diagonal tensors) #133

Open lukstafi opened 1 year ago

lukstafi commented 1 year ago

In the initial implementation of virtual nodes, we require that the getters split into two classes: "recursive" ones share all indices with a setter, and regular caller ones are in a fully disjoint scope. We also require the setter to only have for-loop symbol indices. If there are more complex setters or getters, a node cannot be virtual.

We can extend virtual nodes to partial inlining cases. That's where a setter and a getter share some indices at same positions. If the shared index is a for-loop symbol, the setter and the getter share scope -- the inlined code is smaller than it would be for a disjoint-indices getter.

lukstafi commented 1 year ago

It's easy to handle the Fixed_idx and Task_id cases. Given that for now we don't have tensors defined locally in scope of other tensors' loops), what remains is non-linear symbols. Such use cases can happen when einsum is used to define (partially) diagonal tensors, and would be extra great to optimize since otherwise there's lots of wasted space.