ReWrite: compile-time operator re-writing

[tspooner]

Summary

This issue proposes a design for a rewrite feature in aegir which enables users to transform operator trees according to well-defined rules. The objective is to provide users with a way to "prune" extraneous branches (i.e. zeroed terms) or transform one sub-optimal expression into another that may be more efficient for constructing adjoints.

Overall, this addition is not somethign that we exhaustively implement on a first-pass. Instead, it is something that we work on over time, adding further optimizations as we develop. A limiter here is that the design proposals (see below) may/will rely on the nightly feature "impl specialisation" (see rust-lang/rfcs/pull/1210).

To-Do Tracker

• Implement rewrites.
• Add zero-pruning and evaluate performance.
• Perform cost-benefit analysis of compile-time and run-time.

Motivation

There are two motivations to this proposal:

1. Performance - this proposal should dramatically increase performance by reducing the gap between hand-crafted aegir trees, and those derived directly via Differentiable.
2. Type Explosion - simplification should help avoid the blow-up of operator trees that are common in symbolic-esque autograd frameworks without (common) subexpression caching.

Performance

Type Explosion

Design

Version 1

The base trait itself can be defined very easily as:

pub trait Simplify: Node {
    type Output: Node;

    fn simplify(self) -> Self::Output;
}

We then leverage impl specialisation to set the "default" behaviour to do nothing, just return self, and then let each operator define it's own unique overrides. The complicated part of this proposal comes in two flavours:

1. the optimizations themselves; and
2. in the mechanism used to perform pattern matching over the graph.

With impl specialisation, we should be able to implement the latter via "overloading." However, this does not afford the us, or the user, much control over the precedence rules. That is, Rust's specialisation algorithm will have complete control over the way in which different specialisations are applied. As such, we may need to extend the Simplify trait (see below) to allow some user-controlled variance.

Version 2

The alternative approach to this implementation is via a generic function on the Node trait, and a separate Rule trait for the actual concrete implementations. This model would cleanly separate rewrite rules from concrete types, which would be a good thing. For example, we could have something akin to:

pub trait Rule<N: Node> {
    type Output: Node;

    fn apply(&self, node: N) -> Self::Output;
}

pub trait Node {
    #[inline]
    fn simplify<R: Rule>(self, rule: R) -> R::Output {
        rule.apply(self)
    }

    ...
}