Errors on in-place ops on tensor aliases unresolved by proxy substitution

[shino16]

Repro:

import torch, thunder

def f(a, b):
    return a.exp_() * b.tanh_()

def g(a, _):
    b = a.view(5,5)
    return a.exp_() * b.tanh_()

def h(a, _):
    b = a[0,0]
    return a.exp_() * b.tanh_()

for fn in [f, g, h]:
    jf = thunder.jit(fn)
    x = torch.randn(5, 5, device='cuda')
    x_ = x.detach().clone()
    out = jf(x, x[0, 0])
    out_ = fn(x_, x_[0, 0])

    torch.testing.assert_close(out, out_)
    # AssertionError on f, g and h

Found in #2760 (comment) and #2760 (comment) by @beverlylytle.

Trace of f after update_aliases.py:

# Constructed by Update aliases for in-place ops
import thunder
import thunder.core.prims as prims
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a, b):
  # a: "cuda:0 f32[5, 5]"
  # b: "cuda:0 f32[]"
  (t6,) = prims.update_aliases((a,))

  # /opt/pytorch/lightning-thunder/tmp/main.py:4: 	    return a.exp_() * b.tanh_()
  t1 = ltorch.exp_(t6)  # t1: "cuda:0 f32[5, 5]"
    # t0 = ltorch.exp(t6)  # t0: "cuda:0 f32[5, 5]"
      # t0 = prims.exp(t6)  # t0: "cuda:0 f32[5, 5]"
    # t1 = prims.copy_(t0, t6, grad_enabled=True)  # t1: "cuda:0 f32[5, 5]"
  (t7,) = prims.update_aliases((b,))

  # /opt/pytorch/lightning-thunder/tmp/main.py:4: 	    return a.exp_() * b.tanh_()
  t3 = ltorch.tanh_(t7)  # t3: "cuda:0 f32[]"
    # t2 = ltorch.tanh(t7)  # t2: "cuda:0 f32[]"
      # t2 = prims.tanh(t7)  # t2: "cuda:0 f32[]"
    # t3 = prims.copy_(t2, t7, grad_enabled=True)  # t3: "cuda:0 f32[]"
  t5 = ltorch.mul(t1, t3)  # t5: "cuda:0 f32[5, 5]"
    # t4 = prims.broadcast_in_dim(t3, (5, 5), ())  # t4: "cuda:0 f32[5, 5]"
    # t5 = prims.mul(t1, t4)  # t5: "cuda:0 f32[5, 5]"
  return {'output': (t5,), 'flat_args': [t1, t3]}

Trace after fusion:

# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a, b):
  # a: "cuda:0 f32[5, 5]"
  # b: "cuda:0 f32[]"
  (t10,) = update_aliases((a,))
  del a
  [t1] = nvFusion0(t10)
    # t0 = prims.exp(t10)  # t0: "cuda:0 f32[5, 5]"
    # t1 = prims.copy_(t0, t10, grad_enabled=True)  # t1: "cuda:0 f32[5, 5]"
  del t10
  (t11,) = update_aliases((b,))
  del b
  [t3, t5] = nvFusion1(t11, t1)
    # t2 = prims.tanh(t11)  # t2: "cuda:0 f32[]"
    # t3 = prims.copy_(t2, t11, grad_enabled=True)  # t3: "cuda:0 f32[]"
    # t4 = prims.broadcast_in_dim(t3, (5, 5), ())  # t4: "cuda:0 f32[5, 5]"
    # t5 = prims.mul(t1, t4)  # t5: "cuda:0 f32[5, 5]"
  del t11
  return {'output': (t5,), 'flat_args': [t1, t3]}

The problem here is that nvFusion1 does not know that t11 and t1 share memory.

In order to make sure that t3 = prims.copy_(t2, t11, grad_enabled=True) completes before t5 = prims.mul(t1, t4), we could insert prims.update_aliases before prims.mul, which would fix the bugs because prims.update_aliases is unfusible.

Such solutions create more fusion breaks, so we want to minimize the use of prims.update_aliases. Ideally, we hope to make prims.update_aliases a fusible op and let nvFuser handle memory aliases in its combined region.

[shino16]

This is an example where inserting prims.update_aliases seems more appropriate:

import torch, thunder

