[jit] dispatch and expose linear op #20039

wanchaol · 2019-05-02T00:46:39Z

Stack from ghstack:

[jit] make linear avaiable in tracing #20284 make linear avaiable in tracing
[jit] dispatch and expose linear op #20039 dispatch and expose linear op
[jit] split canonicalize_ops, make a decompose pass #19988 split canonicalize_ops, make a decompose pass
[jit] lower batchmm to non-diff optimization #19987 lower batchmm to non-diff optimization

Summary:
This expose the linear directly when we called the functional interface, the ATen linear op already do the same thing with the functional interface, so there's really no need to duplicate the code. Also, this will expose the higher level aten::linear op up until the custom fusion, so that different backends could know the high level information

Test Plan:
Test the linear op is correctly decomposed in the decomposition pass, which also means it does not get decomposed before that pass.

Currently being blocked from landing by #19769 and #20734.

Differential Revision: D15190354

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

bwasti

generally looks good, one concern re shape inference

bwasti · 2019-05-03T20:35:22Z

torch/csrc/jit/passes/shape_analysis.cpp

      }
+    } else if (node->matches("aten::linear(Tensor input, Tensor weight, Tensor? bias) -> Tensor")) {
+      if (auto type = input_type(0)) {
+        node->output()->setType(type);


is type here a complete tensor or just a dimensioned tensor?

the input shape is not the same as the output shape right?

hmm you are right, actually it should be a complete tensor and the output shape should be different. seems like bilinear below is also wrong, will fix it.

…r op" dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

…atch and expose linear op" dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

zdevito

I don't see any problems here, but because this shuffles around where decompositions happen, I am concerned about test coverage and performance.

Can we make sure the derivative formula for linear is actually being testing
@apaszke can you take a look at this and see if anything jumps out as bugs? Context is that mkldnn backend has good fused linear performance and we are splitting the op up before it can be seen.

apaszke

We should check the AD formulas
This likely breaks a lot of the fusions we used to do. We want to make sure that the bias additions in LSTM get fused, so the fuser needs to be taught to decompose addmms.

apaszke · 2019-05-07T10:44:43Z

torch/csrc/jit/passes/shape_analysis.cpp

+      SHAPE_ASSERT(weight_type->sizes().size() == 2 && sizes.size() >= 2);
+      sizes.at(last_dim) = weight_type->sizes()[0];
+      node->output()->setType(input_type->withSizes(sizes));
+      return true;


I don't think we need to add complete shape prop, since we're not really using it at this point. Also, this looks like the same case as for aten::mm, so maybe we could avoid duplicating them?

aten::mm section requires the tensor inputs to be 2, where aten::linear have 3.. I think I can also do partial shape prop like the aten::bilinear defined here, but this also requires additional code (which is a bit duplicated)..

Edit: do you mean the complete shape prop of aten::mm defined here? they are slightly different semantics in the formula: linear(input, weight) = input.mm(weight.t())

Fair enough, let's leave it.

apaszke · 2019-05-07T10:45:24Z

torch/csrc/jit/passes/decompose_ops.cpp

+      int ndim = input_type->dim();
+      Value* new_output = nullptr;
+      if (ndim == 2 && bias->type()->isSubtypeOf(TensorType::get())) {
+        // if ndim == 2 and bias is statically defined, dispatch to addmm decomposition


statically defined is not very well defined here. You're only checking that we can refine its type, which doesn't tell you much.

This is something I steal from the fuser code, where it tells you if it is statically defined (statically undefined needs more check though). because we are only check the input value types, I guess it would be sufficient to distinguish it.

The next step after this PR, I will move the batchnorm and layernorm decomposition to here as well, then I can use the util function defined above.

My point is that the term statically defined is not used correctly. I'm not even 100% sure what does that mean

From my point of view, statically defined tensor argument means that the value passed in is neither a None constant nor a Optional[Tensor] type, which we could statically know the type passed in rather than a dynamically unknown optional type. So isSubtypeOf(TensorType) should guarantee this definition

wanchaol · 2019-05-07T18:15:49Z

Thanks @zdevito @apaszke ! re concerns on ad formula and performance:

the ad formula is already been checked here, it's checking the input.dim() == 2 pathway, I will add one more test case to cover the matmul AD pathway
I believe by putting the decomposition pass before fuser and the PR added the decomposition from linear to addmm, it should preserve the old behavior as we did the bias add fusion, I checked the forward and backward graph for lstm, it remains the same, so it should be good on fusion and not regressed the LSTM performance. I will make sure to run enough benchmarks before landing this.

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

…it] dispatch and expose linear op" dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

wanchaol · 2019-05-08T09:28:27Z

I checked that the LSTM graph is the same before and after this PR, except the unique names of values are different:

https://gist.github.com/wanchaol/bcd988809a192c56f55319fc6a305637/revisions

Performance: shows no noticeable difference between master and master + this PR:

On Master:

Master + this stacked PR:

So I think this PR is not regressed our current fusion :)

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

apaszke · 2019-05-09T10:03:23Z

Ok I see you have made changes to how the ops are decomposed, so I'd need to read the prior patches to make sure this works ok. If you double checked that we're not loosing LSTM fusion + mm batching, then it might be good to go. I'll try to catch up with the whole stack soon.

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

XiaobingSuper · 2019-07-16T04:49:12Z

@wanchaol, Do you have any progress about this PR?

wanchaol · 2019-07-16T20:05:08Z

@XiaobingSuper this is currently blocked from landing by issue #19769 and #20734, we don't have a good solution to those currently, so this might not be land until we could solve them.

XiaobingSuper · 2019-12-12T08:09:09Z

@wanchaol , #19769 seems fixed, do you have any plan to do it?

wanchaol · 2019-12-12T23:04:05Z

@wanchaol , #19769 seems fixed, do you have any plan to do it?

@XiaobingSuper sorry the issue was closed to dedup, but the underlying problem is not fixed yet, which is #20734.

But I think we recently figured out a plan on how to fix it and I will let you know about the progress once we start working on it :)

