Avoid decomposing _unsafe_index in Inductor (pytorch#107882)

`_unsafe_index` was previously added to the core ATen decomp table in pytorch#106814, but this has performance ramifications for Inductor. Therefore, this diff removes it from the decomposition table used by Inductor.

Differential Revision: [D48649210](https://our.internmc.facebook.com/intern/diff/D48649210/)

Pull Request resolved: pytorch#107882
Approved by: https://github.com/SherlockNoMad

torch/_decomp/__init__.py

@@ -167,6 +167,25 @@ def get_decompositions(
     return decompositions
 
 
+def remove_decompositions(
+    decompositions: Dict[OpOverload, Callable],
+    aten_ops: Sequence[Union[OpOverload, OpOverloadPacket]],
+) -> None:
+    """
+    Given a dictionary of decompositions obtained from get_decompositions(), removes
+    operators associated with a list of operator overloads and overload packets passed
+    as input. If the decomposition dictionary does not contain a decomposition that is
+    specified to be removed, it is silently ignored.
+    """
+    for op in aten_ops:
+        if isinstance(op, OpOverloadPacket):
+            for overload_name in op.overloads():
+                opo = getattr(op, overload_name)
+                decompositions.pop(opo, None)
+        elif isinstance(op, OpOverload):
+            decompositions.pop(op, None)
+
+
 # populate the table
 import torch._decomp.decompositions
 import torch._refs

torch/_inductor/decomposition.py

@@ -7,7 +7,11 @@
 import torch
 import torch._decomp as decomp
 import torch.ao.quantization.fx._decomposed
-from torch._decomp import core_aten_decompositions, get_decompositions
+from torch._decomp import (
+    core_aten_decompositions,
+    get_decompositions,
+    remove_decompositions,
+)
 from torch._decomp.decompositions import pw_cast_for_opmath
 from torch._decomp.decompositions_for_rng import extra_random_decomps
 
@@ -56,6 +60,14 @@
 )
 decompositions = {**core_aten_decompositions(), **inductor_decompositions}
 
+# Remove unwanted decompositions included via the core ATen decompositions from
+# the Inductor decomp table.
+decomps_to_exclude = [
+    aten._unsafe_index,
+]
+
+remove_decompositions(decompositions, decomps_to_exclude)
+
 
 def register_decomposition(ops):
     for op in [ops] if callable(ops) else ops: