april-tools / cirkit

a python framework to build, learn and reason about probabilistic circuits and tensor networks
https://cirkit-docs.readthedocs.io/en/latest/
GNU General Public License v3.0
71 stars 1 forks source link

Implement extendible optimizer for the torch compiler #225

Closed loreloc closed 3 months ago

loreloc commented 4 months ago

There are three computational graph optimization opportunities in circuits:

  1. grouping multiple circuit layers into a single more efficient one, e.g., Kronecker + Dense -> Tucker.
  2. grouping multiple parameterization nodes into a single more efficient one, e.g., Kronecker over categorical's weights + Sum over all variables -> Einsum
  3. ungrouping circuit layers with particular parameterization into more than one layer that are more memory efficient, e.g., Dense parameterized by a Kronecker -> a composition of Dense

The key idea for implementing this is having: I. a registry of optimization rules, which can be possibly extended; II. a sub-graph matcher that grounds the optimization rules; III. an engine that orders the grounded optimization rules, applies them, and returns the optimized circuit.