Repository: pytorch
Issue: Potential rooms for fewer recompilations by introducing higher-level guards
Labels: triaged,oncall: pt2,module: dynamo

### 🐛 Describe the bug

I encountered this while investigating recompilations in #128071. The relevant model code is [here](https://github.com/HazyResearch/based/blob/5cee0bf62be1582580d073af069b96f7fb8dc6b2/based/models/mixers/convolution.py#L131-L146).

## Repro
```python
import torch

@torch.compile(backend="eager")
def f(x, int_dict, n):
    if n in int_dict:
        return x + 1
    return x + 2

x = torch.ones(2)
f(x, {1 : '1'}, 1)
f(x, {1 : '1', 2 : '2'}, 1)
f(x, {2 : '2'}, 2)
```

Running `TORCH_LOGS="recompiles" repro.py` gives 2 recompilations:
```
V1220 10:52:51.577000 12464 torch/_dynamo/guards.py:2817] [0/1] [__recompiles] Recompiling function f in /Users/ryanguo99/Documents/work/scratch/test-dict-contains-guards.py:3
V1220 10:52:51.577000 12464 torch/_dynamo/guards.py:2817] [0/1] [__recompiles]     triggered by the following guard failure(s):
V1220 10:52:51.577000 12464 torch/_dynamo/guards.py:2817] [0/1] [__recompiles]     - 0/0: len(L['int_dict']) == 1
V1220 10:52:51.592000 12464 torch/_dynamo/guards.py:2817] [0/2] [__recompiles] Recompiling function f in /Users/ryanguo99/Documents/work/scratch/test-dict-contains-guards.py:3
V1220 10:52:51.592000 12464 torch/_dynamo/guards.py:2817] [0/2] [__recompiles]     triggered by the following guard failure(s):
V1220 10:52:51.592000 12464 torch/_dynamo/guards.py:2817] [0/2] [__recompiles]     - 0/1: L['n'] == 1
V1220 10:52:51.592000 12464 torch/_dynamo/guards.py:2817] [0/2] [__recompiles]     - 0/0: KeyError on L['int_dict'][1]
```


Relevant guards by running `TORCH_LOGS="guards" python repro.py`: 
- For `f(x, {1 : '1'}, 1)`
```
[__guards] TREE_GUARD_MANAGER:
[__guards] +- RootGuardManager
[__guards] | +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:493 in init_ambient_guards
[__guards] | +- GLOBAL_STATE: ___check_global_state()
[__guards] | +- TORCH_FUNCTION_MODE_STACK: ___check_torch_function_mode_stack()
[__guards] | +- GuardManager: source=L['n'], accessed_by=FrameLocalsGuardAccessor(key='n', framelocals_idx=2)
[__guards] | | +- EQUALS_MATCH: L['n'] == 1                                                   # if n in int_dict:  # scratch/test.py:22 in f
[__guards] | +- GuardManager: source=L['x'], accessed_by=FrameLocalsGuardAccessor(key='x', framelocals_idx=0)
[__guards] | | +- TENSOR_MATCH: check_tensor(L['x'], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, device=None, requires_grad=False, size=[2], stride=[1])  # return x + 1  # scratch/test.py:23 in f
[__guards] | | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False           # return x + 1  # scratch/test.py:23 in f
[__guards] | +- GuardManager: source=L['int_dict'], accessed_by=FrameLocalsGuardAccessor(key='int_dict', framelocals_idx=1)
[__guards] | | +- DICT_LENGTH: len(L['int_dict']) == 1                                       # if n in int_dict:  # scratch/test.py:22 in f
[__guards] | | +- GuardManager: source=L['int_dict'][1], accessed_by=DictGetItemGuardAccessor(1)
[__guards] | | | +- EQUALS_MATCH: L['int_dict'][1] == '1'                                       # if n in int_dict:  # scratch/test.py:22 in f
```

- For `f(x, {2 : '2'}, 2)`
```
[__guards] TREE_GUARD_MANAGER:
[__guards] +- RootGuardManager
[__guards] | +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:493 in init_ambient_guards
[__guards] | +- GLOBAL_STATE: ___check_global_state()
[__guards] | +- TORCH_FUNCTION_MODE_STACK: ___check_torch_function_mode_stack()
[__guards] | +- GuardManager: source=L['n'], accessed_by=FrameLocalsGuardAccessor(key='n', framelocals_idx=2)
[__guards] | | +- TYPE_MATCH: ___check_type_id(L['n'], 4308563856)                          # if n in int_dict:  # scratch/test.py:22 in f
[__guards] | +- GuardManager: source=L['x'], accessed_by=FrameLocalsGuardAccessor(key='x', framelocals_idx=0)
[__guards] | | +- TENSOR_MATCH: check_tensor(L['x'], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, device=None, requires_grad=False, size=[2], stride=[1])  # return x + 2  # scratch/test.py:24 in f
[__guards] | | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False           # return x + 2  # scratch/test.py:24 in f
[__guards] | +- GuardManager: source=L['int_dict'], accessed_by=FrameLocalsGuardAccessor(key='int_dict', framelocals_idx=1)
[__guards] | | +- DICT_LENGTH: len(L['int_dict']) == 1                                       # if n in int_dict:  # scratch/test.py:22 in f
[__guards] | | +- GuardManager: source=L['int_dict'][1], accessed_by=DictGetItemGuardAccessor(1)
[__guards] | | | +- EQUALS_MATCH: L['int_dict'][2] == '2'
[__guards] +- LAMBDA_GUARD: L['n'] == 2  # if n in int_dict:  # scratch/test.py:22 in f (_dynamo/variables/tensor.py:1200 in evaluate_expr)
```

## Thoughts
As shown above, currently Dynamo specializes pretty hard on `int_dict` and `n`, when processing the expression `n in int_dict`; this causes a lot of recompilations both in this contrived example and #128071.

However, in theory we don't care about the specifics of `int_dict` and `n`, rather we just care about whether `int_dict` contains `n`. Thus, we could emit a more general and higher level guard `DICT_CONTAINS` that's parameterized over both the dictionary source and integer source (the current `DICT_CONTAINS` still specializes over the integer source, as we only allow [1 source](https://github.com/pytorch/pytorch/blob/3ee029d4020c40a07c7e20d4f36f08d9697f8d8f/torch/_guards.py#L217) for each guard).

Is this a big problem? Well, for #128071 we would be able to circumvent this problem by fixing the graph breaks, although in other words this rare-looking scenario is exposed by graph breaks in the wild.

Fixing this feels like a non-trivial undertaking, and it's unclear what the ROI is, so I'm creating this issue to track some findings for now.

### Error logs

_No response_

### Versions

main d8ea4ce63, python 3.12

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames