[test] use JK to force graph break on slow aliasing/mutation/dynamic_shape behavior (#155257)

Summary: test to unblock shampoo, needs cleanup

Test Plan:
CI

Rollback Plan:
steps:
  - jk.update:
      jk: pytorch/compiler:aliased_inputs_with_mutation_and_dyn_shapes_killswitch
      constant_bool: null
      consistent_pass_rate: null
      fractional_host_rollout: null
      sampling_rate: null
  - manual.note:
      content: Set it to false.

Reviewed By: c00w

Differential Revision: D76051868

Pull Request resolved: https://github.com/pytorch/pytorch/pull/155257
Approved by: https://github.com/c00w

torch/_dynamo/exc.py

@@ -404,6 +404,7 @@ exceptions_allowed_to_be_fallback = (
     torch._subclasses.fake_tensor.DynamicOutputShapeException,
     torch._subclasses.fake_tensor.UnsupportedOperatorException,
     torch._subclasses.fake_tensor.UnsupportedFakeTensorException,
+    torch._subclasses.fake_tensor.UnsupportedMutationAliasingException,
 )
 
 

torch/_functorch/_aot_autograd/input_output_analysis.py

@@ -300,6 +300,21 @@ def compute_overlapping_inputs(aot_config, fwd_inputs, aliased_input_indices):
         ]
     )
 
+    if torch._inductor.config.is_fbcode():
+        if symbolic and num_aliases > 400:
+            from torch._subclasses.fake_tensor import (
+                UnsupportedMutationAliasingException,
+            )
+            from torch._utils_internal import justknobs_check
+
+            msg = f"Encountered {num_aliases} dynamic, aliased/mutated inputs, consider setting dynamic=False"
+
+            if justknobs_check(
+                "pytorch/compiler:aliased_inputs_with_mutation_and_dyn_shapes_killswitch",
+                False,
+            ):
+                raise UnsupportedMutationAliasingException(msg)
+
     with maybe_suppress_guards():
         aliased_fwd_inputs = [fwd_inputs[i] for i in aliased_input_indices]
         actual_aliased_indices = {

torch/_subclasses/fake_tensor.py

@@ -121,6 +121,11 @@ class UnsupportedOperatorException(RuntimeError):
     func: OpOverload
 
 
+@dataclass
+class UnsupportedMutationAliasingException(RuntimeError):
+    reason: str
+
+
 @dataclass
 class MetadataMismatchError(RuntimeError):
     reason: str