Add sanity check for inplace copy

[kiya00]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes (#265).
Because it's not safe to use nvfuser's inplace update API. A check is added to ensure that the copy_to argument (and output) of prims.copy_ is not used as input for any of its subsequent operators in a nvFusion.

cc @apaz-cli

[jjsjann123]

Thanks for putting the effort into writing the check, it's great to have those for a codegen block!

I'm a bit unclear about our expectation from this check function though.
Feels almost like we are hoping that running this check globally would alway ensure that we have a correct function. But I don't think it's sufficient.

A user script could be using an argument after it has been inplace updated (see example below). That's a legit program, our check will reject it. A false negative check is not ideal but I think it's acceptable.

def fn(a, b):
  c = a + 1
  d = copy_(c, a)
  e = b + a
  return e

The worse scenario is that, our optimization/transformation could re-order things. For the program above, e = b+a doesn't depend on anything and we could move it before the inplace copy. Which alters the semantics of the function, but it would appear as legit program in our check.
False positive check and silent wrong result is bad.

[jjsjann123]

BTW, I'm not blocking this PR. It's an improvement to what we have.

[kiya00]

False positive check and silent wrong result is bad.

yes, I met this problem too

def func1(x, y):
        z = x * y
        o = thunder.core.prims.copy_(z, x)
        return x + y

generate the trace:

@torch.no_grad()
@no_autocast
def computation(x, y):
  # x: "cuda:0 f32[4, 4]"
  # y: "cuda:0 f32[4, 4]"
  [t2] = nvFusion0(x, y)
    # t2 = prims.add(x, y)  # t2: "cuda:0 f32[4, 4]"
    # result = prims.mul(x, y)  # result: "cuda:0 f32[4, 4]"
    # prims.copy_(result, x)
  del x, y
  return t2

which is wrong result.
I think this sanity check is for checking the final execution trace, not the original function. The reorder problem should be solved by the general inplace and aliasing plan(maybe?) cc: @IvanYashchuk

[kiya00]

BTW, I'm not blocking this PR. It's an improvement to what we have.

Not at all, some of the logic is not right, thanks for pointing out @jjsjann123 , I convert it to draft and I'll fix it on Monday :)

[jjsjann123]

False positive check and silent wrong result is bad.

yes, I met this problem too

def func1(x, y):
        z = x * y
        o = thunder.core.prims.copy_(z, x)
        return x + y

generate the trace:

@torch.no_grad()
@no_autocast
def computation(x, y):
  # x: "cuda:0 f32[4, 4]"
  # y: "cuda:0 f32[4, 4]"
  [t2] = nvFusion0(x, y)
    # t2 = prims.add(x, y)  # t2: "cuda:0 f32[4, 4]"
    # result = prims.mul(x, y)  # result: "cuda:0 f32[4, 4]"
    # prims.copy_(result, x)
  del x, y
  return t2

which is wrong result. I think this sanity check is for checking the final execution trace, not the original function. The reorder problem should be solved by the general inplace and aliasing plan(maybe?) cc: @IvanYashchuk

Got'ya. Sounds like our understanding is aligned and that's all that I wanted to confirm. Feel free to ping me when it's ready for another round of review.

[kiya00]

followed the comments, the modifications are:

• check the sharp edge for each nvFusion

• use the utils.consumers and check the index for the relative position

• added a disable_inplace_copy_check in the thunder.jit

@jjsjann123 @IvanYashchuk , could you take a look again

[jjsjann123]

LGTM. Thanks for updating this!

[kiya00]

Hi @t-vi @mruberry , this PR is a simple sanity check on the final trace for nvfuser copy_, could you take a look?

[t-vi]

Thank you @kiya00 @IvanYashchuk @jjsjann123