Skip to content

graph fuser inserts explicit expands where necessary#10325

Closed
zou3519 wants to merge 3 commits intopytorch:masterfrom
zou3519:pytorch-expand-adds
Closed

graph fuser inserts explicit expands where necessary#10325
zou3519 wants to merge 3 commits intopytorch:masterfrom
zou3519:pytorch-expand-adds

Conversation

@zou3519
Copy link
Contributor

@zou3519 zou3519 commented Aug 7, 2018

Fixes #10096

If the only thing preventing a simple mappable operator from being fused
into a fusion group is that its Tensor inputs are not of the same shape as the
output, then the graph fuser inserts explicit expand nodes for those
inputs.
This helps the graph fuser not miss out on any fusion opportunities
involving simple mappable operations that have Tensor inputs. This PR
doesn't do anything for the scalar case; that can be addressed later.

Test Plan

  • Simple expect test case
  • Added expect tests for a raw LSTMCell. The expands help speed up the
    forwards pass by allowing more operations to be fused into the LSTMCell's single
    FusionGroup.

cc @apaszke @zdevito

@zou3519 zou3519 added the oncall: jit Add this issue/PR to JIT oncall triage queue label Aug 7, 2018
@ssnl
Copy link
Collaborator

ssnl commented Aug 14, 2018

need rebase. mac seems broken as well

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.

In general, the approach makes sense but it needs some work to make the expand-generation logic less intrusive in the fusion pass. Right now it makes it increasingly difficult to make further edits to the pass. It might make sense to just establish when FusionGroups should be made and then do a followup pass to issue expands that will occur before the FusionGroup.

This comment was marked as off-topic.

This comment was marked as off-topic.

This comment was marked as off-topic.

This comment was marked as off-topic.

This comment was marked as off-topic.

@zou3519 zou3519 force-pushed the pytorch-expand-adds branch from 3453b27 to 943fcf2 Compare August 16, 2018 18:02
@zou3519
Copy link
Contributor Author

zou3519 commented Aug 16, 2018

I've made the following changes based on @zdevito's comments:

  • isFusable(Node) doesn't check input sizes anymore. A simple mappable
    operation with correct input/output types is fusible.
  • shouldFuse(consumer, producer) checks consumer and producer for
    matching map_size. If the sizes match among other checks,, then fuse(consumer, producer) is called.
  • fuse(consumer, producer) emits expand nodes as necessary to fuse producer and consumer.
  • tryMoveChunk(consumer, chunk) also emits expand nodes as necessary because
    isFusable(Node) doesn't check input sizes anymore. A chunk should only
    be moved if the move results in a node that can be fused with
    consumer; this is why tryMoveChunk needs to emit expand nodes.

This should be more robust than the previous approach, but I'm not sure if the expand-generation logic is still too intrusive or not.

In addition, @apaszke's upcoming refactor of the graph fuser probably will fix this exact problem, so I am also fine with punting on this PR until the refactor and seeing how the state of things are then.

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.

Minor errors that should be fixed before this lands, but generally LGTM.

This comment was marked as off-topic.

This comment was marked as off-topic.

This comment was marked as off-topic.

This comment was marked as off-topic.

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.

Fixes pytorch#10096

If the only thing preventing a simple mappable operator from being fused
into a fusion group is that its Tensor inputs are not of the same shape
as the
output, then the graph fuser inserts explicit expand nodes for those
inputs.
This helps the graph fuser not miss out on any fusion opportunities
involving simple mappable operations that have Tensor inputs.

The main changes are:
- isFusable(Node) doesn't check input sizes anymore. A simple mappable
  operation with correct input/output types is fusible.
- shouldFuse(consumer, producer) checks consumer and producer for
  matching map_size.
- fuse(consumer, producer) emits expand nodes as necessary
- tryMoveChunk(consumer, chunk) also emits expand nodes as necessary because
  isFusable(Node) doesn't check input sizes anymore. A chunk should only
  be moved if the move results in a node that can be fused with
  consumer.
@zou3519 zou3519 force-pushed the pytorch-expand-adds branch from 64148e8 to 61cc7d6 Compare August 17, 2018 21:24
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
Copy link
Contributor Author

zou3519 commented Aug 17, 2018

@pytorchbot retest this please

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[jit] Proposal: fuse add ops that expand tensors

6 participants