pytorch/test/jit
David Berard a237831bc2 [JIT] Optimize DCE by storing a MemoryLocations for an entire set<Value*> (#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 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

Pull Request resolved: https://github.com/pytorch/pytorch/pull/153645
Approved by: https://github.com/eellison
2025-05-19 21:04:59 +00:00
..
_imported_class_test [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
fixtures
fixtures_srcs Revert "Use absolute path path.resolve() -> path.absolute() (#129409)" 2025-01-04 14:17:20 +00:00
xnnpack Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
__init__.py
mydecorator.py
myexception.py
myfunction_a.py
myfunction_b.py
test_alias_analysis.py [BE][Ez]: FURB148 - remove useless enumerate calls (#145619) 2025-01-24 23:37:15 +00:00
test_async.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_aten_pow.py
test_attr.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_autodiff_subgraph_slicing.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_autodiff.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_await.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_backend_nnapi.py Revert "Use absolute path path.resolve() -> path.absolute() (#129409)" 2025-01-04 14:17:20 +00:00
test_backends.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_batch_mm.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_builtins.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_class_type.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_complex.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_complexity.py No actual change, just remove variable contain Tensors from global scope (#143225) 2024-12-17 16:14:25 +00:00
test_convert_activation.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_cuda.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_custom_operators.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_data_parallel.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_dataclasses.py [BE]: Update ruff to 0.11.8 (#153249) 2025-05-12 18:30:52 +00:00
test_dce.py [JIT] Optimize DCE by storing a MemoryLocations for an entire set<Value*> (#153645) 2025-05-19 21:04:59 +00:00
test_decorator.py [BE]: Update ruff to 0.11.8 (#153249) 2025-05-12 18:30:52 +00:00
test_device_analysis.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_dtype_analysis.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_enum.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_exception.py [BE][Ez]: ISC001 Auto concatenate implicit one line strings (#146408) 2025-02-04 19:07:04 +00:00
test_freezing.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_functional_blocks.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_fuser_common.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_generator.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_graph_rewrite_passes.py
test_hash.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_hooks_modules.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_hooks.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_ignorable_args.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_ignore_context_manager.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_isinstance.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_jit_utils.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_list_dict.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_logging.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_misc.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_models.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_module_apis.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_module_containers.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_module_interface.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_modules.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_op_decompositions.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_optimize_for_mobile_preserve_debug_info.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_parametrization.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_pdt.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_peephole.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_profiler.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_python_bindings.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_python_builtins.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_python_ir.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_recursive_script.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_remove_mutation.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_save_load_for_op_version.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_save_load.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_script_profile.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_scriptmod_ann.py Remove outdated test skipif conditions for Python3.9 (#146144) 2025-01-31 19:01:04 +00:00
test_slice.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_sparse.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_string_formatting.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_symbolic_shape_analysis.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_tensor_creation_ops.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_tensor_methods.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_torchbind.py Consistently use load_torchbind_test_lib in tests (#148082) 2025-03-03 19:37:28 +00:00
test_tracer.py [BE][Ez]: ISC001 Auto concatenate implicit one line strings (#146408) 2025-02-04 19:07:04 +00:00
test_type_sharing.py Add None return type to init -- tests rest (#132376) 2024-08-01 15:44:51 +00:00
test_types.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_typing.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00
test_union_pep604.py [BE][Ez]: ISC001 Auto concatenate implicit one line strings (#146408) 2025-02-04 19:07:04 +00:00
test_union.py [BE][Ez]: ISC001 Auto concatenate implicit one line strings (#146408) 2025-02-04 19:07:04 +00:00
test_unsupported_ops.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_upgraders.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_warn.py [BE][Easy][13/19] enforce style for empty lines in import segments in test/j*/ (#129764) 2024-08-01 12:13:42 +00:00
test_with.py Fix unused Python variables in test/[e-z]* (#136964) 2024-12-18 23:02:30 +00:00