Repository: pytorch
Issue: "grapharg" key is missing in FXGraph placeholder nodes when calling torch APIs with "out" syntax in torch.cond and torch.while_loop
Labels: triaged,oncall: pt2,module: higher order operators

### 🐛 Describe the bug

Usage of "out" syntax in torch APIs like `torch.ge(it, a, out=b)` will cause a crash in TorchDynamo when both of the following conditions hold:

1. Those torch APIs are inside `torch.cond`'s `true_fn`, `false_fn`, or `torch.while_loop`'s `body_fn`.
2. The tensors used in torch APIs are from outer scope (e.g. global, class attributes).

`torch.compile()` is done within `torch.cond` and `torch.while_loop`'s eager mode implementation here:
https://github.com/pytorch/pytorch/blob/4e921593a499036429f5a3ac1fb2f06237bf11ad/torch/_higher_order_ops/cond.py#L141-L143
https://github.com/pytorch/pytorch/blob/4e921593a499036429f5a3ac1fb2f06237bf11ad/torch/_higher_order_ops/while_loop.py#L141-L142

Specifically, the error is caused in this if statement:
https://github.com/pytorch/pytorch/blob/4e921593a499036429f5a3ac1fb2f06237bf11ad/torch/_dynamo/variables/torch.py#L807-L814

the "grapharg" key is missing in the placeholder nodes
https://github.com/pytorch/pytorch/blob/4e921593a499036429f5a3ac1fb2f06237bf11ad/torch/_dynamo/output_graph.py#L1340-L1342

I have a reproducer below which outputs the message provided in "Error Logs".

### Error logs

```
Traceback (most recent call last):
  File "/media/data/test/repro/greedy_decoding/torch_compile_grapharg.py", line 71, in <module>
    f(z, a, b)
  File "/media/data/test/repro/greedy_decoding/torch_compile_grapharg.py", line 69, in f
    _ = while_loop(cond_fn, body_fn, (z,))
  File "/media/data/pytorch/torch/_higher_order_ops/while_loop.py", line 140, in while_loop
    return torch.compile(while_loop_op, backend="eager", fullgraph=True)(
  File "/media/data/pytorch/torch/_dynamo/eval_frame.py", line 434, in _fn
    return fn(*args, **kwargs)
  File "/media/data/pytorch/torch/_dynamo/convert_frame.py", line 1121, in __call__
    return self._torchdynamo_orig_callable(
  File "/media/data/pytorch/torch/_dynamo/convert_frame.py", line 472, in __call__
    return _compile(
  File "/media/data/pytorch/torch/_utils_internal.py", line 85, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/media/data/pytorch/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "/media/data/miniconda3/envs/nemo/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/media/data/pytorch/torch/_dynamo/convert_frame.py", line 846, in _compile
    raise InternalTorchDynamoError(str(e)).with_traceback(
  File "/media/data/pytorch/torch/_dynamo/convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/media/data/pytorch/torch/_dynamo/utils.py", line 233, in time_wrapper
    r = func(*args, **kwargs)
  File "/media/data/pytorch/torch/_dynamo/convert_frame.py", line 636, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/media/data/pytorch/torch/_dynamo/bytecode_transformation.py", line 1270, in transform_code_object
    transformations(instructions, code_options)
  File "/media/data/pytorch/torch/_dynamo/convert_frame.py", line 178, in _fn
    return fn(*args, **kwargs)
  File "/media/data/pytorch/torch/_dynamo/convert_frame.py", line 582, in transform
    tracer.run()
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 2476, in run
    super().run()
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 904, in run
    while self.step():
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 816, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 510, in wrapper
    return inner_fn(self, inst)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 1521, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 754, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/media/data/pytorch/torch/_dynamo/variables/lazy.py", line 132, in realize_and_forward
    return getattr(self.realize(), name)(*args, **kwargs)
  File "/media/data/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 46, in graph_break_as_hard_error
    return fn(*args, **kwargs)
  File "/media/data/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 880, in call_function
    ) = speculate_subgraph(
  File "/media/data/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 420, in speculate_subgraph
    output = f.call_function(tx, args, sub_kwargs)
  File "/media/data/pytorch/torch/_dynamo/variables/functions.py", line 295, in call_function
    return super().call_function(tx, args, kwargs)
  File "/media/data/pytorch/torch/_dynamo/variables/functions.py", line 95, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 760, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 2691, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 2807, in inline_call_
    tracer.run()
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 904, in run
    while self.step():
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 816, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 510, in wrapper
    return inner_fn(self, inst)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 1533, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/media/data/pytorch/torch/_dynamo/symbolic_convert.py", line 754, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/media/data/pytorch/torch/_dynamo/variables/torch.py", line 831, in call_function
    and kwargs["out"] in tx.output.graphargs
  File "/media/data/pytorch/torch/_dynamo/output_graph.py", line 1366, in graphargs
    return [node.meta["grapharg"] for node in self.placeholders]
  File "/media/data/pytorch/torch/_dynamo/output_graph.py", line 1366, in <listcomp>
    return [node.meta["grapharg"] for node in self.placeholders]
torch._dynamo.exc.InternalTorchDynamoError: grapharg

from user code:
   File "/media/data/pytorch/torch/_dynamo/external_utils.py", line 38, in inner
    return fn(*args, **kwargs)
  File "/media/data/test/repro/greedy_decoding/torch_compile_grapharg.py", line 63, in body_fn
    torch.ge(it, a, out=b)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

### Minified repro

```
import torch
from torch._higher_order_ops.while_loop import while_loop