dispatch and expose linear op

8f8c93c

pytorchbot added oncall: jit Add this issue/PR to JIT oncall triage queue module: internals Related to internal abstractions in c10 and ATen module: nn Related to torch.nn module: pybind Related to our Python bindings / interactions with other Python libraries labels May 2, 2019

This was referenced May 2, 2019

[jit] lower batchmm to non-diff optimization #19987

Closed

[jit] split canonicalize_ops, make a decompose pass #19988

Closed

move batchnorm and layernorm fusion to decomposition #19989

Closed

wanchaol added 2 commits May 1, 2019 17:48

Update on "dispatch and expose linear op"

5107dce

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

Update on "dispatch and expose linear op"

999ddc5

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

wanchaol requested review from apaszke, bddppq, bwasti and zdevito May 2, 2019 01:10

wanchaol added 4 commits May 1, 2019 18:46

replace inlineGraph on "dispatch and expose linear op"

9358e0f

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

Update on "dispatch and expose linear op"

624af35

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

update callsite to linear on "dispatch and expose linear op"

f2a5f28

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

small update on test on "dispatch and expose linear op"

a2a93b2

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

wanchaol changed the title ~~dispatch and expose linear op~~ [jit] dispatch and expose linear op May 2, 2019

wanchaol added 2 commits May 2, 2019 19:12

fix on unspecialized input on "[jit] dispatch and expose linear op"

2dfef98

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

add more test case on "[jit] dispatch and expose linear op"

9e053e0

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

bwasti approved these changes May 3, 2019

View reviewed changes

wanchaol added 7 commits May 5, 2019 21:15

reorg code, add linear AD formula on "[jit] dispatch and expose linea…

a3934db

…r op" dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

