Anton Korobeynikov a5e83817b2 Reapply "Implement several peephole optimizations to unblock further optimizations of autodiff code" with correctness fix (#62012) ...

* Implement several peephole optimizations to unblock further optimizations of autodiff code

1. Simplify differentiable_function_extract of differentiable_function.
Before:
%x = differentiable_function(%orig, %jvp, %vjp)
%y = differentiable_function_extract [original] %x
After:
%y = %orig

2. Push conversion instructions inside of differentiable_function.
This unblocks inlining and specialization.
Before:
%x = differentiable_function(%orig, %jvp, %vjp)
%y = convert_escape_to_noescape %x
After:
%orig' = convert_escape_to_noescape %orig
%jvp' = convert_escape_to_noescape %jvp
%vjp' = convert_escape_to_noescape %vjp
%y = differentiable_function(%orig', %jvp', %vjp')

3. Another peephole is needed for reordering function conversion instructions to enable full inlining:
(convert_escape_to_noescape (convert_function (thin_to_thick_function x)))
=>
(convert_escape_to_noescape (thin_to_thick_function (convert_function x)))

Co-authored-by: Dan Zheng <danielzheng@google.com>

2022-11-16 23:21:27 -08:00

4.3 KiB

Raw Blame History

View Raw