def f(a, b):
    t2 = a.exp_()
    b.tanh_()
    return t2.sin()

jf = thunder.jit(f)
x = torch.randn(5, 5, device='cuda')
x_ = x.detach().clone()
out = jf(x, x[0, 0])
out_ = f(x_, x_[0, 0])

torch.testing.assert_close(out, out_)
# AssertionError on f

Trace before update_aliases

# Constructed by Remove context manager prims
import thunder
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a, b):
  # a: "cuda:0 f32[5, 5]"
  # b: "cuda:0 f32[]"

  # /opt/pytorch/lightning-thunder/tmp/main.py:4: 	    t2 = a.exp_()
  t2 = ltorch.exp_(a)  # t2: "cuda:0 f32[5, 5]"
    # t0 = ltorch.exp(a)  # t0: "cuda:0 f32[5, 5]"
      # t0 = prims.exp(a)  # t0: "cuda:0 f32[5, 5]"
    # t2 = prims.copy_(t0, a, grad_enabled=True)  # t2: "cuda:0 f32[5, 5]"

  # /opt/pytorch/lightning-thunder/tmp/main.py:5: 	    b.tanh_()
  t4 = ltorch.tanh_(b)  # t4: "cuda:0 f32[]"
    # t3 = ltorch.tanh(b)  # t3: "cuda:0 f32[]"
      # t3 = prims.tanh(b)  # t3: "cuda:0 f32[]"
    # t4 = prims.copy_(t3, b, grad_enabled=True)  # t4: "cuda:0 f32[]"

  # /opt/pytorch/lightning-thunder/tmp/main.py:6: 	    return t2.sin()
  t5 = ltorch.sin(t2)  # t5: "cuda:0 f32[5, 5]"
    # t5 = prims.sin(t2)  # t5: "cuda:0 f32[5, 5]"
  return {'output': (t5,), 'flat_args': [a, b]}

t2 and t4 shares memory, so t4 = ltorch.tanh_(b) must happen before t5 = ltorch.sin(t2). This justifies inserting prims.update_alises((t2, t4)).

Executed trace

# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a, b):
  # a: "cuda:0 f32[5, 5]"
  # b: "cuda:0 f32[]"
  (t10,) = update_aliases((a,))
  del a
  [t2] = nvFusion0(t10)
    # t0 = prims.exp(t10)  # t0: "cuda:0 f32[5, 5]"
    # t2 = prims.copy_(t0, t10, grad_enabled=True)  # t2: "cuda:0 f32[5, 5]"
  del t10
  (t11,) = update_aliases((b,))
  del b
  [t4, t5] = nvFusion1(t11, t2)
    # t3 = prims.tanh(t11)  # t3: "cuda:0 f32[]"
    # t4 = prims.copy_(t3, t11, grad_enabled=True)  # t4: "cuda:0 f32[]"
    # t5 = prims.sin(t2)  # t5: "cuda:0 f32[5, 5]"
  del t11
  return (t5,)

Fusion Definitions

