Skip to content

[jit] Fuse tensor-scalar ops when scalar is constant #10511

Closed
zou3519 wants to merge 8 commits intopytorch:masterfrom
zou3519:pytorch-fusescalar
Closed

[jit] Fuse tensor-scalar ops when scalar is constant #10511
zou3519 wants to merge 8 commits intopytorch:masterfrom
zou3519:pytorch-fusescalar

Conversation

@zou3519
Copy link
Contributor

@zou3519 zou3519 commented Aug 14, 2018 •
edited
Loading

This is on the way to resolving #9940.

Fixes #10501

This PR modifies graph fuser to fuse operations that have constant
scalar arguments. These constant scalar arguments are directly inlined
into the kernel body.

The context for this is that LSTM backward (in particular, sigmoid
backward) has many add(x, 1.) operations. This PR should be sufficient for
LSTM backward to get fused by the graph fuser.

cc @apaszke @zdevito

@zou3519 zou3519 added the oncall: jit Add this issue/PR to JIT oncall triage queue label Aug 14, 2018
zou3519
zou3519 commented Aug 14, 2018
View reviewed changes
test/test_jit.py Outdated

This comment was marked as off-topic.

Sign in to view

@zou3519 zou3519 force-pushed the pytorch-fusescalar branch from 2267979 to a235f9d Compare August 14, 2018 19:01
zdevito
zdevito reviewed Aug 15, 2018
View reviewed changes
Copy link
Contributor

@zdevito zdevito left a comment

Choose a reason for hiding this comment

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

This looks good. I have some small comments. I am surprised that we didn't need to modify the fusion_compiler, did we already have code to emit constants?

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

@zou3519
Copy link
Contributor Author

zou3519 commented Aug 15, 2018 •
edited
Loading

@zdevito Yes there was already code to emit constants inlined into the body of the FusionGroup. In particular, the graph_fuser inlines the alpha argument in add(Tensor, Tensor, Scalar alpha) into the FusionGroup if it is constant. I took this logic and expanded it to work with all constant number arguments.

@zou3519 zou3519 force-pushed the pytorch-fusescalar branch from 389d846 to 0616b4a Compare August 15, 2018 16:37
zou3519
zou3519 commented Aug 15, 2018
View reviewed changes
test/expect/TestScript.test_lstm_fusion_cuda-backward.expect Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

apaszke
apaszke reviewed Aug 16, 2018
View reviewed changes
Copy link
Contributor

@apaszke apaszke left a comment

Choose a reason for hiding this comment

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

Mostly LGTM, but I'd like to remove the hacky scalar checks that don't even try to see what overloads do we handle. Please use the matching syntax, or we'll end up with a lot of bugs again.

test/expect/TestScript.test_lstm_fusion_cuda-backward.expect Outdated

This comment was marked as off-topic.

Sign in to view

test/test_jit.py Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

zou3519
zou3519 commented Aug 17, 2018
View reviewed changes
test/expect/TestScript.test_lstm_fusion_cuda-backward.expect Outdated

This comment was marked as off-topic.

Sign in to view

zou3519
zou3519 commented Aug 17, 2018
View reviewed changes
torch/csrc/jit/register_prim_ops.cpp Outdated

This comment was marked as off-topic.

Sign in to view

facebook-github-bot
facebook-github-bot reviewed Aug 17, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

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

zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

zou3519
zou3519 commented Aug 17, 2018
View reviewed changes
torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

apaszke
apaszke approved these changes Aug 17, 2018
View reviewed changes
Copy link
Contributor

@apaszke apaszke left a comment

Choose a reason for hiding this comment

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

Looks great now! Some minor comments.

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/passes/graph_fuser.cpp Outdated

This comment was marked as off-topic.

Sign in to view

facebook-github-bot
facebook-github-bot reviewed Aug 17, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

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

zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

@zou3519 zou3519 force-pushed the pytorch-fusescalar branch 2 times, most recently from 221c95f to c79b24a Compare August 17, 2018 17:07
zou3519 added 8 commits August 17, 2018 10:38
@zou3519
Fuse tensor-scalar ops when scalar is constant
c86fe64 
This is on the way to resolving pytorch#9940.

This PR modifies graph fuser to fuse operations that have constant
scalar arguments. These constant scalar arguments are directly inlined
into the kernel body.

The context for this is that LSTM backward (in particular, sigmoid
backward) has many add(x, 1.) operations. This PR should be sufficient for
LSTM backward to get fused by the graph fuser.
@zou3519
Expect file updates
7c4c68f
@zou3519
Fix tests
36e4ae1
@zou3519
Fix expect files
3a4c0a5
@zou3519
Addressed comments:
c0f54e3 
- Use WithInsertPoint instead of insertAfter
- Make the compatible devices logic more explicit by adding a Device
  struct and a DeviceType enum. The possible Devices are
  `Unknown | AnyDevice | CPU | CUDA i`
@zou3519
Address comments:
d437193 
- Use ->matches instead of hacky allowing-numbers-in-type-checks
- Rewrite bool compatibleDevices(Node * consumer, Value * producer) to
  be more readable.
@zou3519
Fix mul/div
b4b1fec
@zou3519
Address comments
9d943a6
@zou3519 zou3519 force-pushed the pytorch-fusescalar branch from c79b24a to 9d943a6 Compare August 17, 2018 17:39
facebook-github-bot
facebook-github-bot reviewed Aug 17, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

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

zou3519 is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

@ezyang ezyang added the merged label Jun 26, 2019
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@zdevito zdevito zdevito left review comments

@facebook-github-bot facebook-github-bot facebook-github-bot left review comments

@colesbury colesbury Awaiting requested review from colesbury

@ezyang ezyang Awaiting requested review from ezyang

@gchanan gchanan Awaiting requested review from gchanan

@soumith soumith Awaiting requested review from soumith

+1 more reviewer

@apaszke apaszke apaszke approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

oncall: jit Add this issue/PR to JIT oncall triage queue

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[jit] Fuse sub(Tensor, Tensor, Scalar) nodes

5 participants

@zou3519 @zdevito @apaszke @facebook-github-bot @ezyang