torch.split / torch.chunk cause "too complex strides" for Inductor + CUDA graph

When a model uses output of torch.split / torch.chunk ops as input, or has those ops internally, torch._debug_has_internal_overlap can complain about the size + strides of those tensors being too hard to determine whether memory overlap exists, preventing enablement of CUDA graph:

Example 1: using output of torch.split as input to graph

all_embs = torch.randn(8, 101168)
emb_split = [98400, 340, 40, 328, 380, 1320, 360]
split_emb = torch.split(all_embs, emb_split, dim=1)

split_emb now has tensors of size and stride:

torch.Size([8, 98400])
(101168, 1)
torch.Size([8, 340])
(101168, 1)
torch.Size([8, 40])
(101168, 1)
torch.Size([8, 328])
(101168, 1)
torch.Size([8, 380])
(101168, 1)
torch.Size([8, 1320])
(101168, 1)
torch.Size([8, 360])
(101168, 1)

Passing split_emb as input to “Inductor + CUDA graph” backend will cause CUDA graph to be disabled.

Example 2: using torch.chunk within a graph

_2384 = torch.randn(8, 4096)
chunk_list = torch.chunk(_2384, 16, dim=1)

chunk_list now has 16 tensors of same size and stride:

torch.Size([8, 256])
(4096, 1)

If we do this torch.chunk op within the graph and apply “Inductor + CUDA graph” backend, CUDA graph will be disabled.

cc. @ngimel

[ngimel]

Yes, looks good, you can assume there are no negative strides.

[jansel]

We should stop using torch._debug_has_internal_overlap() and write our own check that is more selective.

Perhaps:

x.numel() != <last element ptr> - <first element ptr> + 1

Then switch our cudagraphs wrapper to just copy the underlying storage and use as_strided().