[Dynamo] Implement graph region tracking for deduplication #141381

mlazos · 2024-11-22T21:28:54Z

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows:

GraphRegionTracker tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes.

Once graph capture is complete, these nodes are organized into region groups. A region group looks like this:
[[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The get_identical_regions function is the main entry point which will be used by the graph replacement algorithm in #141383

Edge cases to add more testing for in future PRs (in progress):

~~multiple nodes on the same line~~ (implemented)
~~dynamic shapes checking (need to verify symbolic inputs are the same across subgraphs)~~ (implemented)
ensure we don't expand regions where it will create a cycle during subgraph replacement
ensure outputs are always tensors (or tuples of tensors iirc)
~~out of order kwargs, unevenly nested kwargs~~ (implemented)
input aliasing - TBD, we may add support for this in invoke_subgraph or reuse the aliasing analysis here to not form regions with these properties
~~all global state~~ (implemented)

Other followups:

consolidate global state checking across all caching infra

Stack from ghstack (oldest at bottom):

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames

Fixes for bfs Initial region tracking tests hashing + test fixes fix for hash more tests more fixes region tracker updates Update tests Update tests2 [ghstack-poisoned]

pytorch-bot · 2024-11-22T21:28:59Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/141381

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit f475866 with merge base 117b6c3 ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

Lint / workflow-checks / linux-job (gh) (#142485)
##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 [ghstack-poisoned]

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 Edge cases to add more testing for in future PRs (in progress): * multiple nodes on the same line * dynamic shapes banning * ensure we don't expand regions where it will create a cycle during subgraph replacement * ensure outputs are always tensors * out of order kwargs, unevenly nested kwargs [ghstack-poisoned]

torch/_dynamo/graph_region_tracker.py

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 Edge cases to add more testing for in future PRs (in progress): * multiple nodes on the same line * dynamic shapes checking (need to verify symbolic inputs are the same across subgraphs) * ensure we don't expand regions where it will create a cycle during subgraph replacement * ensure outputs are always tensors (or tuples of tensors iirc) * out of order kwargs, unevenly nested kwargs [ghstack-poisoned]

This PR implements the deduplication pass (blocked by config currently) for dynamo where identical regions from #141381 are replaced with a common subgraph. The two phases of deduplication are explained below. **Subgraph creation**: Subgraph creation works by taking one representative region from each region group and creating a subgraph from it, which will then be used to replace all regions in the group. This is implemented by first copying all nodes of the region to the new subgraph and then finding all inputs which are not within the region and creating placeholders for them. For the outputs, all regions in a region group need to be scanned to ensure the largest set of outputs is found, and then an output node is created which returns a tuple of all outputs. **Graph replacement**: To replace each region with the extracted subgraph, the node index in the region and argument index within the node's flattened args and kwargs are recorded once during subgraph creation. This allows us to determine which (external to the region) nodes and in which order these nodes are passed as inputs. For the outputs, getitem nodes are created for each output, and all nodes in the region with external outputs are replaced by the proper getitem node. Finally, all original nodes are erased (there should be no uses of these left in the graph). cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 Edge cases to add more testing for in future PRs (in progress): * multiple nodes on the same line * dynamic shapes checking (need to verify symbolic inputs are the same across subgraphs) * ensure we don't expand regions where it will create a cycle during subgraph replacement * ensure outputs are always tensors (or tuples of tensors iirc) * out of order kwargs, unevenly nested kwargs [ghstack-poisoned]

This PR implements the deduplication pass (blocked by config currently) for dynamo where identical regions from #141381 are replaced with a common subgraph. The two phases of deduplication are explained below. **Subgraph creation**: Subgraph creation works by taking one representative region from each region group and creating a subgraph from it, which will then be used to replace all regions in the group. This is implemented by first copying all nodes of the region to the new subgraph and then finding all inputs which are not within the region and creating placeholders for them. For the outputs, all regions in a region group need to be scanned to ensure the largest set of outputs is found, and then an output node is created which returns a tuple of all outputs. **Graph replacement**: To replace each region with the extracted subgraph, the node index in the region and argument index within the node's flattened args and kwargs are recorded once during subgraph creation. This allows us to determine which (external to the region) nodes and in which order these nodes are passed as inputs. For the outputs, getitem nodes are created for each output, and all nodes in the region with external outputs are replaced by the proper getitem node. Finally, all original nodes are erased (there should be no uses of these left in the graph). cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 Edge cases to add more testing for in future PRs (in progress): * multiple nodes on the same line * dynamic shapes checking (need to verify symbolic inputs are the same across subgraphs) * ensure we don't expand regions where it will create a cycle during subgraph replacement * ensure outputs are always tensors (or tuples of tensors iirc) * out of order kwargs, unevenly nested kwargs [ghstack-poisoned]

This PR implements the deduplication pass (blocked by config currently) for dynamo where identical regions from #141381 are replaced with a common subgraph. The two phases of deduplication are explained below. **Subgraph creation**: Subgraph creation works by taking one representative region from each region group and creating a subgraph from it, which will then be used to replace all regions in the group. This is implemented by first copying all nodes of the region to the new subgraph and then finding all inputs which are not within the region and creating placeholders for them. For the outputs, all regions in a region group need to be scanned to ensure the largest set of outputs is found, and then an output node is created which returns a tuple of all outputs. **Graph replacement**: To replace each region with the extracted subgraph, the node index in the region and argument index within the node's flattened args and kwargs are recorded once during subgraph creation. This allows us to determine which (external to the region) nodes and in which order these nodes are passed as inputs. For the outputs, getitem nodes are created for each output, and all nodes in the region with external outputs are replaced by the proper getitem node. Finally, all original nodes are erased (there should be no uses of these left in the graph). cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 Edge cases to add more testing for in future PRs (in progress): * ~~multiple nodes on the same line~~ (implemented) * ~~dynamic shapes checking (need to verify symbolic inputs are the same across subgraphs)~~ (implemented) * ensure we don't expand regions where it will create a cycle during subgraph replacement * ensure outputs are always tensors (or tuples of tensors iirc) * ~~out of order kwargs, unevenly nested kwargs~~ (implemented) * input aliasing - TBD, we may add support for this in `invoke_subgraph` or reuse the aliasing analysis here to not form regions with these properties * ~~all global state~~ (implemented) Other followups: * consolidate global state checking across all caching infra [ghstack-poisoned]

This PR implements the deduplication pass (blocked by config currently) for dynamo where identical regions from #141381 are replaced with a common subgraph. The two phases of deduplication are explained below. **Subgraph creation**: Subgraph creation works by taking one representative region from each region group and creating a subgraph from it, which will then be used to replace all regions in the group. This is implemented by first copying all nodes of the region to the new subgraph and then finding all inputs which are not within the region and creating placeholders for them. For the outputs, all regions in a region group need to be scanned to ensure the largest set of outputs is found, and then an output node is created which returns a tuple of all outputs. **Graph replacement**: To replace each region with the extracted subgraph, the node index in the region and argument index within the node's flattened args and kwargs are recorded once during subgraph creation. This allows us to determine which (external to the region) nodes and in which order these nodes are passed as inputs. For the outputs, getitem nodes are created for each output, and all nodes in the region with external outputs are replaced by the proper getitem node. Finally, all original nodes are erased (there should be no uses of these left in the graph). cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 Edge cases to add more testing for in future PRs (in progress): * ~~multiple nodes on the same line~~ (implemented) * ~~dynamic shapes checking (need to verify symbolic inputs are the same across subgraphs)~~ (implemented) * ensure we don't expand regions where it will create a cycle during subgraph replacement * ensure outputs are always tensors (or tuples of tensors iirc) * ~~out of order kwargs, unevenly nested kwargs~~ (implemented) * input aliasing - TBD, we may add support for this in `invoke_subgraph` or reuse the aliasing analysis here to not form regions with these properties * ~~all global state~~ (implemented) Other followups: * consolidate global state checking across all caching infra [ghstack-poisoned]

This PR implements the deduplication pass (blocked by config currently) for dynamo where identical regions from #141381 are replaced with a common subgraph. The two phases of deduplication are explained below. **Subgraph creation**: Subgraph creation works by taking one representative region from each region group and creating a subgraph from it, which will then be used to replace all regions in the group. This is implemented by first copying all nodes of the region to the new subgraph and then finding all inputs which are not within the region and creating placeholders for them. For the outputs, all regions in a region group need to be scanned to ensure the largest set of outputs is found, and then an output node is created which returns a tuple of all outputs. **Graph replacement**: To replace each region with the extracted subgraph, the node index in the region and argument index within the node's flattened args and kwargs are recorded once during subgraph creation. This allows us to determine which (external to the region) nodes and in which order these nodes are passed as inputs. For the outputs, getitem nodes are created for each output, and all nodes in the region with external outputs are replaced by the proper getitem node. Finally, all original nodes are erased (there should be no uses of these left in the graph). cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

This PR implements graph region tracking for later extraction into common subgraphs. The algorithm is as follows: `GraphRegionTracker` tracks each node added to the output graph and generates a key based on the source location, instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with the same key are grouped together in a list of identical nodes. Once graph capture is complete, these nodes are organized into region groups. A region group looks like this: [[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group (starting at the topologically latest region group), the inner regions are gradually expanded one node at time from args and kwargs of the node in each region provided that for all regions in the group, the nodes being added are also identical (ie have the same key computed above). The `get_identical_regions` function is the main entry point which will be used by the graph replacement algorithm in #141383 Edge cases to add more testing for in future PRs (in progress): * ~~multiple nodes on the same line~~ (implemented) * ~~dynamic shapes checking (need to verify symbolic inputs are the same across subgraphs)~~ (implemented) * ensure we don't expand regions where it will create a cycle during subgraph replacement * ensure outputs are always tensors (or tuples of tensors iirc) * ~~out of order kwargs, unevenly nested kwargs~~ (implemented) * input aliasing - TBD, we may add support for this in `invoke_subgraph` or reuse the aliasing analysis here to not form regions with these properties * ~~all global state~~ (implemented) Other followups: * consolidate global state checking across all caching infra [ghstack-poisoned]

This PR implements the deduplication pass (blocked by config currently) for dynamo where identical regions from #141381 are replaced with a common subgraph. The two phases of deduplication are explained below. **Subgraph creation**: Subgraph creation works by taking one representative region from each region group and creating a subgraph from it, which will then be used to replace all regions in the group. This is implemented by first copying all nodes of the region to the new subgraph and then finding all inputs which are not within the region and creating placeholders for them. For the outputs, all regions in a region group need to be scanned to ensure the largest set of outputs is found, and then an output node is created which returns a tuple of all outputs. **Graph replacement**: To replace each region with the extracted subgraph, the node index in the region and argument index within the node's flattened args and kwargs are recorded once during subgraph creation. This allows us to determine which (external to the region) nodes and in which order these nodes are passed as inputs. For the outputs, getitem nodes are created for each output, and all nodes in the region with external outputs are replaced by the proper getitem node. Finally, all original nodes are erased (there should be no uses of these left in the graph). cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

This PR adds debug logging for the region expansion algorithm. Pull Request resolved: #141382 Approved by: https://github.com/williamwen42 ghstack dependencies: #141381

This PR implements the deduplication pass (blocked by config currently) for dynamo where identical regions from #141381 are replaced with a common subgraph. The two phases of deduplication are explained below. **Subgraph creation**: Subgraph creation works by taking one representative region from each region group and creating a subgraph from it, which will then be used to replace all regions in the group. This is implemented by first copying all nodes of the region to the new subgraph and then finding all inputs which are not within the region and creating placeholders for them. For the outputs, all regions in a region group need to be scanned to ensure the largest set of outputs is found, and then an output node is created which returns a tuple of all outputs. **Graph replacement**: To replace each region with the extracted subgraph, the node index in the region and argument index within the node's flattened args and kwargs are recorded once during subgraph creation. This allows us to determine which (external to the region) nodes and in which order these nodes are passed as inputs. For the outputs, getitem nodes are created for each output, and all nodes in the region with external outputs are replaced by the proper getitem node. Finally, all original nodes are erased (there should be no uses of these left in the graph). Pull Request resolved: #141383 Approved by: https://github.com/zou3519 ghstack dependencies: #141381, #141382

…141384) Replaced the function in HOP infra with a method on output graph to make it more general and accessible. Pull Request resolved: #141384 Approved by: https://github.com/zou3519 ghstack dependencies: #141381, #141382, #141383

[Dynamo] Implement graph region tracking for deduplication

8238650

Fixes for bfs Initial region tracking tests hashing + test fixes fix for hash more tests more fixes region tracker updates Update tests Update tests2 [ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo labels Nov 22, 2024

This was referenced Nov 22, 2024

[Dynamo] add debug logging for graph region expansion #141382

Closed

[Dynamo] Initial deduplication pass impl #141383

Closed

[Dynamo] Refactor to use install subgraph method in higher order ops #141384

Closed

mlazos requested a review from anijain2305 November 22, 2024 21:31

mlazos added the release notes: dynamo label Nov 22, 2024

mlazos requested a review from zou3519 November 25, 2024 19:04

zou3519 reviewed Nov 26, 2024

View reviewed changes

torch/_dynamo/graph_region_tracker.py Outdated Show resolved Hide resolved

zou3519 reviewed Nov 26, 2024

View reviewed changes

torch/_dynamo/graph_region_tracker.py Show resolved Hide resolved

zou3519 reviewed Nov 26, 2024

View reviewed changes

torch/_dynamo/graph_region_tracker.py Outdated Show resolved Hide resolved

pytorchmergebot closed this in 96c36a6 Dec 11, 2024

pytorchmergebot added the Merged label Dec 11, 2024

github-actions bot deleted the gh/mlazos/102/head branch January 11, 2025 02:10

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Dynamo] Implement graph region tracking for deduplication #141381

[Dynamo] Implement graph region tracking for deduplication #141381

Uh oh!

mlazos commented Nov 22, 2024 •

edited by pytorch-bot bot

Loading

Uh oh!

pytorch-bot bot commented Nov 22, 2024 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[Dynamo] Implement graph region tracking for deduplication #141381

[Dynamo] Implement graph region tracking for deduplication #141381

Uh oh!

Conversation

mlazos commented Nov 22, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/141381

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mlazos commented Nov 22, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Nov 22, 2024 •

edited

Loading