Repository: pytorch/torchdynamo

Issue:
[HF: AllenaiLongformerBase, GoogleFnet] RuntimeError: CUDA error: an illegal memory access was encountered
Repro:
```
benchmarks/huggingface.py -d cuda --inductor --training --float32 --use-eval-mode -k   AllenaiLongformerBase
benchmarks/huggingface.py -d cuda --inductor --training --float32 --use-eval-mode -k   GoogleFnet
```

Error:
```
RuntimeError: CUDA error: an illegal memory access was encountered
```

The error disappears if we set `config.triton.cudagraph` to False.


Discussion:
minified:
```
import torch
from torch import tensor, device
import torch.fx as fx
from torchdynamo.testing import rand_strided
from math import inf
from torch.fx.experimental.proxy_tensor import make_fx

class Repro(torch.nn.Module):
    def forward(self, arg2_1, arg3_1):
        add_tensor = torch.ops.aten.add.Tensor(arg3_1, arg2_1);  arg3_1 = arg2_1 = None
        return (add_tensor,)
    
args = [((1, 512, 768), (786432, 1536, 2), torch.float32, 'cuda'), ((1, 512, 768), (0, 768, 1), torch.float32, 'cuda')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
mod = make_fx(Repro())(*args)

from torchinductor.compile_fx import compile_fx_inner

compiled = compile_fx_inner(mod, args)
compiled(*args)
```
Corresponding triton code:
```
from ctypes import c_void_p, c_long
import torch
import random
from torch import empty_strided, as_strided, device
from torchinductor.codecache import CppCodeCache, TritonCodeCache

aten = torch.ops.aten

import triton
import triton.language as tl

from torchinductor.triton_ops.autotune import pointwise_heuristics
from torchinductor.triton_ops.autotune import reduction_heuristics
from torchinductor.triton_ops.autotune import grid

@pointwise_heuristics(size_hints=[524288])
@triton.jit
def kernel0(in_ptr0, in_ptr1, out_ptr0, ks0, xnumel, XBLOCK : tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK])
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl.load(in_ptr1 + ks0*x0, xmask)
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr0 + x0 + tl.zeros([XBLOCK], tl.int32), tmp2, xmask)


def call(arg0_1, arg1_1):
    arg0_1_size = arg0_1.size()
    s0 = arg0_1_size[1]
    s1 = arg0_1_size[2]
    arg0_1_stride = arg0_1.stride()
    s2 = arg0_1_stride[2]
    buf0 = empty_strided((1, s0, s1), (s0*s1, s1, 1), device='cuda', dtype=torch.float32)
    kernel0_xnumel = s0*s1
    kernel0[grid(kernel0_xnumel)](arg1_1, arg0_1, buf0, s2, kernel0_xnumel)
    return (buf0, )


if __name__ == "__main__":
    from torchdynamo.testing import rand_strided
    from torchinductor.utils import print_performance
    arg0_1 = rand_strided((1, 512, 768), (786432, 1536, 2), device='cuda', dtype=torch.float32)
    arg1_1 = rand_strided((1, 512, 768), (0, 768, 1), device='cuda', dtype=torch.float32)
    print_performance(lambda: call(arg0_1, arg1_1))
```
@ngimel , any suggestion on how to debug a cuda graph problem like this one?
```
Traceback (most recent call last):
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 133, in cudagraphify
    static_outputs = model(*static_inputs)
  File "/tmp/torchinductor_binbao/if/cifymydxfc4uki3tenttzslkzauhfd3uwbjw4vfjdst6h2w3l5ad.py", line 38, in call
    kernel0[grid(kernel0_xnumel)](arg1_1, arg0_1, buf0, s2, kernel0_xnumel)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 999, in __call__
    return self.kernel(*wargs, **kwargs, grid=self.grid)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1433, in fun
    return kernel(*args, **meta)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 988, in __call__
    return _triton.runtime.launch(wargs, self.fn.do_not_specialize, cache_key, self.fn.arg_names,
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 956, in add_to_cache
    return self.fn._warmup(key, arg_types=arg_types, device=device_idx, attributes=attributes, constants=constants, num_warps=num_warps, num_stages=num_stages,
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1285, in _warmup
    binary = self._compile(**compile)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1320, in _compile
    name, asm, shared_mem = _triton.code_gen.compile_ttir(backend, generator.module, device, num_warps, num_stages, extern_libs)
RuntimeError: unimplemented code: 900

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "5.py", line 24, in <module>
    compiled = compile_fx_inner(mod, args)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/debug_utils.py", line 249, in debug_wrapper
    raise e
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/debug_utils.py", line 237, in debug_wrapper
    compiled_fn = compiler(gm, example_inputs, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 96, in compile_fx_inner
    compiled_fn = cudagraphify(
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 133, in cudagraphify
    static_outputs = model(*static_inputs)
  File "/fsx/users/binbao/pytorch-release/torch/cuda/graphs.py", line 159, in __exit__
    self.cuda_graph.capture_end()
  File "/fsx/users/binbao/pytorch-release/torch/cuda/graphs.py", line 81, in capture_end
    super(CUDAGraph, self).capture_end()
RuntimeError: CUDA error: operation failed due to a previous error during capture
```
Seems related to https://github.com/pytorch/torchdynamo/issues/601#issuecomment-1189443580, it also had strided inputs, but I haven't looked into it, this is very weird. 
This is due to inputs for cuda graphs copied to memory dense locations, and that causes triton recompilation during graph capture. We need to make sure that contiguity of tensors during warmup/compile and actual graph capture is the same.  

