8000 [JIT] Optimize DCE by storing a MemoryLocations for an entire set<Value*> by davidberard98 · Pull Request #153645 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Value*> #153645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 1 commit into from
Closed

[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Value*> #153645

wants to merge 1 commit into from

Conversation

davidberard98
Copy link
Contributor
@davidberard98 davidberard98 commented May 15, 2025 •
edited
Loading

Summary:
TL;DR: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

Details
The goal of this PR is to optimize this function from AliasDb:

bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}

In the DCE use case, we have a ValueSet of live values, into which we insert Value*s; and sometimes need to check whether a node mutates any of the live values using writesToAlias.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

Implementation: To avoid exposing too many details of AliasDb, I introduce a friend class ValueAndMemoryLocationSet, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use ValueAndMemoryLocationSet if we're using AliasDb for analysis, and otherwise use a Set<Value*> if we don't have AliasDb.

Test Plan: Rely on unit tests.

Differential Revision: D74827086

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel

@pytorch-bot PyTorch Bot
Copy link
pytorch-bot bot commented May 15, 2025 •
edited
Loading

🔗 Helpful Links

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

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 26d4aec with merge base 6487ea3 (image):

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

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

@pytorch-bot pytorch-bot bot added the release notes: jit release notes category label May 15, 2025
@facebook-github-bot facebook-github-bot added oncall: jit Add this issue/PR to JIT oncall triage queue fb-exported labels May 15, 2025
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74827086

