thorin::plug::autodiff::AutoDiffEval Class Reference

This pass is the heart of AD. More...

#include <thorin/plug/autodiff/pass/autodiff_eval.h>

Inheritance diagram for thorin::plug::autodiff::AutoDiffEval:

Public Member Functions
	AutoDiffEval (PassMan &man)
Ref	rewrite (Ref) override
	Detect autodiff calls.
Ref	derive (Ref)
	Acts on toplevel autodiff on closed terms:
Ref	derive_ (Ref)
	Additionally to the derivation, the pullback is registered and the maps are initialized.
Ref	augment (Ref, Lam *, Lam *)
	Applies to (open) expressions in a functional context.
Ref	augment_ (Ref, Lam *, Lam *)
	Rewrites the given definition in a lambda environment.
Ref	augment_var (const Var *, Lam *, Lam *)
	helper functions for augment
Ref	augment_lam (Lam *, Lam *, Lam *)
Ref	augment_extract (const Extract *, Lam *, Lam *)
Ref	augment_app (const App *, Lam *, Lam *)
Ref	augment_lit (const Lit *, Lam *, Lam *)
Ref	augment_tuple (const Tuple *, Lam *, Lam *)
Ref	augment_pack (const Pack *pack, Lam *f, Lam *f_diff)
Public Member Functions inherited from thorin::RWPass< AutoDiffEval, Lam >
	RWPass (PassMan &man, std::string_view name)
bool	inspect () const override
	Should the PassMan even consider this pass?
Lam *	curr_mut () const
Public Member Functions inherited from thorin::Pass
	Pass (PassMan &, std::string_view name)
virtual	~Pass ()=default
World &	world ()
PassMan &	man ()
const PassMan &	man () const
std::string_view	name () const
size_t	index () const
virtual Ref	rewrite (const Var *var)
virtual Ref	rewrite (const Proxy *proxy)
virtual undo_t	analyze (Ref)
virtual undo_t	analyze (const Var *)
virtual undo_t	analyze (const Proxy *)
virtual bool	fixed_point () const
virtual void	enter ()
	Invoked just before Pass::rewriteing PassMan::curr_mut's body.
virtual void	prepare ()
	Invoked once before entering the main rewrite loop.
const Proxy *	proxy (Ref type, Defs ops, u32 tag=0)
const Proxy *	isa_proxy (Ref def, u32 tag=0)
	Check whether given def is a Proxy whose Proxy::pass matches this Pass's IPass::index.
const Proxy *	as_proxy (Ref def, u32 tag=0)

Detailed Description

This pass is the heart of AD.

We replace an autodiff fun call with the differentiated function.

Definition at line 11 of file autodiff_eval.h.

Constructor & Destructor Documentation

AutoDiffEval()

thorin::plug::autodiff::AutoDiffEval::AutoDiffEval ( PassMan &  man )

inline

Definition at line 13 of file autodiff_eval.h.

Member Function Documentation

augment()

Ref thorin::plug::autodiff::AutoDiffEval::augment	(	Ref	def,
		Lam *	f,
		Lam *	f_diff
	)

Applies to (open) expressions in a functional context.

Returns the rewritten expressions and augments the partial and modular pullbacks. The rewrite is identity on the term up to renaming of variables. Otherwise, only pullbacks are added. To do so, some calls (e.g. axioms) are replaced by their derivatives. This transformation can be seen as an augmentation with a dual computation that generates the derivatives.

Definition at line 15 of file autodiff_eval.cpp.

References augment_().

Referenced by augment_app(), augment_extract(), augment_lam(), augment_pack(), augment_tuple(), and derive_().

augment_()

Ref thorin::plug::autodiff::AutoDiffEval::augment_	(	Ref	def,
		Lam *	f,
		Lam *	f_diff
	)

Rewrites the given definition in a lambda environment.

Definition at line 317 of file autodiff_rewrite_inner.cpp.

References augment_app(), augment_extract(), augment_lam(), augment_lit(), augment_pack(), augment_tuple(), augment_var(), thorin::plug::autodiff::autodiff_type_fun(), thorin::World::external(), thorin::find_and_replace(), thorin::Pass::index(), thorin::Def::isa_mut(), thorin::Def::node_name(), thorin::World::sym(), thorin::Def::type(), thorin::Pass::world(), and thorin::Def::world().

Referenced by augment(), and augment_pack().

augment_app()

Ref thorin::plug::autodiff::AutoDiffEval::augment_app	(	const App *	app,
		Lam *	f,
		Lam *	f_diff
	)

Definition at line 209 of file autodiff_rewrite_inner.cpp.

References thorin::World::app(), thorin::App::arg(), augment(), thorin::App::callee(), thorin::compose_cn(), thorin::World::debug_dump(), thorin::Pi::isa_basicblock(), thorin::Pi::isa_cn(), thorin::World::mut_lam(), thorin::plug::direct::op_cps2ds_dep(), thorin::Lam::set(), thorin::Def::type(), thorin::World::type(), thorin::Def::var(), thorin::Pass::world(), and thorin::Def::world().

Referenced by augment_().

augment_extract()

Ref thorin::plug::autodiff::AutoDiffEval::augment_extract	(	const Extract *	ext,
		Lam *	f,
		Lam *	f_diff
	)

