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 |
|
Anton Korobeynikov
|
2f5492f572
|
Implement several peephole optimizations to unblock further optimizations of autodiff code (#60520)
* 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)))
* Remove broken disabled folding transform:
- It is disabled
- It is broken
- It is certainly not something mandatory
Co-authored-by: Dan Zheng <danielzheng@google.com>
|
2022-09-26 09:44:04 -07:00 |
|
