I'm a little iffy about this since dst becomes triton_kernel_wrapper_functional_proxy here but with above check that should be safe

seems a bit confusing to me, since getitem is the actual node that we're copying into the input, not triton_kernel_wrapper_functional_proxy. But I guess it's not wrong 🤔

hmm isn't this flipped? I think the user code dst.copy_(src) will show up in the graph as:

torch.ops.aten.copy_.default(dst, src)

I think this check will miss some cases. In particular, mutated_arg can have more than 2 users, but it can still be ok to reinplace. Example:

def f(x, out):
    # x will now have **three** users: `add`, `triton_kernel`, and `copy_`
    tmp = torch.add(x, 1)
    triton_kernel[grid](inp=x, out=out)
    return out, tmp

Hmm, I think what we probably want is: if we look at users of mutated_arg later in the graph than the triton kernel, copy_() should be the only user (if this is the case, then it's safe to reinplace).

I'm not sure if there's a builtin way in FX to check "number of users after a given node", but there's a util here that does something similar: https://github.com/pytorch/pytorch/blob/main/torch/fx/passes/reinplace.py#L176

cc @Chillee (lmk if this sounds right to you)

Yeah, @Chillee and I talked about this. As you mentioned, this could be improved by doing exactly what you suggested. I could do that as follow up.

I think the concrete check is

"are there any views of this node that are used after the mutation"

But yeah, like Oguz said, mainly didn't do it out of laziness/simplicity.

This check is probably also too conservative. if there happen to be any views of the input in the graph, then this will also be > 2 and we won't reinplace.

Handling the aliasing case is a bit tricky though (but definitely feels solvable). Do we want to try to handle reinplacing in 100% of cases as part of this PR? Followup also feels totally reasonable, since this seems more like it's just the existing state of the reinplacing pass.

In particular: we'd ideally like to remove the copy_() on x for a case like this:

def f(x, out):
    x_view = x.view(-1)
    triton_kernel[grid](inp=x, out=out)
    out2 = x_view.mul(2)
    return out, out2

But we want to make sure not to remove the copy_() for a case like this (where we actually do mutate the alias of x later):

def f(x, out):
    x_view = x.view(-1)
    triton_kernel[grid](inp=x, out=out)
    x_view.mul_(2)
    return out

yup, this could be a good follow up too. For now, i kept the reinplace_scatters as is