# nvFusion0
def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[5, 5], contiguity=[True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.ops.exp(tv0)
    tv2 = fd.ops.cast(tv1, dtype=DataType.Float)
    fd.add_output(tv2, tv0)
    fd.add_output(tv2)

# nvFusion1
def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[], contiguity=[], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[5, 5], contiguity=[True, True], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.ops.tanh(tv0)
    tv3 = fd.ops.cast(tv2, dtype=DataType.Float)
    tv4 = fd.ops.sin(tv1)
    fd.add_output(tv3, tv0)
    fd.add_output(tv3)
    fd.add_output(tv4)

[shino16]

diff --git a/thunder/core/update_aliases.py b/thunder/core/update_aliases.py
index 3a679bcf..a37cc40d 100644
--- a/thunder/core/update_aliases.py
+++ b/thunder/core/update_aliases.py
@@ -150,7 +150,10 @@ def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[li
             out_tensors = set(map(variableify, filter(lambda p: isinstance(p, TensorProxy), bsym.flat_proxy_outs)))
             if _is_inplace_op(bsym):
                 inplace_inputs.add(in_tensor)
-                out_tensors = set()
+                if len(out_tensors) == 1 and all(len(group) <= 1 or in_tensor not in group for group in view_groups):
+                    out_tensors = set()
             for group in view_groups:
                 if in_tensor in group:
                     group.update(out_tensors)

This is a small patch on #2760 in my mind that works for all the repros here. And I think this is natural and reasonable to do.

To sort things out, we can eliminate out_tensors being aliased to in_tensor if and only if

• in_tensor is going to be swapped with out_tensor by _update_swap_map in subsequent bsyms, and
• in_tensor has no other aliases so far.

These were the hidden assumptions behind update_aliases.py, and the repros here are their counterexamples. The if in the patch above tests if these assumptions are met.

We eliminate input aliases by replacing one with another in replace_args_with_alias_map, and we eliminate aliases created by in-place ops by replacing in_tensor with out_tensor in subsequent bsyms using swap_map. But when view-creating ops create aliases, they are unresolvable, so we handle it using view_groups and prims.update_aliases. In #2760, the same treatment is applied when we come across unresolvable input aliases. So, when aliases created by in-place ops are unresolvable, it makes sense to put them into view_groups too.

We do aim at causing fewer fusion breaks, but in the repros here, avoiding fusion break seems as difficult as other unresolvable aliases. View-creating ops, unresolvable input aliases and unresolvable in-place ops, all seem to be on the same line.

Note that the patch above does not introduce unnecessary fusion breaks on in-place ops that we have considered, which satisfy the assumptions above.

import torch, thunder

def fn(a, b):
    return a.exp_() * b.tanh_()

jf = thunder.jit(fn)
x = torch.randn(5, 5, device='cuda')
y = torch.randn(5, 5, device='cuda')
x_ = x.detach().clone()
y_ = y.detach().clone()
out = jf(x, y)
out_ = fn(x_, y_)

torch.testing.assert_close((x, y, out), (x_, y_, out_))
print(thunder.last_traces(jf)[-1])

# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a, b):
  # a: "cuda:0 f32[5, 5]"
  # b: "cuda:0 f32[5, 5]"
  (t9,) = update_aliases((a,))
  del a
  [t1] = nvFusion0(t9)
    # t0 = prims.exp(t9)  # t0: "cuda:0 f32[5, 5]"
    # t1 = prims.copy_(t0, t9, grad_enabled=True)  # t1: "cuda:0 f32[5, 5]"
  del t9
  (t10,) = update_aliases((b,))
  del b
  [t3, t4] = nvFusion1(t10, t1)
    # t2 = prims.tanh(t10)  # t2: "cuda:0 f32[5, 5]"
    # t3 = prims.copy_(t2, t10, grad_enabled=True)  # t3: "cuda:0 f32[5, 5]"
    # t4 = prims.mul(t1, t3)  # t4: "cuda:0 f32[5, 5]"
  del t10
  return (t4,)

[beverlylytle]

I'm a little confused by the fn defined in the last comment because there are no views of any tensors, and so the view groups created are uncomplicated, and the effect of the patch isn't obvious.

Consider instead something we looked at in another issue

def f(a, _):
    b = a[0,0]
    c = a.exp_()
    b.tanh_()
    return c.sin()

With the above patch we create the following trace

import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a):
  # a: "cuda:0 f32[5, 5]"
  (t73,) = update_aliases((a,))
  del a

  # /home/blytle/scratch/sym_mod.py:52:             b = a[0,0]
  t83 = torch_slice_prim_impl(t73, [0, 0], [1, 1], [1, 1])  # t83: "cuda:0 f32[1, 1]"
  t84 = torch.squeeze(t83, (0, 1))  # t84: "cuda:0 f32[]"
    # t84 = ltorch.squeeze(t83, (0, 1))  # t84: "cuda:0 f32[]"
      # t84 = prims.squeeze(t83, (0, 1))  # t84: "cuda:0 f32[]"
  del t83
  (t74, t75) = update_aliases((t73, t84))
  del t73, t84
  [t32] = nvFusion0(t74)
    # t29 = prims.exp(t74)  # t29: "cuda:0 f32[5, 5]"
    # t32 = prims.copy_(t29, t74, grad_enabled=True)  # t32: "cuda:0 f32[5, 5]"
  del t74
  (t76, t77, t78) = update_aliases((t32, t75, t32))
  del t76, t32, t75
  [t45] = nvFusion1(t77)
    # t40 = prims.tanh(t77)  # t40: "cuda:0 f32[]"
    # t45 = prims.copy_(t40, t77, grad_enabled=True)  # t45: "cuda:0 f32[]"
  del t77
  (t79, t80, t81, t82) = update_aliases((t78, t45, t78, t45))
  del t79, t80, t82, t78, t45
  [t55] = nvFusion2(t81)
    # t55 = prims.sin(t81)  # t55: "cuda:0 f32[5, 5]"
  return (t55,)

Note (t76, t77, t78) = update_aliases((t32, t75, t32)) updates t32 twice, and again it updates twice in the next update_aliases. This makes the trace very confusing to read and reason about, and I believe it could potentially result in some unexpected behavior. This is coming from the fact that this patch adds both a and an alias of a to the view groups. I guess one could filter out the repeated symbols in the creation of the update_aliases bsym, but it feels inefficient and wrong to add something to view groups and then filter it out.

Moreover, this patch does result in excess instances of update_aliases, very weird ones. The function

def f(a, _):
    b = a.view(5,5)
    c = a.exp_()
    b.tanh_()
    return c.sin() * 2 * b

produces the trace

import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a):
  # a: "cuda:0 f32[5, 5]"
  (t97,) = update_aliases((a,))
  del a

  # /home/blytle/scratch/sym_mod.py:52:             b = a.view(5,5)
  t111 = shallow_copy(t97)  # t111: "cuda:0 f32[5, 5]"
  (t98, t99) = update_aliases((t111, t97))
  del t111, t97
  [t33] = nvFusion0(t99)
    # t30 = prims.exp(t99)  # t30: "cuda:0 f32[5, 5]"
    # t33 = prims.copy_(t30, t99, grad_enabled=True)  # t33: "cuda:0 f32[5, 5]"
  del t99
  (t100, t101, t102) = update_aliases((t33, t98, t33))
  del t100, t33, t98
  [t46] = nvFusion1(t101)
    # t41 = prims.tanh(t101)  # t41: "cuda:0 f32[5, 5]"
    # t46 = prims.copy_(t41, t101, grad_enabled=True)  # t46: "cuda:0 f32[5, 5]"
  del t101
  (t103, t104, t105, t106) = update_aliases((t102, t46, t46, t102))
  del t103, t104, t102, t46
  (t107, t108, t109, t110) = update_aliases((t106, t105, t105, t106))
  del t107, t108, t105
  [t72] = nvFusion3(t106, t109)
    # t56 = prims.sin(t106)  # t56: "cuda:0 f32[5, 5]"
    # t60 = prims.mul(t56, 2.0)  # t60: "cuda:0 f32[5, 5]"
    # t72 = prims.mul(t60, t109)  # t72: "cuda:0 f32[5, 5]"
  del t106, t109
  return (t72,)