Definition at line 87 of file autodiff_rewrite_inner.cpp.

References augment(), thorin::World::call(), thorin::World::extract(), thorin::Pass::index(), thorin::Extract::index(), thorin::World::insert(), thorin::World::mut_lam(), thorin::plug::autodiff::pullback_type(), thorin::Lam::set(), thorin::Def::set(), thorin::Extract::tuple(), thorin::Def::type(), thorin::World::var(), thorin::Pass::world(), and thorin::Def::world().

Referenced by augment_().

augment_lam()

Ref thorin::plug::autodiff::AutoDiffEval::augment_lam	(	Lam *	lam,
		Lam *	f,
		Lam *	f_diff
	)

Definition at line 28 of file autodiff_rewrite_inner.cpp.

References augment(), thorin::plug::autodiff::autodiff_type_fun(), thorin::Lam::body(), thorin::World::call(), thorin::Pi::dom(), thorin::Lam::filter(), thorin::Lam::isa_basicblock(), thorin::World::mut_con(), thorin::plug::autodiff::pullback_type(), thorin::Def::sym(), thorin::Lam::type(), thorin::Def::var(), thorin::Pass::world(), thorin::Def::world(), and thorin::plug::autodiff::zero_pullback().

Referenced by augment_().

augment_lit()

Ref thorin::plug::autodiff::AutoDiffEval::augment_lit	(	const Lit *	lit,
		Lam *	f,
		Lam *
	)

Definition at line 15 of file autodiff_rewrite_inner.cpp.

References thorin::plug::autodiff::zero_pullback().

Referenced by augment_().

augment_pack()

Ref thorin::plug::autodiff::AutoDiffEval::augment_pack	(	const Pack *	pack,
		Lam *	f,
		Lam *	f_diff
	)

Definition at line 166 of file autodiff_rewrite_inner.cpp.

References thorin::World::annex(), thorin::World::app(), thorin::Def::arity(), thorin::World::arr(), augment(), augment_(), thorin::Pack::body(), thorin::World::extract(), thorin::World::mut_lam(), thorin::World::mut_pack(), thorin::plug::direct::op_cps2ds_dep(), thorin::World::pack(), thorin::plug::autodiff::pullback_type(), thorin::Pack::set(), thorin::Lam::set(), thorin::plug::autodiff::tangent_type_fun(), thorin::Def::type(), thorin::Pass::world(), and thorin::Def::world().

Referenced by augment_().

augment_tuple()

Ref thorin::plug::autodiff::AutoDiffEval::augment_tuple	(	const Tuple *	tup,
		Lam *	f,
		Lam *	f_diff
	)

Definition at line 130 of file autodiff_rewrite_inner.cpp.

References augment(), thorin::World::mut_lam(), thorin::plug::autodiff::op_sum(), thorin::Def::projs(), thorin::plug::autodiff::pullback_type(), thorin::Lam::set(), thorin::Def::set(), thorin::plug::autodiff::tangent_type_fun(), thorin::World::tuple(), thorin::Def::type(), thorin::World::var(), thorin::Pass::world(), and thorin::Def::world().

Referenced by augment_().

augment_var()

Ref thorin::plug::autodiff::AutoDiffEval::augment_var	(	const Var *	var,
		Lam *	,
		Lam *
	)

helper functions for augment

Definition at line 21 of file autodiff_rewrite_inner.cpp.

Referenced by augment_().

derive()

Ref thorin::plug::autodiff::AutoDiffEval::derive

(

Ref

def

)

Acts on toplevel autodiff on closed terms:

• Replaces lambdas, operators with the appropriate derivatives.
• Creates new lambda, calls associate variables, init maps, calls augment.

Definition at line 21 of file autodiff_eval.cpp.

References derive_().

Referenced by rewrite().

derive_()

Ref thorin::plug::autodiff::AutoDiffEval::derive_

(

Ref

def

)

Additionally to the derivation, the pullback is registered and the maps are initialized.

Definition at line 7 of file autodiff_rewrite_toplevel.cpp.

References thorin::Def::as_mut(), augment(), thorin::plug::autodiff::autodiff_type_fun_pi(), thorin::plug::autodiff::id_pullback(), thorin::World::mut_lam(), thorin::Lam::set(), thorin::Def::set(), thorin::World::tuple(), thorin::Def::type(), thorin::Def::var(), thorin::Pass::world(), thorin::Def::world(), and thorin::plug::autodiff::zero_pullback().

Referenced by derive().

rewrite()

Ref thorin::plug::autodiff::AutoDiffEval::rewrite ( Ref  def )

overridevirtual

Detect autodiff calls.

Reimplemented from thorin::Pass.

Definition at line 27 of file autodiff_eval.cpp.

References derive(), and thorin::Pass::world().

---

The documentation for this class was generated from the following files:

• include/thorin/plug/autodiff/pass/autodiff_eval.h
• src/thorin/plug/autodiff/auxiliary/autodiff_rewrite_inner.cpp
• src/thorin/plug/autodiff/auxiliary/autodiff_rewrite_toplevel.cpp
• src/thorin/plug/autodiff/pass/autodiff_eval.cpp