Update on "[jit] dispatch and expose linear op"

228c94d

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

Update on "[jit] dispatch and expose linear op"

944147b

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

update shape prop on "[jit] dispatch and expose linear op"

8e7979d

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

Update on "[jit] dispatch and expose linear op"

4a5080b

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

Update on "[jit] dispatch and expose linear op"

2342cf9

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

update ad formula, add constructor to compilation unit on "[jit] disp…

6c9a152

…atch and expose linear op" dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

zdevito reviewed May 7, 2019

View reviewed changes

apaszke reviewed May 7, 2019

View reviewed changes

wanchaol added 4 commits May 7, 2019 19:25

address comments on "[jit] dispatch and expose linear op"

dbf478e

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

fix tests on "[jit] dispatch and expose linear op"

87d76eb

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

change profiled shapes since we directly call the cpp function on "[j…

5ea5a48

…it] dispatch and expose linear op" dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

minor update on "[jit] dispatch and expose linear op"

aed834f

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

Update on "[jit] dispatch and expose linear op"

6bd430c

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

wanchaol mentioned this pull request May 8, 2019

[jit] make linear avaiable in tracing #20284

Closed

wanchaol added 2 commits May 8, 2019 13:42

Update on "[jit] dispatch and expose linear op"

80787f4

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

small AD optimize on "[jit] dispatch and expose linear op"

76a1516

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

add onnx symbolic for linear on "[jit] dispatch and expose linear op"

f65b756

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

wanchaol mentioned this pull request May 15, 2019

[jit] Autodiff bug on double backward for early expiration of grad_accumulator #19769

Closed

rebase with master on "[jit] dispatch and expose linear op"

e20ea0b

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

pytorchbot added the module: autograd Related to torch.autograd, and the autograd engine in general label Jun 4, 2019

Update on "[jit] dispatch and expose linear op"

f7ca26e

dispatch and expose linear op gh-metadata: pytorch pytorch 20039 gh/wanchaol/4/head

XiaobingSuper mentioned this pull request Jul 16, 2019

Add aten mkldnn backward ops: relu, linear and reshape #20570

Closed

6 tasks

facebook-github-bot added the cla signed label Oct 30, 2020

skyline75489 mentioned this pull request Jan 25, 2021

Cleanup git branches #41232

Open

wanchaol closed this Nov 18, 2021

facebook-github-bot deleted the gh/wanchaol/4/head branch December 19, 2021 15:17

[jit] dispatch and expose linear op #20039

[jit] dispatch and expose linear op #20039

Uh oh!

Conversation

wanchaol commented May 2, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bwasti left a comment

Choose a reason for hiding this comment

Uh oh!

bwasti May 3, 2019

Choose a reason for hiding this comment

Uh oh!

wanchaol May 3, 2019

Choose a reason for hiding this comment

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

apaszke left a comment

Choose a reason for hiding this comment

Uh oh!

apaszke May 7, 2019

Choose a reason for hiding this comment

Uh oh!

wanchaol May 8, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

apaszke May 9, 2019

Choose a reason for hiding this comment

Uh oh!

apaszke May 7, 2019

Choose a reason for hiding this comment

Uh oh!

wanchaol May 8, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

apaszke May 9, 2019

Choose a reason for hiding this comment

Uh oh!

wanchaol May 9, 2019

Choose a reason for hiding this comment

Uh oh!

wanchaol commented May 7, 2019

Uh oh!

wanchaol commented May 8, 2019

Uh oh!

apaszke commented May 9, 2019

Uh oh!

XiaobingSuper commented Jul 16, 2019

Uh oh!

wanchaol commented Jul 16, 2019

Uh oh!

XiaobingSuper commented Dec 12, 2019

Uh oh!

wanchaol commented Dec 12, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

wanchaol commented May 2, 2019 •

edited

Loading

wanchaol May 8, 2019 •

edited

Loading

wanchaol May 8, 2019 •

edited

Loading