[shino16]

For fn in my last comment, please ignore it because I just wanted to clarify that no additional update_aliases is introduced in such cases.

And indeed, having one tensor appearing in update_aliases((...)) is pretty awkward (although I haven't come up with a problematic case). To prevent it, we would need some machinery to update the alias relationship from {a, b} to {c, b}, instead of just adding c to {a, b}, or take care of it in _get_update_bsym, which seems simpler.

For the last case with consecutive update_aliases, I think there is a yet another separate issue, so I wrote #2768. Without my patch, we can reproduce the same behavior by replacing c = a.exp_() with c = a.view(5, 5),

def fn(a, _):
    b = a.view(5,5)
    c = a.view(5,5)
    b.tanh_()
    return c.sin() * 2 * b

# ...

# Constructed by Dead Code Elimination (took 0 milliseconds)
import thunder
import thunder.core.prims as prims
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a):
  # a: "cuda:0 f32[5, 5]"
  (t7,) = prims.update_aliases((a,))

  # /opt/pytorch/lightning-thunder/tmp/main.py:10: 	    b = a.view(5,5)
  b = ltorch.view(t7, 5, 5)  # b: "cuda:0 f32[5, 5]"
    # b = ltorch.reshape(t7, (5, 5))  # b: "cuda:0 f32[5, 5]"
      # b = prims.shallow_copy(t7)  # b: "cuda:0 f32[5, 5]"
  (t8, t9) = prims.update_aliases((t7, b))

  # /opt/pytorch/lightning-thunder/tmp/main.py:11: 	    c = a.view(5,5)
  c = ltorch.view(t8, 5, 5)  # c: "cuda:0 f32[5, 5]"
    # c = ltorch.reshape(t8, (5, 5))  # c: "cuda:0 f32[5, 5]"
      # c = prims.shallow_copy(t8)  # c: "cuda:0 f32[5, 5]"
  (t10, t11, t12) = prims.update_aliases((t8, c, t9))

  # /opt/pytorch/lightning-thunder/tmp/main.py:12: 	    b.tanh_()
  t3 = ltorch.tanh_(t12)  # t3: "cuda:0 f32[5, 5]"
    # t2 = ltorch.tanh(t12)  # t2: "cuda:0 f32[5, 5]"
      # t2 = prims.tanh(t12)  # t2: "cuda:0 f32[5, 5]"
    # t3 = prims.copy_(t2, t12, grad_enabled=True)  # t3: "cuda:0 f32[5, 5]"
  (t13, t14, t15) = prims.update_aliases((t10, t11, t3))

  # /opt/pytorch/lightning-thunder/tmp/main.py:13: 	    return c.sin() * 2 * b
  t4 = ltorch.sin(t14)  # t4: "cuda:0 f32[5, 5]"
    # t4 = prims.sin(t14)  # t4: "cuda:0 f32[5, 5]"
  t5 = ltorch.mul(t4, 2)  # t5: "cuda:0 f32[5, 5]"
    # t5 = prims.mul(t4, 2.0)  # t5: "cuda:0 f32[5, 5]"
  (t16, t17, t18) = prims.update_aliases((t13, t14, t15))

  # /opt/pytorch/lightning-thunder/tmp/main.py:13: 	    return c.sin() * 2 * b
  t6 = ltorch.mul(t5, t18)  # t6: "cuda:0 f32[5, 5]"
    # t6 = prims.mul(t5, t18)  # t6: "cuda:0 f32[5, 5]"
  return {'output': (t6,), 'flat_args': [t16]}

