torch fx cannot trace assert for some cases

[YinglinSun]

🐛 Describe the bug

torch fx cannot trace assert for some cases.

import torch
from torch.fx import Tracer

def test(x):
    H, W = x.shape
    assert (H, W) == (2, 3), 'haha'

tracer = Tracer()
tracer.trace_asserts = True
graph = tracer.trace(test)
print(graph)

It failed for this case.

(ai-0401)yinsun@se02ln001:~/tmp/txp$ python test.py
Traceback (most recent call last):
  File "test.py", line 10, in <module>
    graph = tracer.trace(test)
  File "/home/sa/ac-ap-ci/.conda/envs/ai-0401/lib/python3.8/site-packages/torch/fx/_symbolic_trace.py", line 566, in trace
    self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
  File "test.py", line 6, in test
    assert (H, W) == (2, 3), 'haha'
  File "/home/sa/ac-ap-ci/.conda/envs/ai-0401/lib/python3.8/site-packages/torch/fx/proxy.py", line 278, in __bool__
    return self.tracer.to_bool(self)
  File "/home/sa/ac-ap-ci/.conda/envs/ai-0401/lib/python3.8/site-packages/torch/fx/proxy.py", line 154, in to_bool
    raise TraceError('symbolically traced variables cannot be used as inputs to control flow')
torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow

Versions

(ai-0401)yinsun@se02ln001:~/tmp/txp$ python collect_env.py
Collecting environment information...
PyTorch version: 1.11.0a0+gitbc2c6ed
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: 14.0.6 (https://github.com/conda-forge/clangdev-feedstock 28f7809e7f4286b203af212a154f5a8327bd6fd6)
CMake version: version 3.19.1
Libc version: glibc-2.31

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-64-generic-x86_64-with-glibc2.10
Is CUDA available: False
CUDA runtime version: 11.4.120
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] pytorch3d==0.6.2
[pip3] torch==1.11.0+bc2c6ed.cuda114.cudnn841.se02s03.ap
[pip3] torch-scatter==2.0.8
[pip3] torch-tb-profiler==0.4.0
[pip3] torchfile==0.1.0
[pip3] torchvision==0.9.0a0+8fb5838
[conda] magma-cuda111             2.5.2                         1    pytorch
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.19.5           py38h8246c76_3    conda-forge
[conda] pytorch3d                 0.6.2                    pypi_0    pypi
[conda] torch                     1.11.0+bc2c6ed.cuda114.cudnn841.se02s03.ap          pypi_0    pypi
[conda] torch-scatter             2.0.8                    pypi_0    pypi
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchvision               0.9.0a0+8fb5838          pypi_0    pypi

cc @ezyang @SherlockNoMad @soumith

[YinglinSun]

The issue seems BUILD_TUPLE instruction messed up the logic for locating POP_JUMP_IF_TRUE.

[YinglinSun]

@zdevito @jamesr66a