a = torch.zeros((4,), device="cuda")
b = torch.full((4,), fill_value=0, dtype=torch.bool, device="cuda")
z = torch.tensor(0, device="cuda")

def f(z, a, b):
    def cond_fn(it):
        return it.sum() < 5

    def body_fn(it):
        torch.ge(it, a, out=b)
        return (it + 1)

    _ = while_loop(cond_fn, body_fn, (z,))

f(z, a, b)
```

### Versions

```
Collecting environment information...
PyTorch version: 2.5.0a0+gitcc449f2
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.29.3
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jul  5 2023, 18:54:27) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L40S
Nvidia driver version: 550.54.15
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   45 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          25
On-line CPU(s) list:             0-24
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 9454 48-Core Processor
CPU family:                      25
Model:                           17
Thread(s) per core:              1
Core(s) per socket:              25
Socket(s):                       1
Stepping:                        1
BogoMIPS:                        5491.74
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext invpcid_single ssbd ibrs ibpb vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 clzero wbnoinvd arat avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid overflow_recov succor fsrm flush_l1d
Hypervisor vendor:               VMware
Virtualization type:             full
L1d cache:                       800 KiB (25 instances)
L1i cache:                       800 KiB (25 instances)
L2 cache:                        25 MiB (25 instances)
L3 cache:                        32 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-24
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP disabled, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] onnx==1.16.0
[pip3] open-clip-torch==2.24.0
[pip3] pytorch-lightning==2.2.4
[pip3] torch==2.5.0a0+gitcc449f2
[pip3] torchdiffeq==0.2.3
[pip3] torchmetrics==1.4.0
[pip3] torchsde==0.2.6
[pip3] triton==2.3.1
[conda] blas                      1.0                         mkl  
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.8           py310h5eee18b_0  
[conda] mkl_random                1.2.4           py310hdb19cb5_0  
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpy-base                1.24.3          py310hb5e798b_1  
[conda] open-clip-torch           2.24.0                   pypi_0    pypi
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch
[conda] pytorch-lightning         2.2.4                    pypi_0    pypi
[conda] torch                     2.5.0a0+gitcc449f2           dev_0    <develop>
[conda] torchdiffeq               0.2.3                    pypi_0    pypi
[conda] torchmetrics              1.4.0                    pypi_0    pypi
[conda] torchsde                  0.2.6                    pypi_0    pypi
[conda] torchtriton               2.3.0                     py310    pytorch
```

cc @ezyang @anijain2305 @chauhang @penguinwu @zou3519 @ydwu4