[inductor] BatchLinearLHSFusion: also match torch._C._nn.linear#3232
Open
reger-men wants to merge 1 commit into
Open
[inductor] BatchLinearLHSFusion: also match torch._C._nn.linear#3232reger-men wants to merge 1 commit into
reger-men wants to merge 1 commit into
Conversation
|
Jenkins build for 1bc8980a52488cd6e77abdbf7c0dab34272e910f commit finished as FAILURE |
1bc8980 to
aed5e1b
Compare
|
Jenkins build for aed5e1bb3da2b806c42c10dc85a62dc5301b5f48 commit finished as FAILURE |
aed5e1b to
d763082
Compare
Dynamo inlines torch.nn.functional.linear into torch._C._nn.linear on hot paths, so BatchLinearLHSFusion.match silently misses the inlined form. This one-line change extends CallFunctionVarArgs to also accept torch._C._nn.linear. The sibling fusion class PreGradBatchLinearFusion in the same file already accepts both call targets (line 626 in current main); this brings BatchLinearLHSFusion in line with that. The fusion is gated by pre_grad_fusion_options["batch_linear_lhs"] (unset by default), so this is invisible to users who haven't opted in. is_linear_node_can_be_fused still filters out shapes that can't actually be fused (different weight shapes per head, etc). Adds a regression test that builds a module out of nn.Linear modules (Dynamo inlines those to _C._nn.linear). Without this change the test expects 0 fires; with it, 1 fire.
|
Jenkins build for d7630821e00abbdc75ac9c8ceab013feab3772a5 commit finished as FAILURE |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dynamo inlines
torch.nn.functional.linearintotorch._C._nn.linearon hot paths, soBatchLinearLHSFusion.matchsilently misses the inlined form. This one-line change extendsCallFunctionVarArgsto also accepttorch._C._nn.linear.The sibling fusion class
PreGradBatchLinearFusionin the same file already accepts both call targets (line 626 in current main); this bringsBatchLinearLHSFusionin line with that.The fusion is gated by
pre_grad_fusion_options["batch_linear_lhs"](unset by default), so this is invisible to users who haven't opted in.is_linear_node_can_be_fusedstill filters out shapes that can't actually be fused (different weight shapes per head, etc).Test plan
test_batch_linear_lhs_fusion_nn_linear_inlinedbuilds an nn.Linear module list (Dynamo inlines those to_C._nn.linear); without this change the matcher fires 0 times, with it 1 timetest_batch_linear_lhs_fusioncontinues to pass (regression for theF.linearform)compare_gradients/compare_parameters