pytorch/torch/utils
Blaine Burton Rister a1bfb39a31 [Inductor] Expand Identity ops prior to block pattern matching (#146000)
# Feature

Inductor sometimes uses `Identity` functions to group various terms of an expression. While this is convenient in some scenarios, it can frustrate pattern matching. For example, when we're matching an indexing expression to tell if it can be represented as a block pointer, that analysis should be invariant to `Identity`'s.

This PR adds a few features to achieve this invariance.
 - Create a new expansion mode `expr.expand(identity=True)`, which removes all `Identity` functions from the expression.
 -  Preprocess the expression with this expansion prior to pattern matching.
 - Bonus: create a new test utility function called `dummy_graph()`, which creates a simple `GraphLowering`. This is useful for testing the pattern matcher, as we need to initialize `V.graph` before we can access `V.graph.sizevars`.

# Test plan
This PR adds a few new unit tests:
 - Added a unit test specifically for `expr.expand(identity=True)`.
 - Added a new unit test module for the block pattern matcher. Tested that we can correctly match some example patterns containing Identity ops.

I originally intended to add an end to end test compiling pointwise cat, and mapping the corresponding memory accesses to block pointers. However, it looks like that will take more work, since the [relevant code path](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/codegen/triton.py#L1306) disables block pointer analysis. It might be better to defer that to a future PR.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/146000
Approved by: https://github.com/eellison, https://github.com/jansel
2025-02-08 18:11:53 +00:00
..
_strobelight PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_sympy [Inductor] Expand Identity ops prior to block pattern matching (#146000) 2025-02-08 18:11:53 +00:00
backcompat
benchmark PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
bottleneck PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
data Fixed a typo in dataset.py (#146600) 2025-02-07 05:09:51 +00:00
hipify PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
jit PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
model_dump PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
serialization Add option to serialization config to reduce random reads from get_record_offset when loading with mmap=True (#143880) 2025-01-31 17:09:20 +00:00
tensorboard PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
viz PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
__init__.py [BE] enable UFMT in torch.utils.data (#127705) 2024-06-27 23:16:24 +00:00
_backport_slots.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_config_module.py Add check that envvar configs are boolean (#145454) 2025-02-05 19:40:10 +00:00
_config_typing.pyi
_content_store.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_contextlib.py
_cpp_embed_headers.py PEP585: Missed conversions (#145342) 2025-01-29 05:24:36 +00:00
_cpp_extension_versioner.py Avoid file encoding issues when loading cpp extensions (#138565) 2024-10-28 14:06:34 +00:00
_cxx_pytree.py [AOTI] Cache treespec_loads calculation (#145815) 2025-01-31 06:38:21 +00:00
_device.py Revert "Fix torch.normal ignores default_device (#144070)" 2025-01-14 17:41:58 +00:00
_exposed_in.py Type exposed_in decorator (#141894) 2024-12-03 16:28:17 +00:00
_filelock.py filelock: Make waitcounter variant to use (#139816) 2024-12-12 01:18:34 +00:00
_foreach_utils.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_freeze.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_functools.py dynamo tracing perf: cache on import_source: 52.9 -> 52.58 (#143058) 2024-12-13 18:20:48 +00:00
_get_clean_triton.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_import_utils.py remove allow-untyped-defs from utils/_import_utils.py (#144089) 2025-01-03 01:21:13 +00:00
_mode_utils.py Flip default value for mypy disallow_untyped_defs [9/11] (#127846) 2024-06-08 18:50:06 +00:00
_ordered_set.py [BE]: Update typing of OrderedSet ancestor (#145783) 2025-01-28 04:43:49 +00:00
_python_dispatch.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_pytree.py [TreeSpec] Add custom comparision function (#146442) 2025-02-07 22:39:21 +00:00
_stats.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_thunk.py Refactor thunkify to return proper thunk abstraction (#132407) 2024-08-06 02:35:45 +00:00
_traceback.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_triton.py Adding lowering to persistent-tma device kernel for _scaled_mm (#142045) 2024-12-09 01:48:40 +00:00
_typing_utils.py Revert "Fix type annotation of Linear.bias (#142326)" 2025-01-26 03:41:00 +00:00
_zip.py Flip default value for mypy disallow_untyped_defs [9/11] (#127846) 2024-06-08 18:50:06 +00:00
backend_registration.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
bundled_inputs.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
checkpoint.py [BE] typing for decorators (#144161) 2025-01-04 16:40:09 +00:00
collect_env.py Revert "Collect packages with importlib in collect_env (#144616)" 2025-01-13 03:11:04 +00:00
cpp_backtrace.py
cpp_extension.py Set -DPy_LIMITED_API flag for py_limited_api=True extensions (#145764) 2025-01-28 20:11:05 +00:00
deterministic.py
dlpack.py
file_baton.py
flop_counter.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
hooks.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
mkldnn.py
mobile_optimizer.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
model_zoo.py
module_tracker.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
show_pickle.py Use typing.IO[bytes] instead of io.BytesIO in annotations (#144994) 2025-01-27 18:08:07 +00:00
throughput_benchmark.py Flip default value for mypy disallow_untyped_defs [10/11] (#127847) 2024-06-08 18:50:06 +00:00
weak.py [BE]: Update mypy to 1.11.2 (#133816) 2024-09-16 19:44:11 +00:00