# ...

# Constructed by Unwrap the actual return value
from torch import Tensor
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a):
  # a: "cuda:0 f32[5, 5]"
  (t31,) = update_aliases((a,))
  del a

  # /opt/pytorch/lightning-thunder/tmp/main.py:10: 	    b = a.view(5,5)
  t43 = Tensor.view(t31, 5, 5)  # t43: "cuda:0 f32[5, 5]"
    # t43 = ltorch.view(t31, 5, 5)  # t43: "cuda:0 f32[5, 5]"
      # t43 = ltorch.reshape(t31, (5, 5))  # t43: "cuda:0 f32[5, 5]"
        # t43 = prims.shallow_copy(t31)  # t43: "cuda:0 f32[5, 5]"
  (t32, t33) = update_aliases((t31, t43))
  del t31, t43

  # /opt/pytorch/lightning-thunder/tmp/main.py:11: 	    c = a.view(5,5)
  t44 = Tensor.view(t32, 5, 5)  # t44: "cuda:0 f32[5, 5]"
    # t44 = ltorch.view(t32, 5, 5)  # t44: "cuda:0 f32[5, 5]"
      # t44 = ltorch.reshape(t32, (5, 5))  # t44: "cuda:0 f32[5, 5]"
        # t44 = prims.shallow_copy(t32)  # t44: "cuda:0 f32[5, 5]"
  (t34, t35, t36) = update_aliases((t32, t44, t33))
  del t32, t44, t33
  [t3] = nvFusion0(t36)
    # t2 = prims.tanh(t36)  # t2: "cuda:0 f32[5, 5]"
    # t3 = prims.copy_(t2, t36, grad_enabled=True)  # t3: "cuda:0 f32[5, 5]"
  del t36
  (t37, t38, t39) = update_aliases((t34, t35, t3))
  del t34, t35, t3
  (t40, t41, t42) = update_aliases((t37, t38, t39))
  del t41, t37, t39
  [t6] = nvFusion2(t38, t42)
    # t4 = prims.sin(t38)  # t4: "cuda:0 f32[5, 5]"
    # t5 = prims.mul(t4, 2.0)  # t5: "cuda:0 f32[5, 5]"
    # t6 = prims.mul(t5, t42)  # t6: "cuda:0 f32[5, 5]"
  del t38, t42
  return (t6,)