Resolving patch:
diff --git a/torchinductor/compile_fx.py b/torchinductor/compile_fx.py
--- a/torchinductor/compile_fx.py
+++ b/torchinductor/compile_fx.py
@@ -18,6 +18,7 @@
 from . import overrides
 from .decomposition import select_decomp_table
 from .graph import GraphLowering
+from .utils import ceildiv
 from .utils import gen_gm_and_inputs
 from .virtualized import V
 
@@ -97,9 +98,22 @@ def cudagraphify(model, inputs, static_input_idxs=()):
     """
     Assumes inputs[static_input_idxs[i]] are always the same memory address
     """
+
+    def static_input(x):
+        # make sure alignment and contiguity of inputs is preserved
+        needed_size = (
+            sum((shape - 1) * stride for shape, stride in zip(x.size(), x.stride())) + 1
+        )
+        needed_size = ceildiv(needed_size, 32) * 32
+        buffer = torch.zeros(needed_size, dtype=x.dtype, device=x.device)
+        cache_line_offset = (
+            (x.data_ptr() - buffer.data_ptr()) % 32
+        ) // x.element_size()
+        return torch.as_strided(buffer, x.size(), x.stride(), cache_line_offset)
+
     assert isinstance(inputs, (list, tuple))
     static_inputs = [
-        torch.zeros_like(x) if idx not in static_input_idxs else inputs[idx]
+        static_input(x) if idx not in static_input_idxs else inputs[idx]
         for idx, x in enumerate(inputs)
     ]
 


Test patch:
diff --git a/test/test_torchinductor.py b/test/test_torchinductor.py
--- a/test/test_torchinductor.py
+++ b/test/test_torchinductor.py
@@ -2139,6 +2139,18 @@ def fn(a):
 
         self.common(fn, (torch.randn([8, 1, 1]),))
 
+    @patch.object(config.triton, "cudagraphs", True)
+    def test_strided_inputs(self):
+        @torchdynamo.optimize("inductor")
+        def fn(x, y):
+            return x + y
+
+        inputs = (
+            rand_strided((8, 16), (32, 2), device=self.device),
+            rand_strided((8, 16), (16, 1), device=self.device),
+        )
+        self.assertTrue(same(fn(*inputs), inputs[0] + inputs[1]))
+
     @patch.object(config.triton, "cudagraphs", True)
     @patch.object(functorch_config, "use_fake_tensor", True)
     def test_input_mutation1(self):