This pull-request enables some simple automatic differentiation for thorin programs.
Overview
The most important new feature is the op_grad axiom in World. Lambdas can be passed to op_grad so that a new pass GradGen will replace it by a lambda that calculates the gradients.
Currently this only works for functions taking some reals and returning one real. There are no ifs or loops or function calls allowed yet, this will follow later. Also note that without -Othorin the grad axiom will stay unchanged in the thorin code.
The implementation is mainly based on this paper about AD for SSA form programs.
Example
Here is an example (from the mentioned paper) that already works:
fn example(a: f64, b: f64) -> f64 { a / (a + b * b) }
fn main() -> i32 {
// If the grad-gen works, we get da=1/9 and db=4/9
let (da, db) = grad(example)(-10.0, 5.0);
(db / da) as i32
}
The generated code for the gradients of example is:
Simple gradient generation
This pull-request enables some simple automatic differentiation for thorin programs.
Overview
The most important new feature is the
op_grad
axiom in World. Lambdas can be passed toop_grad
so that a new passGradGen
will replace it by a lambda that calculates the gradients.Currently this only works for functions taking some reals and returning one real. There are no ifs or loops or function calls allowed yet, this will follow later. Also note that without
-Othorin
thegrad
axiom will stay unchanged in the thorin code.The implementation is mainly based on this paper about AD for SSA form programs.
Example
Here is an example (from the mentioned paper) that already works:
The generated code for the gradients of
example
is:And in fact the inliner is able to generate this code for
main
(notice howdb/da
is replaced by4
):Future features
This merge-request only contains the most basic AD functionality. Features I plan to implement next are:
sin
,cos
, ...)if
expressions in the target lambdai32
)