And this is part of #2768. The view group {a, c, b} was already established, so prims.update_aliases was inserted before both multiplications.

[IvanYashchuk]

Thanks for the detailed write-up @shino16, and for tracing this down with @beverlylytle!

Masato's approach is conceptually consistent. As Masato notes, this treats unresolvable in-place op aliases the same way we treat unresolvable view-creating ops and input aliases. They're all in the same category of "aliases we can't eliminate through substitution". The trade-off is that we accept fusion breaks for the problematic cases (which seems unavoidable without nvFuser-level aliasing support), while preserving fusion for the common case where tensors don't share memory. This looks good to me as a fix on top of #2760, even though traces become quite verbose in the number of "update_aliases". However, let's revisit what we're trying to solve here.

The core problem is clear from the traces: as you say, nvFuser doesn't know that t11 and t1 share memory when they're passed as separate inputs to nvFusion1 (and Thunder-jitted function). Looking at the trace:

[t1] = nvFusion0(t10)
  # t0 = prims.exp(t10)
  # t1 = prims.copy_(t0, t10, grad_enabled=True)
...
[t3, t5] = nvFusion1(t11, t1)
  # t2 = prims.tanh(t11)
  # t3 = prims.copy_(t2, t11, grad_enabled=True)
  # t5 = prims.mul(t1, t4)

When nvFuser executes prims.copy_ to t11, it should invalidate t1, but it has no way to know they alias each other. This should be fixed, but how?

On alias detection at trace initialization

Repeating some of the information from "Alias Detection" from the design doc. The question of whether input tensors should be grouped as aliases at trace initialization is worth investigating. PyTorch provides Tensor.data_ptr() and Tensor.storage_offset() for this. However, as the NumPy mem_overlap.c implementation shows, detecting whether two strided arrays actually share memory elements (not just storage) is NP-hard in the general case.

For Thunder's purposes, I think we should consider two approaches:

1. Conservative (may_share_memory style): Check if memory bounds overlap using storage_offset and extent calculations. This is O(ndim) and gives us a safe over-approximation - some non-overlapping tensors might be grouped together, but we won't miss actual aliases. This is what NumPy's may_share_memory does with max_work=0. This is better than conservatively assuming all input tensors may share memory.

2. Exact (shares_memory style): Solve the Diophantine equation to determine actual overlap. This is expensive and only needed if we want to avoid unnecessary fusion breaks for non-overlapping slices of the same storage.

For correctness, option 1 is sufficient and cheap. We could detect at prologue time which input tensors may share memory, then:

• Group them together so the aliasing relationship is known throughout the trace
• Either insert update_aliases barriers before operations that read from potentially-modified aliases, or
• Pass this grouping information to nvFuser so it can handle sequencing internally

The reshape replacement pattern works when tensors have the same numel, but for cases like x[0,0], we need a different mechanism. One option: could we pass aliasing metadata to FusionDefinitionWrapper, similar to how we pass shape/stride info? Then, nvFuser could at least sequence operations correctly within a fusion.

For the immediate fix, detecting may-share-memory at input processing time and inserting update_aliases before the mul would work (accepting the fusion break). The three test cases (f, g, h) cover the important scenarios and should be added as regression tests.

@beverlylytle and @shino16, could you please check separately with @jjsjann123 on what it would take for nvFuser to handle aliased inputs within a fusion? If Thunder can provide all alias groups, would nvFuser be able to use them?

[beverlylytle]

I submit #2769 for your consideration as an alternative WIP to the patch suggested in the comment above.

#2768 is very interesting to me. I'm curious how to resolve that.

I am interested to see how the discussion with the nvFuser team resolves for situations where either the above patch or #2769 does not resolve the aliased-input to a fusion problem.

[mattteochen]

nvFuser segmentation behaviour could increase the difficulty of a possible update_aliases within fusion regions.

For example, the test function:

def f(a, b):
    result = a.exp_() * b.tanh_()
    a.add_(1)
    return result, a

shape = (2, 2)
a = torch.ones(shape, ...)
b = a[0, 0]

