Add a compiler and tracer, each creating OpPrograms

[fritzo]

Addresses pyro-ppl/pyro#2929
pair coded with @eb8680

This aims to work around backend jit issues and improve speed of Funsor used in Pyro. The approach is to first perform symbolic computations among Funsors, then lower to simple Funsor expressions, then compile the lowered Funsor expression to a straight-line program involving only backend ops, then optionally convert to Python code. The final program depends only on funsor.ops. This eliminates interpreter overhead, but does not eliminate op dispatch overhead.

The immediate application is to speed up Pyro's AutoGaussian guide pyro-ppl/pyro#2929, but this will also require symbolic Gaussians #556.

Tasks

• compile lowered funsors to Funsor-free programs
• lower Contraction to Binary
• trace simple op graphs
• support tuples in compile_function() for functions with multiple outputs (e.g. forward filter backward precondition).

Tasks deferred to follow-up PRs

• lower bound variables, maybe via Lambda or ops.einsum? E.g. to eliminate the bound variable i:

  Variable("f", Reals[3])["i"].reduce(ops.add, "i")  ⟶  Variable("f", Reals[3]).sum(0)

• fix contraction to use funsor.ops internally (not opt_einsum with direct backends), or make opt_einsum an op or sth
• fix Gaussian to use funsor.ops internally
• support tuples in trace_function() to support multiple outputs

Tested

• unit tests for compile_funsor()
• unit tests for trace_function()

[eb8680]

Looks great, and definitely feasible in some form. What about tracing individual funsor.Ops rather than Funsor terms/rewrite rules? The interface might a bit more awkward, but we could have a very low-overhead untyped graph representation that way.

[eb8680]

I guess that approach would require strict use of funsor.ops for every eager computation in every pattern, which is probably not realistic or desirable.

[fritzo]

What about tracing individual funsor.Ops

🤔 Interesting, I guess that would avoid both a lowering stage and the need to refactor Gaussian to be symbolic.

strict use of funsor.ops for every eager computation

Yes, I guess we'd need to manually desugar all math e.g. x + y --> ops.add(x, y), even in Gaussian, which IMO would make math less maintainable. It might actually be easier to refactor Gaussian to be symbolic. I wonder if funsor.syntax could help, since it desugars to use funsor.ops 😬

[eb8680]

We could also experiment with torch.fx for even lower level tracing, which would save us having to change a bunch of Funsor code. That would make this module specific to the PyTorch backend for now, but that doesn't seem so bad since it's currently our only use case.

Refactoring Gaussian seems worthwhile regardless of how we go about this, but would be lower priority.

[fritzo]

@eb8680 WDYT about merging this partial implementation (w/o lower() support for bound variables), and simultaneously working on all the solutions you proposed: this funsor.compiler (this PR and subsequent pair coding sessions), an op-level funsor.ops.trace module, and torch.fx (would this be a backend?)?

[eb8680]

Looks good per zoom review