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
_sympy [Inductor] Expand Identity ops prior to block pattern matching (#146000) 2025-02-08 18:11:53 +00:00
backcompat
benchmark
bottleneck
data Fixed a typo in dataset.py (#146600) 2025-02-07 05:09:51 +00:00
hipify
jit
model_dump
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
viz
__init__.py
_backport_slots.py
_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
_contextlib.py
_cpp_embed_headers.py PEP585: Missed conversions (#145342) 2025-01-29 05:24:36 +00:00
_cpp_extension_versioner.py
_cxx_pytree.py [AOTI] Cache treespec_loads calculation (#145815) 2025-01-31 06:38:21 +00:00
_device.py
_exposed_in.py
_filelock.py
_foreach_utils.py
_freeze.py
_functools.py
_get_clean_triton.py
_import_utils.py
_mode_utils.py
_ordered_set.py
_python_dispatch.py
_pytree.py [TreeSpec] Add custom comparision function (#146442) 2025-02-07 22:39:21 +00:00
_stats.py
_thunk.py
_traceback.py
_triton.py
_typing_utils.py
_zip.py
backend_registration.py
bundled_inputs.py
checkpoint.py
collect_env.py
cpp_backtrace.py
cpp_extension.py
deterministic.py
dlpack.py
file_baton.py
flop_counter.py
hooks.py
mkldnn.py
mobile_optimizer.py
model_zoo.py
module_tracker.py
show_pickle.py
throughput_benchmark.py
weak.py