Repository: pytorch
Issue: Custom Operator Design for torch.compile: Must Output Tensors Always Be Returned?
Labels: triaged,module: custom-operators,actionable,oncall: pt2,module: pt2-dispatcher,module: reinplacing

### 🐛 Describe the bug

Hello PyTorch Community,

I am currently developing a custom operator and considering its compatibility with torch.compile. I have a specific question regarding the function signature of the operator: Does the output tensor need to be explicitly returned by the function, or can it be passed as an input argument?

If I pass the output as an input parameter when creating a custom operator, with the operator interface having no return value, I encounter incorrect results when running with torch.compile, but correct results in eager mode. Based on my analysis, torch.compile seems to optimize away the custom operators that do not return any value, leading to incorrect outcomes.

### Error logs

Traceback (most recent call last):
  File "/tmp/test.py", line 40, in <module>
    assert torch.allclose(out, x * y)
AssertionError

### Minified repro

```python
import torch
import numpy as np                                                                                                                                                           
                                                                                                                                                                             
def custom_func(x, y, out):                                                                                                            
    torch._check(x.shape == y.shape)
    torch._check(x.device == y.device)
    x_np = x.numpy()
    y_np = y.numpy()
    z_np = np.multiply(x_np, y_np)
    out.copy_(torch.from_numpy(z_np))
    return

torch.library.define("mylib::custom_func", "(Tensor x, Tensor y, Tensor out) -> None")
 # Add the implementation of the custom op                                                                                                                                 
torch.library.impl("mylib::custom_func", "default", custom_func)

# Add an abstract impl that describes what the properties of the output  tensor are, given the properties of the input Tensors.                                                                                                                     
@torch.library.impl_abstract("mylib::custom_func")
def custom_func_abstract(x, y, out):
    torch._check(x.shape == y.shape)
    torch._check(x.device == y.device)
    torch._check(out.shape == y.shape)
    torch._check(out.device == y.device)
    return
                                                                                                      
@torch.compile(backend="inductor", fullgraph=True)
def f(x, y, out):
    return torch.ops.mylib.custom_func.default(x, y, out)

x = torch.randn(3)
y = torch.randn(3)
out = torch.empty_like(x)
z = f(x, y, out)
assert torch.allclose(out, x * y)
```

### Versions

PyTorch version: 2.3.0a0+40ec155e58.nv24.03                                                                                                                                  
Is debug build: False                                                                                                                                                        
CUDA used to build PyTorch: 12.4                                                                                                                                             
ROCM used to build PyTorch: N/A                                                                                                                                              
                                                                                                                                                                             
OS: Ubuntu 22.04.4 LTS (x86_64)                                                                                                                                              
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0                                                                                                                           
Clang version: Could not collect                                                                                                                                             
CMake version: version 3.28.3                                                                                                                                                
Libc version: glibc-2.35                                                                                                                                                     
                                                                                                                                                                             
Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)                                                                                          
Python platform: Linux-5.4.0-150-generic-x86_64-with-glibc2.35                                                                                                               
Is CUDA available: True                                                                                                                                                      
CUDA runtime version: 12.4.99                                                                                                                                                
CUDA_MODULE_LOADING set to: LAZY                                                                                                                                             
GPU models and configuration:                                                                                                                                                
GPU 0: NVIDIA A100 80GB PCIe                                                                                                                                                 
GPU 1: NVIDIA A100 80GB PCIe                                                                                                                                                 
GPU 2: NVIDIA A100 80GB PCIe                                                                                                                                                 
GPU 3: NVIDIA A100 80GB PCIe                                                                                                                                                 
GPU 4: NVIDIA A100 80GB PCIe                                                                                                                                                 
GPU 5: NVIDIA A100 80GB PCIe                                                                                                                                                 
GPU 6: NVIDIA A100 80GB PCIe                                                                                                                                                 
GPU 7: NVIDIA A100 80GB PCIe                                                                                                                                                 
                                                                                                                                                                             
Nvidia driver version: 525.89.02                                                                                                                                             
cuDNN version: Probably one of the following:                                                                                                                                
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.0.0                                                                                                                                  
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.0.0                                                                                                                              
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.0.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] onnx==1.15.0rc2
[pip3] optree==0.10.0
[pip3] pytorch-quantization==2.1.2
[pip3] pytorch-triton==2.2.0+e28a256d7
[pip3] torch==2.3.0a0+40ec155e58.nv24.3
[pip3] torch-tensorrt==2.3.0a0
[pip3] torchdata==0.7.1a0
[pip3] torchtext==0.17.0a0
[pip3] torchvision==0.18.0a0
[conda] Could not collect

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @zou3519