From f0a5681c585158c61e3a177c96f973ef041cdfdd Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Thu, 20 Nov 2025 10:42:04 -0800 Subject: [PATCH 1/3] Insert prims.shape before reshape bsym in update_aliases.py --- thunder/core/update_aliases.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/thunder/core/update_aliases.py b/thunder/core/update_aliases.py index ffbb59e632..9f04b1b678 100644 --- a/thunder/core/update_aliases.py +++ b/thunder/core/update_aliases.py @@ -1,5 +1,6 @@ from functools import reduce, partial +from thunder.core.compile_data import using_symbolic_values import thunder.core.prims as prims from thunder.core.proxies import TensorProxy, variableify, unvariableify from thunder.core.pytree import tree_flatten @@ -76,21 +77,23 @@ def replace_args_with_alias_map( reshaped_arg = arg if arg_to_replace.shape != arg.shape: with tracectx(computation_trace): - reshaped_arg = prims.reshape.meta(arg, arg_to_replace.shape) - arg_to_optional_bsyms[variableify(arg_to_replace)] = prims.reshape.bind( - arg, - arg_to_replace.shape, - output=reshaped_arg, - ) + shape = prims.shape.meta(arg_to_replace) + reshaped_arg = prims.reshape.meta(arg, shape) + reshape_bsym = prims.reshape.bind(arg, shape, output=reshaped_arg) + if using_symbolic_values(): + shape_bsym = prims.shape.bind(arg_to_replace, output=shape) + arg_to_optional_bsyms[variableify(arg_to_replace)] = (shape_bsym, reshape_bsym) + else: + arg_to_optional_bsyms[variableify(arg_to_replace)] = (reshape_bsym,) swap_map_for_aliases[variableify(arg_to_replace)] = reshaped_arg appended_bsyms = {} for bsym in computation_trace.bound_symbols: for arg in filter(lambda p: isinstance(p, TensorProxy), bsym.flat_args): - reshape_bsym = arg_to_optional_bsyms.get(variableify(arg)) - if reshape_bsym is not None: - if reshape_bsym not in appended_bsyms: - bsyms.append(reshape_bsym) - appended_bsyms[reshape_bsym] = arg + reshape_bsyms = arg_to_optional_bsyms.get(variableify(arg)) + if reshape_bsyms is not None: + if reshape_bsyms not in appended_bsyms: + bsyms.extend(reshape_bsyms) + appended_bsyms[reshape_bsyms] = arg if replaced_args_map := { x.name: swap_map_for_aliases[variableify(x)].name for x in filter(lambda p: isinstance(p, TensorProxy), bsym.flat_args) From 1a2c7d972197ea0a59e1016af56a9d8c78d8d549 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Thu, 20 Nov 2025 10:55:41 -0800 Subject: [PATCH 2/3] Add tests Co-authored-by: beverlylytle --- thunder/tests/test_update_aliases.py | 33 ++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/thunder/tests/test_update_aliases.py b/thunder/tests/test_update_aliases.py index 238da88a4f..4d293aae21 100644 --- a/thunder/tests/test_update_aliases.py +++ b/thunder/tests/test_update_aliases.py @@ -328,8 +328,9 @@ def g(a, b): @instantiate( dtypes=NOTHING, + decorators=(pytest.mark.parametrize("cache", ("constant values", "symbolic values")),), ) -def test_aliased_input(executor, device, dtype): +def test_aliased_input(executor, device, dtype, cache): def f(x, y, z): return y.exp_().add(x) + z.exp() @@ -339,7 +340,7 @@ def f(x, y, z): a_ = a.clone().detach() b_ = b.clone().detach() c_ = c.clone().detach() - jfn = executor.make_callable(f) + jfn = executor.make_callable(f, cache=cache) actual = jfn(a, b, c) expected = f(a_, b_, c_) torch.testing.assert_close(actual, expected) @@ -350,8 +351,9 @@ def f(x, y, z): @instantiate( dtypes=NOTHING, + decorators=(pytest.mark.parametrize("cache", ("constant values", "symbolic values")),), ) -def test_write_to_intermediate_result(executor, device, dtype): +def test_write_to_intermediate_result(executor, device, dtype, cache): if executor == nvFuserExecutor: pytest.xfail("nvFuser does not support writing to intermediate results") @@ -361,7 +363,7 @@ def fn(x): return y a = make_tensor((2, 3), dtype=torch.float32, device=device) - jfn = executor.make_callable(fn, skip_inplace_alias_updates=True) + jfn = executor.make_callable(fn, cache=cache) actual = jfn(a) expected = fn(a) torch.testing.assert_close(actual, expected) @@ -517,3 +519,26 @@ def foo(x): expected_grad = torch.autograd.grad(expected, c, g) torch.testing.assert_close(actual_grad_fx, expected_grad) torch.testing.assert_close(actual_grad_jit, expected_grad) + + +@instantiate( + dtypes=(dtypes.float32,), + decorators=(pytest.mark.parametrize("cache", ("constant values", "symbolic values")),), +) +def test_aliasing_for_viewed_input_of_different_shapes(executor, device, dtype, cache): + def f(x, y, z): + return x + 2, y.add_(z) + + a = make_tensor((2, 3), dtype=dtypes.to_torch_dtype(dtype), device=device) + b = a[0, :] + c = a[1, :] + a_ = a.clone().detach() + b_ = a_[0, :] + c_ = a_[1, :] + jfn = executor.make_callable(f, cache=cache) + actual = jfn(a, b, c) + expected = f(a_, b_, c_) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(a, a_) + torch.testing.assert_close(b, b_) + torch.testing.assert_close(c, c_) From b60ad529247d240ba95cdbf60d125eac56ffe0df Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Wed, 26 Nov 2025 07:28:00 -0800 Subject: [PATCH 3/3] empty commit