Rematerialization
- Dominant language
- C++
- Stars
- 1.9k
- Forks
- 283
- Avg merge
- 2d 10h
- Merged PRs (30d)
- 135
Description
In register allocation, rematerialization is an optimization in which variables with long lifetimes are recomputed at the use point to avoid register pressure. Analogously, in the XLS pipeline scheduler, if a node of bitwidth `b` is used at pipeline stage `y` is created at pipeline stage `x`, then we have to pay for `b * (y - x)` bits worth of registers to carry that value along the pipeline, so it sometimes makes sense to replace those registers with additional nodes in pipeline stage `y` that can recompute the value from other available values.
In general, there are a few ways to recompute such a value. In order of complexity:
1. The optimization only triggers if the value is some function of values that are already _available_ in pipeline stage `y`. This is how classical rematerialization works.
2. In addition to available values, we can use reverse computation (as described in #543) to get additional values from the available ones, albeit at additional cost in terms of nodes.
3. We can use values from pipeline stages that are later than `x`, at the cost of additional registers, as long as area is saved in total.
These techniques should be used in some currently unspecified way to generate a list of "optimization opportunities", each of which consists of an interstage edge to be replaced and the code that will be added to replace it. Once we have a list of possible rematerialization optimization opportunities, we should optimize the code for each one of them in isolation. It may also make sense to have optimizations that take into account the other opportunities, e.g.: rewriting the code added by opportunities to be more similar to one another in the hopes that the final CSE merges them together.
We now need to evaluate the opportunities for quality. The main factor involved is area: rematerialization removes register area but adds combinational area. Additionally, we probably do not want the maximum delay of the pipeline stages to increase after rematerialization, so we should prune optimization opportunities that pose a risk of causing the pipeline to not meet timing (note that this implies that a rematerialization pass should happen after scheduling, when we have access to such information).
Next, we need to choose some subset of these rematerialization opportunities to actually enact in the IR graph. Some of these choices are mutually exclusive with one another (e.g.: 1 vs 2 vs 3 above), whereas others are independent. We should first make the mutually exclusive choices using our notion of quality above. Then, we can choose among the independent ones using _submodular minimization_.
A [submodular function](https://en.wikipedia.org/wiki/Submodular_set_function) is a function from a set to the reals that has "diminishing returns" or "diminishing costs". Formally, this means that for any submodular `f : 2ᴱ → R` and `X ⊆ Y ⊆ E` and `x ∈ E \ Y`, `f(X ∪ {x}) - f(X) ≥ f(Y ∪ {x}) - f(Y)`. If you have a set of pieces of code `C` and you choose some subset of `C` and run common subexpression elimination on it, the size of the resulting code (as a function of the subset chosen) is a submodular function. Additionally, if you have a function `f : E → R` and you create a set function `g` by applying `f` to all elements of the subset and then summing the results, then `g` will be a _modular function_, which is defined the same way as submodularity except replacing `≥` with `=`. Modular functions are closed under negation and submodular functions are closed under nonnegative linear combination and monotone postcomposition.
We can now finally construct a cost function for our problem that looks like `f(S) = area(CSE(S)) - λ * sum(map(quality, S))` where `quality` is the measure of quality discussed earlier. This function is submodular since it is a linear combination of a submodular function (assuming `area` is monotone) and a modular function (all modular functions are submodular). Note that submodular optimization does not require expressing the submodular function in any particular form; it is a black-box optimization technique.
Submodular functions can be minimized exactly (globally) in `O(n log(n))` time with a [complicated algorithm](https://arxiv.org/abs/1610.09800), or approximately in `O(n²)` with a [very simple algorithm](http://proceedings.mlr.press/v28/iyer13-supp.pdf) that is similar to gradient descent (submodular functions are similar in many respects to convex ones). The simple algorithm requires a starting set, so it is amenable to randomization in case of poor results. To avoid running common subexpression elimination repeatedly during optimization, we can run it once with all the rematerialization code snippets, and each evaluation of the cost model consists of doing some depth first traversals of this graph to determine the set of nodes actually used by the given subset of snippets.
Finally, now that we have chosen some subset of rematerialization opportunities, we can enact them in the IR graph and run CSE or other optimization passes (so long as they are guaranteed not to affect the schedule).
Contributor guide
Assessment
This issue has not been assessed yet.