Thunder checked out to #2760 generates:

def computation(a, b):
  # a: "cuda:0 f32[5, 5]"
  # b: "cuda:0 f32[]"
  (t20, t21) = update_aliases((b, a))
  del b, a
  [t1] = nvFusion0(t21)
    # t0 = prims.exp(t21)  # t0: "cuda:0 f32[5, 5]"
    # t1 = prims.copy_(t0, t21, grad_enabled=True)  # t1: "cuda:0 f32[5, 5]"
  del t21
  (t22, t23) = update_aliases((t20, t1))
  del t20, t1
  [t3, result] = nvFusion1(t22, t23)
    # t2 = prims.tanh(t22)  # t2: "cuda:0 f32[]"
    # t3 = prims.copy_(t2, t22, grad_enabled=True)  # t3: "cuda:0 f32[]"
    # t4 = prims.broadcast_in_dim(t3, (5, 5), ())  # t4: "cuda:0 f32[5, 5]"
    # result = prims.mul(t23, t4)  # result: "cuda:0 f32[5, 5]"
  del t22
  (t24, t25) = update_aliases((t3, t23))
  del t3, t23
  [t7] = nvFusion2(t25)
    # t6 = prims.add(t25, 1.0)  # t6: "cuda:0 f32[5, 5]"
    # t7 = prims.copy_(t6, t25, grad_enabled=True)  # t7: "cuda:0 f32[5, 5]"
  del t25
  return (result, t7)

Now, if we focus on the second fusion region, which generates the wrong output, the FusionDefinition looks like:

def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[], contiguity=[], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[5, 5], contiguity=[True, True], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.ops.tanh(tv0)
    tv3 = fd.ops.cast(tv2, dtype=DataType.Float)
    tv4 = fd.ops.broadcast(tv3, is_broadcast_dim=[True, True])
    tv5 = fd.ops.expand(tv4, shape=[5, 5])
    tv6 = fd.ops.mul(tv1, tv5)
    fd.add_output(tv3, tv0)
    fd.add_output(tv3)
    fd.add_output(tv6)

Due to its SSA structure, it's clear now that it doesn't know that tv1 is invalidated after tanh.

If we analyze the generated kernels, we see how the in-place copy_ writeback is deferred as the last executed kernel of this fusion definition; moreover, it's handled separately from the main computation kernel:

// Codegen generated code
__global__ void nvfuser_no_op_f0_c1_r0_g0(Tensor<float, 0, 0> T8, Tensor<float, 0, 0> T0, Tensor<float, 0, 0> T9) {
  T9[0]
     = T8[0];
}

// Codegen generated code
__global__ void nvfuser_no_op_f0_c1_r0_g1(Tensor<float, 0, 0> T0, Tensor<float, 0, 0> T8) {
  Array<float, 1, 1> T2;
  T2[0] = 0;
  T2[0]
     = tanhf(T0[0]);
  Array<float, 1, 1> T3;
  T3[0]
     = T2[0];
  T8[0]
     = T3[0];
}

// Codegen generated code
__global__ void nvfuser_pointwise_f0_c1_r0_g2(Tensor<float, 0, 0> T0, Tensor<float, 2, 2> T1, Tensor<float, 2, 2> T6) {
  nvfuser_index_t i0;
  i0 = ((nvfuser_index_t)threadIdx.x) + (128 * ((nvfuser_index_t)blockIdx.x));
  Array<float, 1, 1> T2;
  T2[0] = 0;
  T2[0]
     = tanhf(T0[0]);
  Array<float, 1, 1> T7;
  T7[0]
     = T2[0];
  Array<float, 1, 1> T5;
  T5[0]
     = T7[0];
  if ((i0 < 25)) {
    Array<float, 1, 1> T10;
    T10[0] = 0;
    T10[0]
       = T1[i0];
    Array<float, 1, 1> T11;
    T11[0]
      = T10[0]
      * T5[0];
    T6[i0]
       = T11[0];
  }
}

From nsys is clear how the value of the in-place copy_ is detached from whatever happens in xxxx_g2:

This results in a race condition if update_aliases can be fused, merging fusion1 and fusion2 (potentially also fusion0) because the writeback by xxx_g0 uses a value precomputed separately from whatever happens in xxx_g2.