[Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff #6078

yzhliu · 2020-07-16T15:45:28Z

Co-authored-by: Sergei Grechanik sergei.grechanik+h@gmail.com

This is the PR that aim to remove the intermediate large Jacobian tensor in te.autodiff, as well as to do some optimizations such as inline. We have made a series of PRs for https://discuss.tvm.ai/t/rfc-bring-in-tensor-expression-autodiff and this will be the last one.

@sgrechanik-h @MarisaKirisame @junrushao1994 @xqdan @tqchen Please take a look.

Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com>

src/te/autodiff/ad_simplify.cc

include/tvm/arith/int_solver.h

MarisaKirisame · 2020-07-29T04:06:42Z

I have read the paper and roughly understand how it work. will continue reviewing shortly after.

src/te/autodiff/ad_simplify.cc

MarisaKirisame · 2020-07-29T12:02:23Z

src/te/autodiff/ad_simplify.cc

+
+    // For Div we simply use the condition of the numerator.
+
+    if (nz_a.value.same_as(op->a)) {


refactor this into a function.

do you mean provide a function bool value_equals(const PrimExpr&) in NonzeroConditionResult?

return ReuseNZ(nz_a, op);

If this is the only place to use the function, do we really need to create a new function? correct me if I misunderstood.

MarisaKirisame · 2020-07-29T12:27:29Z

src/te/autodiff/ad_simplify.cc

+  return NonzeroConditionFunctor().NonzeroCondition(expr);
+}
+
+struct FactorOutAtomicFormulasResult {


perhaps you should store them in conjunctive normal form.

do you mean to change the struct & member name to mention CNF?

your approach is pure conjunction. I was suggesting maybe a disjunction of conjunction will allow more optimization ability. it is fine if you dont do it.

I see. I'll keep in mind and maybe change it later.

src/te/autodiff/ad_simplify.cc

MarisaKirisame · 2020-08-01T09:12:28Z

src/te/autodiff/ad_simplify.cc

+  return NonzeroConditionFunctor().NonzeroCondition(expr);
+}
+
+struct FactorOutAtomicFormulasResult {


your approach is pure conjunction. I was suggesting maybe a disjunction of conjunction will allow more optimization ability. it is fine if you dont do it.

src/te/autodiff/ad_simplify.cc

yzhliu · 2020-08-07T20:49:35Z

@MarisaKirisame @sergei-grechanik @tqchen I addressed most of the comments, please take a look again.

sergei-grechanik

Thank you @yzhliu, looks good to me. My only concern is that without performance testing the quality of the generated code may deteriorate with subsequent changes of the simplifiers, but that problem probably shouldn't be addressed in this PR.

yzhliu · 2020-08-10T19:12:59Z

@sergei-grechanik I agree. Perhaps we need performance integration test run periodically as mentioned in https://discuss.tvm.ai/t/efforts-on-benchmarking-for-tvm/

@tqchen could you also take a look?

tqchen · 2020-08-11T22:51:12Z

include/tvm/node/container.h

+   * \param other Map to be merged.
+   * @return The merged Array. Original Map is kept unchanged.
+   */
+  Map<K, V> Merge(const Map<K, V>& other) const {


Shall we make it as a global function instead of member? so it is not ambiguitous (that the result is a new map)

tqchen · 2020-08-11T22:51:41Z

include/tvm/runtime/container.h

+   * \param other Array to be concatenated.
+   * @return The concatenated Array. Original Array is kept unchanged.
+   */
+  Array<T> Concat(const Array<T>& other) const {


Consider make it as a global function, which also enables copy on write on the lhs(this)

yzhliu · 2020-08-12T22:14:20Z

@tqchen I make them static now.

tqchen · 2020-08-13T00:11:14Z

include/tvm/runtime/container.h

+   * \param lhs second Array to be concatenated.
+   * \return The concatenated Array. Original Arrays are kept unchanged.
+   */
+  static Array<T> Concat(Array<T> lhs, const Array<T>& rhs) {


We can make it a global template fucntion, so you can do Concat(lhs, rhs) without having to refer to the original type. Note that the type signature should be able to only run the function for the specific type so we won't risk over generalization

tqchen · 2020-08-13T03:26:05Z

include/tvm/node/container.h

+template <typename K, typename V,
+          typename = typename std::enable_if<std::is_base_of<ObjectRef, K>::value>::type,
+          typename = typename std::enable_if<std::is_base_of<ObjectRef, V>::value>::type>
+static Map<K, V> Merge(Map<K, V> lhs, const Map<K, V>& rhs) {


static->inline as it is in the header

@tqchen my bad I forgot to when copy-pasting. modified.

yzhliu · 2020-08-14T17:29:28Z

@tqchen it's ready

yzhliu · 2020-08-18T17:44:30Z

@tqchen kindly ping.

tqchen · 2020-08-18T23:19:53Z

Thanks @yzhliu @sergei-grechanik @MarisaKirisame !

…apache#6078) * [Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com> * fix lint * fix clang-format * add comments and magic number * clang-lint * address some comments * remove FreeVarsVisitor * fix constexpr lint * fix lint * fix lint * add Map.Merge * lint * change Array::Concat & Map::Merge to global functions * fix lint * move functions to global * static -> inline Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com>

yzhliu and others added 4 commits July 16, 2020 08:28

[Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff

7f6c5fc

Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com>

fix lint

5972033

Merge remote-tracking branch 'upstream/master' into opt-autodiff

018f67b

fix clang-format

2bd109c

MarisaKirisame requested changes Jul 19, 2020

View reviewed changes

src/te/autodiff/ad_simplify.cc Outdated Show resolved Hide resolved

src/te/autodiff/ad_simplify.cc Outdated Show resolved Hide resolved

src/te/autodiff/ad_simplify.cc Outdated Show resolved Hide resolved

yzhliu added 5 commits July 27, 2020 10:58

add comments and magic number

004ae0b

clang-lint

737cdf9

Merge remote-tracking branch 'upstream/master' into opt-autodiff

a3e6b7e

Merge remote-tracking branch 'upstream/master' into opt-autodiff

d158851

Merge remote-tracking branch 'upstream/master' into opt-autodiff

f9f4c18

tqchen requested changes Jul 29, 2020

View reviewed changes

include/tvm/arith/int_solver.h Outdated Show resolved Hide resolved

MarisaKirisame requested changes Jul 29, 2020

View reviewed changes

yzhliu added 2 commits July 31, 2020 23:27

address some comments

5951a94

Merge remote-tracking branch 'upstream/master' into opt-autodiff

914de04

MarisaKirisame reviewed Aug 1, 2020

View reviewed changes

sergei-grechanik reviewed Aug 2, 2020

View reviewed changes

src/te/autodiff/ad_simplify.cc Outdated Show resolved Hide resolved

src/te/autodiff/ad_simplify.cc Outdated Show resolved Hide resolved

src/te/autodiff/ad_simplify.cc Outdated Show resolved Hide resolved

src/te/autodiff/ad_simplify.cc Outdated Show resolved Hide resolved

yzhliu added 8 commits August 5, 2020 18:15

remove FreeVarsVisitor

69ea436

fix constexpr lint

e866720

Merge remote-tracking branch 'upstream/master' into opt-autodiff

d7178b3

fix lint

6596d7d

fix lint

e6c0b9b

add Map.Merge

79172d8

Merge remote-tracking branch 'upstream/master' into opt-autodiff

bb19a32

lint

9731d0f

MarisaKirisame approved these changes Aug 8, 2020

View reviewed changes

sergei-grechanik approved these changes Aug 10, 2020

View reviewed changes

yzhliu assigned tqchen Aug 11, 2020

tqchen requested changes Aug 11, 2020

View reviewed changes

yzhliu added 2 commits August 12, 2020 15:09

change Array::Concat & Map::Merge to global functions

47a5852

Merge remote-tracking branch 'upstream/master' into opt-autodiff

fd3a4a5

fix lint

4ede7fb

tqchen requested changes Aug 13, 2020

View reviewed changes

move functions to global

357510b

tqchen requested changes Aug 13, 2020

View reviewed changes

yzhliu added 2 commits August 12, 2020 22:14

static -> inline

c80ef74

Merge remote-tracking branch 'upstream/master' into opt-autodiff

e5746ed

tqchen approved these changes Aug 18, 2020

View reviewed changes

tqchen merged commit ff2a76f into apache:master Aug 18, 2020

ZihengJiang mentioned this pull request Sep 25, 2020

TVM v0.7 Release Note Candidate #6486

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff #6078

[Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff #6078

yzhliu commented Jul 16, 2020

MarisaKirisame commented Jul 29, 2020

MarisaKirisame Jul 29, 2020

yzhliu Jul 31, 2020

MarisaKirisame Aug 1, 2020

yzhliu Aug 7, 2020

MarisaKirisame Jul 29, 2020

yzhliu Aug 1, 2020

MarisaKirisame Aug 1, 2020

yzhliu Aug 7, 2020

MarisaKirisame Aug 1, 2020

yzhliu commented Aug 7, 2020

sergei-grechanik left a comment

yzhliu commented Aug 10, 2020

tqchen Aug 11, 2020

tqchen Aug 11, 2020

yzhliu commented Aug 12, 2020

tqchen Aug 13, 2020

tqchen Aug 13, 2020

yzhliu Aug 13, 2020

yzhliu commented Aug 14, 2020

yzhliu commented Aug 18, 2020

tqchen commented Aug 18, 2020


		// For Div we simply use the condition of the numerator.

		if (nz_a.value.same_as(op->a)) {

[Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff #6078

[Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff #6078

Conversation

yzhliu commented Jul 16, 2020

MarisaKirisame commented Jul 29, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yzhliu commented Aug 7, 2020

sergei-grechanik left a comment

Choose a reason for hiding this comment

yzhliu commented Aug 10, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yzhliu commented Aug 12, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yzhliu commented Aug 14, 2020

yzhliu commented Aug 18, 2020

tqchen commented Aug 18, 2020