Improve reasoning for size oblivious equations involving max()

[ezyang]

🐛 Describe the bug

Internal xref: https://fb.workplace.com/groups/6829516587176185/posts/7172183699576137/

This program fails to compile:

@torch.compile(fullgraph=True, backend="eager")
def cf(x):
    u0, u1 = x.tolist()
    torch._check_is_size(u0)
    torch._check_is_size(u1)
    torch._check(u0 + u1 == 20)
    if guard_size_oblivious(torch.sym_max(1, u0 + u1) == 20):
        return torch.tensor(True)
    else:
        return torch.tensor(False)

@run_test
def test_symmax():
    assert cf(torch.tensor([10, 10])).item()

Actually, we should be able to make the inference here, because u0 and u1 are size-like in a size oblivious, so we assume they are >= 2, which means that the Max should evaporate, but we are unable to do this.

cc @gchanan @zou3519 @kadeng @msaroufim @bdhirsh @anijain2305 @chauhang @lezcano

Versions

main

[ezyang]

cc @shazqadeer

[ezyang]

The harder version of this is when 20 is replaced with a (typically backed) symbolic variable s0.

[lezcano]

I would argue that, since you are literally checking u0+u1 == 20 in the line before, we should be able to pattern match this and use it. We should pattern match something that is expression == constant at the very least.