mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
a1bfb39a31
# 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 |
||
|---|---|---|
| .. | ||
| aoti_runtime | ||
| cuda | ||
| rocm | ||
| xpu | ||
| __init__.py | ||
| aoti_hipify_utils.py | ||
| block_analysis.py | ||
| common.py | ||
| cpp.py | ||
| cpp_bmm_template.py | ||
| cpp_flex_attention_template.py | ||
| cpp_gemm_template.py | ||
| cpp_grouped_gemm_template.py | ||
| cpp_micro_gemm.py | ||
| cpp_prefix.h | ||
| cpp_template.py | ||
| cpp_template_kernel.py | ||
| cpp_utils.py | ||
| cpp_wrapper_cpu.py | ||
| cpp_wrapper_cpu_array_ref.py | ||
| cpp_wrapper_gpu.py | ||
| cpu_device_op_overrides.py | ||
| cuda_combined_scheduling.py | ||
| debug_utils.py | ||
| halide.py | ||
| memory_planning.py | ||
| mps.py | ||
| mps_device_op_overrides.py | ||
| multi_kernel.py | ||
| simd.py | ||
| simd_kernel_features.py | ||
| triton.py | ||
| triton_combo_kernel.py | ||
| triton_split_scan.py | ||
| triton_utils.py | ||
| wrapper.py | ||