Initial version of reverse mode autodiff #227

apaszke · 2018-03-27T17:17:44Z

There are still a few things that could have been improved, but I think this can be done in subsequent PRs, and I wanted to get some feedback at this point. Some issues I see:

It might be nicer to link it as a part of lang. On the other hand, this needs tc2halide, which is part of the core, and that would create a dependency cycle between libtc_lang.so and libtc_core.so.
Not sure where the Python bindings should go. It's not really part of the engine, but I don't think we have a different pybind file that fits it better. I can create a new one if you want, just let me know.

prigoyal · 2018-03-27T20:40:41Z

yay :) @apaszke , I was wondering if you you give some high level idea of the approach for education purposes? Thanks for adding this :)

cc @abadams who might also have some idea on the autodiff related to Halide.

apaszke · 2018-03-27T21:02:53Z

I can of course answer some more specific questions, but most of this patch is basically a tiny bit of reverse mode AD code, and a ton of defensive programming to perform bookkeeping and shield it from unsupported features.

prigoyal · 2018-03-27T21:07:19Z

oh very cool, thanks for the reference @apaszke :) I'll look at that tutorial and then look over this PR.

prigoyal · 2018-03-27T21:20:49Z

test/test_autodiff.cc

+
+using namespace tc;
+
+void generalTest() {


I am curious about two things atm so far from looking at the test case:

how do we handle temporaries like the example https://facebookresearch.github.io/TensorComprehensions/framework/pytorch_integration/autograd_with_tc.html#reordering-grad-outputs ? would it be possible to demonstrate such use case with test example?

How is the inputs/outputs order handled for backwards?

Is it possible for users to specify that they want grad for only some inputs and not all?

Regarding 1 and 2:

The test case has temporaries, so you can actually see this in the example. In general the grad TC currently takes in:

inputs

outputs (including intermediates)

grads w.r.t. outputs (including intermediates)

and returns grads w.r.t. inputs. Inputs and outputs are ordered in the same way as they were specified in the original TC.

Regarding 3:

No, that's not currently possible, but is a very straightforward extension (only need a few extra lines of code).

abadams · 2018-03-30T21:15:50Z

The bulk of the work in the Halide autodiff was handling non-trivial indexing in ways that preserve parallelism. E.g. consider:

A(i) = 2 * B(some nasty expression in i)

You don't want to compute the derivatives as:

B'(some nasty expression in i) = A'(i) / 2

because you can't necessarily parallelize that over i, so it's going to be slow. This sort of thing comes up as soon as you have a stencil or broadcast (one input influences multiple outputs, when reversed, introduces a race condition).

Dealing with bounds and shapes of the computed derivatives was also interesting in the Halide work. I think that's simpler in TC, because tensors have finite size.

I think this PR is the right sort of approach, and we should do the rewrite of B(nasty) = A'(i) / 2 into B(j) = A'(nasty) / 2 as a later pass inside lowering, using the Halide solver or polyhedral tools.

apaszke · 2018-03-31T10:09:10Z

@abadams Thanks for the review! Great, it seems that we're on the same page! 😄

I agree that shifting the formulas to lhs isn't perfect, but it could easily be done as a post processing step. I know how to implement this, it's just that I'd need a linear equation solver (which is not available at this level from what I understand). For now I decided to bail and put the indexing expressions there, but it's likely to cause errors downstream anyway (I don't think it's supported in TC/Halide). It would be useful to support this in general too.

Bounds are another problem, and you can see that I had to implement a few checks to make sure that they can be still auto-inferred (usedIndexVars + requireAllIndexVarsOf). This is still not perfect because there are a few simple things like where ranges, that also cause an error for now, but are fairly easy to add later. I didn't want to complicate the initial patch so we can quickly get this in, and then work on improvements.

abadams · 2018-04-01T00:05:47Z

Equation solvers exist once you get down to Halide IR or polyhedral IR. It would be silly to implement yet another one at the TC-front-end-IR level. I'd just allow complex index expressions on the LHS for now, which I think is what you're doing? They may barf inside the tc2halide layer right now, but support for those is an outstanding near-term TODO - we need them for things like histogram computation.

apaszke · 2018-04-04T11:29:12Z

Yep, that's exactly what I've been thinking. Great that we're on the same page!

nicolasvasilache · 2018-05-19T18:30:05Z

src/core/autodiff.cc

+  throw std::runtime_error("Unknown Halide type");
+}
+
+void findAccessedTensors(


We have been implementing this type of functions as:

std::unordered_set<std::string> collectAccessedTensors(const TreeRef& tree) {...}

This way it can compose.
Generally we try not to pass out parameters as mutable in but just return the result.

nicolasvasilache · 2018-05-19T18:30:46Z

src/core/autodiff.cc

+void assertNoWriteAfterRead(Def def) {
+  std::unordered_set<std::string> read_only;
+  // Inputs are always read-only
+  for (Param input : def.params())


braces even around 1-liner for/if please

nicolasvasilache · 2018-05-19T18:35:09Z

src/core/autodiff.cc

+    auto lhs_name = comp.ident().name();
+    if (read_only.count(lhs_name) > 0)
+      throw std::runtime_error(
+          "AD not supported in TCs that write to a value after reading it");


So this would not work for FCRelu?

def fcrelu(float(N, M) I, float(M, P) W, float(P) Bias) -> (O) { O(n, p) +=! W(r_m, p) * I(n, r_m) O(n, p) = O(n, p) + B(p) }

Is this a fundamental limitation or a future TODO?

nicolasvasilache · 2018-05-19T18:44:22Z

src/core/autodiff.cc

+}
+
+void findAccessedTensors(
+    std::unordered_set<std::string>& read_only,


Also, why convert to string at this point, returning TreeRef seems fine to me?

nicolasvasilache · 2018-05-19T18:47:14Z

src/core/autodiff.cc

+// that can happen is that we will be too conservative and throw, so it's ok.
+std::unordered_set<std::string> usedIndexVars(Comprehension comp) {
+  std::unordered_set<std::string> index_vars;
+  for (Ident idx : comp.indices())


nicolasvasilache · 2018-05-19T18:49:31Z

src/core/autodiff.cc

+    return has_zero_grad_.count(name) > 0;
+  }
+  void markZeroGrad(const std::string& name) {
+    has_zero_grad_.count(name);


does this function do anything?

nicolasvasilache · 2018-05-19T18:51:02Z

test/lang_utils.h

@@ -0,0 +1,110 @@
+/**


This extraction is useful by itself and independent, please split in a separate commit

nicolasvasilache · 2018-05-19T18:52:26Z

src/core/autodiff.cc

+
+std::string differentiate(const std::string& source) {
+  // Parse and check the source
+  auto def = Def(Sema().checkFunction(Parser(source).parseFunction()));


What happens if the TC has more than 1 def?

nicolasvasilache · 2018-05-19T18:54:17Z

src/core/autodiff.cc

+
+using namespace lang;
+
+static const lang::SourceRange dummyRange{std::make_shared<std::string>(""),


no freestanding globals please, make it a static definition inside a function in an anon namespace and return that instead

nicolasvasilache · 2018-05-19T18:59:27Z

src/core/autodiff.cc

+  auto body = def.statements();
+  auto it = body.end();
+  if (it == body.begin())
+    throw std::runtime_error("empty body");


braces around if
append TC to the message?

nicolasvasilache · 2018-05-19T19:02:10Z

src/core/autodiff.cc

+  auto it = body.end();
+  if (it == body.begin())
+    throw std::runtime_error("empty body");
+  do {


can we use a for / rbegin / rend loop?

nicolasvasilache · 2018-05-19T19:02:36Z

src/core/autodiff.cc

+
+    int assign_kind = comp.assignment()->kind();
+    if (assign_kind != '=' && assign_kind != TK_PLUS_EQ_B &&
+        assign_kind != TK_PLUS_EQ)


nicolasvasilache · 2018-05-19T19:02:40Z

src/core/autodiff.cc

+        assign_kind != TK_PLUS_EQ)
+      throw ErrorReport(comp)
+          << "Only =, += and +=! assignments are supported in AD";
+    if (comp.whereClauses().size() > 0 || comp.equivalent().present())


nicolasvasilache · 2018-05-19T19:05:52Z

src/core/autodiff.cc

+  }
+}
+
+void findIndexVars(


Same comment as above re return values vs mutable reference as input.

nicolasvasilache · 2018-05-19T19:12:31Z

src/lang/test_expected/general_ad.expected

+def grad_mm(double(M, K) A, double(K, N) B, double(M, N) O, double(M, N) W, double(M, N) seed_d_O, double(M, N) seed_d_W) -> (d_A, d_B) {
+  d_O(i, j) = seed_d_d_O(i, j)
+  d_O(i, j) += (seed_d_W(i, j) * 2)
+  d_A(i, k) = 0


In order to be parsed by TC this will need either a where clause or to be folded into +=!.

More generally I'd make the test run parsing + sema on the produced TC and print / compare against the canonical TC. Dogfooding would have caught the where case here.

nicolasvasilache · 2018-05-19T19:17:17Z

src/core/autodiff.cc

+// XXX: Sema isn't nilpotent, so we have to reparse the source
+std::vector<TreeRef> inferOutputTypes(const std::string& source) {
+  auto halide_def =
+      tc2halide::translate(isl::with_exceptions::globalIslCtx(), source, true);


So this now makes lang depend on ISL which is unnecessary.
We should extract what we need from halide for range inference.

nicolasvasilache · 2018-05-19T19:30:05Z

src/core/autodiff.cc

+        ListView<TreeRef>::create(dummyRange, TreeList{}),
+        Compound::create(TK_OPTION, dummyRange, {}),
+        ListView<TreeRef>::create(dummyRange, TreeList{})));
+    if (usedIndexVars(Comprehension(grad_comps_.back())) !=


nicolasvasilache

Made a first quick pass, generally this look great, thanks for your contribution.
I made a few comments, I'll finish reviewing early next week.

I have one question regarding gathers, I did not seem to see a guard against those in AD (i.e. A(B(i))), can you point where such cases are caught and what we throw on them?

facebook-github-bot · 2018-05-25T13:19:02Z

Thank you for your pull request. We require contributors to sign our Contributor License Agreement, and yours has expired.

Before we can review or merge your code, we need you to email cla@fb.com with your details so we can update your status.

facebook-github-bot added the CLA Signed label Mar 27, 2018

apaszke force-pushed the autodiff_init_pr branch from 3a92daa to 95dcba6 Compare March 27, 2018 17:30

Initial version of reverse mode autodiff

01697c6

apaszke force-pushed the autodiff_init_pr branch from 95dcba6 to 01697c6 Compare March 27, 2018 17:55

prigoyal reviewed Mar 27, 2018

View reviewed changes

prigoyal requested review from zdevito and abadams March 30, 2018 20:56

nicolasvasilache reviewed May 19, 2018

View reviewed changes

src/core/autodiff.cc

}

}

void findIndexVars(

Copy link

Contributor

nicolasvasilache May 19, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comment as above re return values vs mutable reference as input.

nicolasvasilache reviewed May 19, 2018

View reviewed changes

zdevito removed request for zdevito and abadams March 2, 2020 21:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Initial version of reverse mode autodiff #227

Initial version of reverse mode autodiff #227

apaszke commented Mar 27, 2018

prigoyal commented Mar 27, 2018 •

edited

Loading

apaszke commented Mar 27, 2018 •

edited

Loading

prigoyal commented Mar 27, 2018

prigoyal Mar 27, 2018 •

edited

Loading

apaszke Mar 29, 2018 •

edited

Loading

abadams commented Mar 30, 2018

apaszke commented Mar 31, 2018

abadams commented Apr 1, 2018

apaszke commented Apr 4, 2018 •

edited

Loading

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache May 19, 2018

nicolasvasilache left a comment

facebook-github-bot commented May 25, 2018


		using namespace lang;

		static const lang::SourceRange dummyRange{std::make_shared<std::string>(""),

Initial version of reverse mode autodiff #227

Are you sure you want to change the base?

Initial version of reverse mode autodiff #227

Conversation

apaszke commented Mar 27, 2018

prigoyal commented Mar 27, 2018 • edited Loading

apaszke commented Mar 27, 2018 • edited Loading

prigoyal commented Mar 27, 2018

prigoyal Mar 27, 2018 • edited Loading

Choose a reason for hiding this comment

apaszke Mar 29, 2018 • edited Loading

Choose a reason for hiding this comment

Regarding 1 and 2:

Regarding 3:

abadams commented Mar 30, 2018

apaszke commented Mar 31, 2018

abadams commented Apr 1, 2018

apaszke commented Apr 4, 2018 • edited Loading

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

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

nicolasvasilache left a comment

Choose a reason for hiding this comment

facebook-github-bot commented May 25, 2018

prigoyal commented Mar 27, 2018 •

edited

Loading

apaszke commented Mar 27, 2018 •

edited

Loading

prigoyal Mar 27, 2018 •

edited

Loading

apaszke Mar 29, 2018 •

edited

Loading

apaszke commented Apr 4, 2018 •

edited

Loading