davidberard98 added a commit to davidberard98/pytorch that referenced this pull request May 15, 2025
@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
d00c136 
…ue*> (pytorch#153645)

Summary:

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
davidberard98 added a commit to davidberard98/pytorch that referenced this pull request May 15, 2025
@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
edaba45 
…ue*> (pytorch#153645)

Summary:

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74827086

davidberard98 added a commit to davidberard98/pytorch that referenced this pull request May 15, 2025
@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
e7d589a 
…ue*> (pytorch#153645)

Summary:
Pull Request resolved: pytorch#153645

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74827086

davidberard98 added a commit to davidberard98/pytorch that referenced this pull request May 15, 2025
@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
77438e3 
…ue*> (pytorch#153645)

Summary:
Pull Request resolved: pytorch#153645

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
@davidberard98 davidberard98 added the ciflow/trunk Trigger trunk jobs on your pull request label May 15, 2025
davidberard98 added a commit to davidberard98/pytorch that referenced this pull request May 15, 2025
@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
5169ce2 
…ue*> (pytorch#153645)

Summary:

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
davidberard98 added a commit to davidberard98/pytorch that referenced this pull request May 15, 2025
@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
da49722 
…ue*> (pytorch#153645)

Summary:

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74827086

davidberard98 added a commit to davidberard98/pytorch that referenced this pull request May 15, 2025
@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
2b6ee35 
…ue*> (pytorch#153645)

Summary:
Pull Request resolved: pytorch#153645

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
eellison
eellison approved these changes May 19, 2025
View reviewed changes
Copy link
Contributor
@eellison eellison left a comment

Choose a reason for hiding this comment

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

Looks good !

torch/csrc/jit/ir/alias_analysis.h Outdated
Comment on lines 333 to 335
// It is tied to an instance of AliasDb but does not own the ALiasDb. It
// is the user's responsibility to ensure that the AliasDb outlives the
// ValuesAndMemoryLocationSet
Copy link
Contributor

Choose a reason for hiding this comment

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

Lets also add a comment that this is only useable for passes that do not do mutation of the AliasDb -

otherwise the mapping of Value->MemLocation will be stale: https://github.com/davidberard98/pytorch/blob/2b6ee35adcac92bf92f5ae96d707502e3e755c37/torch/csrc/jit/ir/alias_analysis.h#L159-L175

@davidberard98 @facebook-github-bot
[JIT] Optimize DCE by storing a MemoryLocations for an entire set<Val…
26d4aec 
…ue*> (pytorch#153645)

Summary:

**TL;DR**: make DCE faster by replacing a Set<Value*> with a MemoryLocations sparse bitset (representing all the memory locations stored by the collection of all values in the set).

**Details**
The goal of this PR is to optimize this function from AliasDb:

```
bool AliasDb::writesToAlias(Node* n, const ValueSet& vs) const {
  const auto writtenTo = getWrites(n);
  if (writtenTo.empty()) {
    return false;
  }

  MemoryLocations locs;
  for (const auto v : vs) {
    auto it = elementMap_.find(v);
    if (it != elementMap_.end()) {
      const auto& vlocs = memoryDAG_->getMemoryLocations(it->second);
      if (writtenTo.intersects(vlocs)) {
        return true;
      }
    }
  }

  return false;
}
```

In the DCE use case, we have a ValueSet of live values, into which we insert `Value*`s; and sometimes need to check whether a node mutates any of the live values using `writesToAlias`.

Looping through all the values in the ValueSet and indexing into the elementMap_ is slow; so if we can pre-compute the MemoryLocations set, this speeds up the function. In some large model examples, I see ~15-25x speedups from this change.

**Implementation**: To avoid exposing too many details of AliasDb, I introduce a friend class `ValueAndMemoryLocationSet`, which is an insert-only set of Values, which also maintains the corresponding MemoryLocations.

Then in AliasDb, I use `ValueAndMemoryLocationSet` if we're dealing with

Test Plan: Rely on unit tests.

Differential Revision: D74827086
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74827086

@facebook-github-bot
Copy link
Contributor

@pytorchbot merge

(Initiating merge automatically since Phabricator Diff has merged)

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

malfet added a commit that referenced this pull request May 19, 2025
@malfet
[BE] Fix -Wextra-semi warning
568e9f2 
Introduced by #153645

Semicolon is not needed after closing curly bracket defining a class method.

Not sure why CI did not catch it, but my local builds are now erroring out with
```
[19/97] Building CXX object caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/passes/dead_code_elimination.cpp.o
In file included from /Users/nshulga/git/pytorch/pytorch/torch/csrc/jit/passes/dead_code_elimination.cpp:4:
/Users/nshulga/git/pytorch/pytorch/torch/csrc/jit/ir/alias_analysis.h:356:64: warning: extra ';' after member function definition [-Wextra-semi]
  356 |   ValueAndMemoryLocationSet(const AliasDb* db) : aliasDb_(db){};
      |                                                                ^
```
@malfet malfet mentioned this pull request May 19, 2025
Closed
pytorchmergebot pushed a commit that referenced this pull request May 19, 2025
@malfet @pytorchmergebot
[BE] Fix -Wextra-semi warning (#153887)
08e716f 
Introduced by #153645

Semicolon is not needed after closing curly bracket defining a class method.

Not sure why CI did not catch it, but my local builds are now erroring out with
```
[19/97] Building CXX object caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/passes/dead_code_elimination.cpp.o
In file included from /Users/nshulga/git/pytorch/pytorch/torch/csrc/jit/passes/dead_code_elimination.cpp:4:
/Users/nshulga/git/pytorch/pytorch/torch/csrc/jit/ir/alias_analysis.h:356:64: warning: extra ';' after member function definition [-Wextra-semi]
  356 |   ValueAndMemoryLocationSet(const AliasDb* db) : aliasDb_(db){};
      |                                                                ^
```

Fixes #ISSUE_NUMBER

Pull Request resolved: #153887
Approved by: https://github.com/wdvr, https://github.com/davidberard98
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@eellison eellison eellison approved these changes

@mikaylagawarecki mikaylagawarecki Awaiting requested review from mikaylagawarecki mikaylagawarecki is a code owner

@gmagogsfm gmagogsfm Awaiting requested review from gmagogsfm

@robieta robieta Awaiting requested review from robieta

Assignees
No one assigned
Labels
ciflow/trunk Trigger trunk jobs on your pull request fb-exported Merged oncall: jit Add this issue/PR to JIT oncall triage queue release notes: jit release notes category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@davidberard98 @facebook-github-bot @pytorchmergebot @eellison
0