From 9d4110eb7cfbd63a05121a8b419ca3d6269c1510 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sun, 30 Nov 2025 08:45:21 -0800 Subject: [PATCH 01/46] Register methods to NumberProxy --- thunder/clang/__init__.py | 6 +++--- thunder/core/proxies.py | 19 +++++++++---------- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py index 333086b906..395dbb50b4 100644 --- a/thunder/clang/__init__.py +++ b/thunder/clang/__init__.py @@ -1453,7 +1453,7 @@ def tanh(a): ) -@clangop() +@clangop(method_name="trunc") def trunc(a: TensorLike | Number) -> TensorLike | Number: # Short-circuits on unsigned inputs (which are already trivially truncated) if dtypes.is_exact_dtype(dtypes.to_dtype(a)): @@ -1793,7 +1793,7 @@ def zeta(a, b): ) -@clangop() +@clangop(method_name="bitwise_left_shift") def bitwise_left_shift(a, b): return _elementwise_binary_wrapper( a, @@ -1803,7 +1803,7 @@ def bitwise_left_shift(a, b): ) -@clangop() +@clangop(method_name="bitwise_right_shift") def bitwise_right_shift(a, b): return _elementwise_binary_wrapper( a, diff --git a/thunder/core/proxies.py b/thunder/core/proxies.py index d9f2bcdd5e..d2f8e58e2c 100644 --- a/thunder/core/proxies.py +++ b/thunder/core/proxies.py @@ -30,7 +30,7 @@ TagBase, ) import thunder.core.baseutils as baseutils -from thunder.core.langctxs import resolve_method, get_langctx +from thunder.core.langctxs import LanguageContext, resolve_method, get_langctx import thunder.core.devices as devices import thunder.core.dtypes as dtypes @@ -739,7 +739,7 @@ def _elementwise_unary_helper(a, name, fn, type_promotion_kind=None): vala = pyval(a) trace: None | TraceCtx = get_tracectx() - lang: None | LangCtx = None + lang: None | LanguageContext = None try: lang = get_langctx() except LookupError: @@ -775,7 +775,7 @@ def __neg__(self): return self._elementwise_unary_helper(self, "neg", operator.neg) def __pos__(self): - return self._elementwise_unary_helper(self, "pos", operator.pos) + return self # See https://docs.python.org/3/reference/datamodel.html#object.__round__ def __round__(self): @@ -797,7 +797,7 @@ def _elementwise_binary_helper(a, b, name, fn, type_promotion_kind=None): valb = pyval(b) if isinstance(b, NumberProxy) else b trace: None | TraceCtx = get_tracectx() - lang: None | LangCtx = None + lang: None | LanguageContext = None try: lang = get_langctx() except LookupError: @@ -954,16 +954,16 @@ def __rxor__(self, other): # tracks implementing these def __lshift__(self, other): - raise NotImplementedError + return self._elementwise_binary_helper(self, other, "bitwise_left_shift", operator.lshift) def __rlshift__(self, other): - raise NotImplementedError + return self._elementwise_binary_helper(other, self, "bitwise_left_shift", operator.lshift) def __rshift__(self, other): - raise NotImplementedError + return self._elementwise_binary_helper(self, other, "bitwise_right_shift", operator.rshift) def __rrshift__(self, other): - raise NotImplementedError + return self._elementwise_binary_helper(other, self, "bitwise_right_shift", operator.rshift) # # Casts to Python numbers @@ -1676,8 +1676,7 @@ def __neg__(self): return method(self) def __pos__(self): - method = resolve_method("pos", self) - return method(self) + return self def __round__(self): method = resolve_method("round", self) From 23a677b2a037ffeaa4e1b9d2b79a6b6c0a99b51d Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sun, 30 Nov 2025 08:50:50 -0800 Subject: [PATCH 02/46] Register Python's scalar ops for pythonex --- thunder/executors/pythonex.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/thunder/executors/pythonex.py b/thunder/executors/pythonex.py index 0083853f59..6ccd732123 100644 --- a/thunder/executors/pythonex.py +++ b/thunder/executors/pythonex.py @@ -257,6 +257,10 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None: a.clear() +ceil = ex.register_operator("ceil", like=prims.ceil, module=math) +floor = ex.register_operator("floor", like=prims.floor, module=math) +trunc = ex.register_operator("trunc", like=prims.trunc, module=math) +py_round = ex.register_operator("round", like=prims.round, module=builtins) acos = ex.register_operator("acos", like=prims.acos, module=math) acosh = ex.register_operator("acosh", like=prims.acosh, module=math) asin = ex.register_operator("asin", like=prims.asin, module=math) @@ -273,6 +277,10 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None: "clear_mutable_collection", meta=_clear_mutable_collection_meta, fn=_clear_mutable_collection_prim_impl ) +ex.register_implementation(prims.ceil, ceil, checker=_elementwise_unary_checker) +ex.register_implementation(prims.floor, floor, checker=_elementwise_unary_checker) +ex.register_implementation(prims.trunc, trunc, checker=_elementwise_unary_checker) +ex.register_implementation(prims.round, py_round, checker=_elementwise_unary_checker) ex.register_implementation(prims.acos, acos, checker=_elementwise_unary_checker) ex.register_implementation(prims.acosh, acosh, checker=_elementwise_unary_checker) ex.register_implementation(prims.asin, asin, checker=_elementwise_unary_checker) @@ -287,8 +295,6 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None: ex.register_implementation(prims.signbit, signbit, checker=_elementwise_unary_checker) -# # bitwise_not = _elementwise_unary_factory("invert", operator) -# # ceil = _elementwise_unary_factory("ceil", math) # # cos = _elementwise_unary_factory("cos", math) # # cosh = _elementwise_unary_factory("cosh", math) # # erf = _elementwise_unary_factory("erf", math) @@ -298,7 +304,6 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None: # # exp = _elementwise_unary_factory("exp", math) # # exp2 = None # # expm1 = _elementwise_unary_factory("expm1", math) -# # floor = _elementwise_unary_factory("floor", math) # # isfinite = _elementwise_unary_factory("isfinite", cmath) # # lgamma = _elementwise_unary_factory("lgamma", math) # # log = _elementwise_unary_factory("log", math) @@ -307,8 +312,6 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None: # # log2 = _elementwise_unary_factory("log2", math) # # ndtri = None # # reciprocal = None -# # # NOTE pythonex_round to avoid a name conflict with the builtin round -# # pythonex_round = _elementwise_unary_factory("round", builtins) # # rsqrt = None # # sign = None # # sin = _elementwise_unary_factory("sin", math) @@ -316,7 +319,6 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None: # # sqrt = _elementwise_unary_factory("sqrt", math) # # tan = _elementwise_unary_factory("tan", math) # # tanh = _elementwise_unary_factory("tanh", math) -# # trunc = _elementwise_unary_factory("trunc", math) # # Elementwise binary primitives @@ -332,6 +334,9 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten bitwise_and = ex.register_operator("bitwise_and", like=prims.bitwise_and, fn=operator.and_) bitwise_or = ex.register_operator("bitwise_or", like=prims.bitwise_or, fn=operator.or_) bitwise_xor = ex.register_operator("bitwise_xor", like=prims.bitwise_xor, fn=operator.xor) +bitwise_not = ex.register_operator("bitwise_not", like=prims.bitwise_not, fn=operator.inv) +bitwise_left_shift = ex.register_operator("bitwise_left_shift", like=prims.bitwise_left_shift, fn=operator.lshift) +bitwise_right_shift = ex.register_operator("bitwise_right_shift", like=prims.bitwise_right_shift, fn=operator.rshift) eq = ex.register_operator("eq", like=prims.eq, module=operator) py_floordiv = ex.register_operator("floordiv", like=prims.py_floordiv, module=operator) fmod = ex.register_operator("fmod", like=prims.fmod, module=math) @@ -356,6 +361,9 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten ex.register_implementation(prims.bitwise_and, bitwise_and, checker=_elementwise_binary_checker) ex.register_implementation(prims.bitwise_or, bitwise_or, checker=_elementwise_binary_checker) ex.register_implementation(prims.bitwise_xor, bitwise_xor, checker=_elementwise_binary_checker) +ex.register_implementation(prims.bitwise_not, bitwise_not, checker=_elementwise_unary_checker) +ex.register_implementation(prims.bitwise_left_shift, bitwise_left_shift, checker=_elementwise_binary_checker) +ex.register_implementation(prims.bitwise_right_shift, bitwise_right_shift, checker=_elementwise_binary_checker) ex.register_implementation(prims.eq, eq, checker=_elementwise_binary_checker) ex.register_implementation(prims.py_floordiv, py_floordiv, checker=_elementwise_binary_checker) ex.register_implementation(prims.fmod, fmod, checker=_elementwise_binary_checker) From 68a5bcc8a3928044b0d14d8ee5da9c9301ad26a8 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sun, 30 Nov 2025 08:52:50 -0800 Subject: [PATCH 03/46] Python's floor/ceil/round/trunc returns int --- thunder/clang/__init__.py | 4 ++-- thunder/core/prims.py | 10 +++++++--- thunder/core/utils.py | 5 ++++- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py index 395dbb50b4..90fe53fdf5 100644 --- a/thunder/clang/__init__.py +++ b/thunder/clang/__init__.py @@ -1370,7 +1370,7 @@ def round(a: TensorLike | Number) -> TensorLike | Number: return _elementwise_unary_wrapper( a, prim=prims.round, - type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.NUMBER_TO_INT, ) @@ -1462,7 +1462,7 @@ def trunc(a: TensorLike | Number) -> TensorLike | Number: return _elementwise_unary_wrapper( a, prim=prims.trunc, - type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.NUMBER_TO_INT, ) diff --git a/thunder/core/prims.py b/thunder/core/prims.py index 90e1f3689b..e4c535808f 100644 --- a/thunder/core/prims.py +++ b/thunder/core/prims.py @@ -1945,9 +1945,12 @@ def _numpy_array_to_torch_tensor_meta(a: TensorProxy, /) -> TensorProxy: # usually produce an output with that same datatype (SAME). # Sometimes, however, elementwise operations can produce an output with a different # datatype than the inputs. For example, comparison operations like eq and lt always -# produce boolean results (ALWAYS_BOOL), math.ceil/math.floor produces integer outputs for number inputs while preserves datatype for tensor inputs, and other operations, like abs, map -# complex numbers to floats (COMPLEX_TO_FLOAT). +# produce boolean results (ALWAYS_BOOL), ceil/floor produces integer outputs for +# number inputs while preserves datatype for tensor inputs (INT_FOR_NUMBER), and +# other operations, like abs, map complex numbers to floats (COMPLEX_TO_FLOAT). # The ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND enum describes these three behaviors so that +# operations, like abs, map complex numbers to floats (COMPLEX_TO_FLOAT). +# The ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND enum describes these four behaviors so that # elementwise operations can rely on helper functions to implement this behavior. class ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND(Enum): SAME = auto() @@ -2316,6 +2319,7 @@ def frexp_meta(a: TensorProxy, /) -> (TensorProxy, TensorProxy): "round", number_fn=builtins.round, supported_input_dtypes=fp_math_dtypes, + output_dtype_kind=ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND.INT_FOR_NUMBER, ) rsqrt = _make_elementwise_unary_prim( @@ -2385,12 +2389,12 @@ def _signbit_number(a: Number) -> bool: supported_input_dtypes=fp_math_dtypes, ) -# NOTE This trunc preserves the dtype of its input trunc = _make_elementwise_unary_prim( PrimIDs.TRUNC, "trunc", supported_input_dtypes=fp_math_dtypes, number_fn=math.trunc, + output_dtype_kind=ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND.INT_FOR_NUMBER, ) diff --git a/thunder/core/utils.py b/thunder/core/utils.py index 40b3e0eb13..14a1ccb41d 100644 --- a/thunder/core/utils.py +++ b/thunder/core/utils.py @@ -458,6 +458,9 @@ def elementwise_type_promotion(*args, type_promotion_kind: ELEMENTWISE_TYPE_PROM ALWAYS_BOOL is like PRESERVE, except the result dtype is always bool. + NUMBER_TO_INT is like DEFAULT, except float promotion dtypes *with no tensor inputs* use int + for their result dtypes. This absorbs the difference between e.g. math.ceil and torch.ceil. + Example operators for each type promotion option: DEFAULT : add @@ -504,7 +507,7 @@ def elementwise_type_promotion(*args, type_promotion_kind: ELEMENTWISE_TYPE_PROM and is_float_dtype(promotiontype) and all_number_type ): - return int, int + return promotiontype, int # Falls through to DEFAULT if is_low_precision_dtype(promotiontype): From 88d971f39e8c1fd114d39edf92c35ebb1499eea9 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sun, 30 Nov 2025 08:56:30 -0800 Subject: [PATCH 04/46] Support all-number inputs to where, register to pythonex Needed for floordiv --- thunder/core/prims.py | 18 +++++++++--------- thunder/executors/pythonex.py | 18 ++++++++++++++++++ 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/thunder/core/prims.py b/thunder/core/prims.py index e4c535808f..80a782bdf8 100644 --- a/thunder/core/prims.py +++ b/thunder/core/prims.py @@ -2808,7 +2808,6 @@ def _lerp_meta(start: TensorProxy, end: TensorProxy, weight: Number | TensorProx ) -# TODO Restore Number x Number x Number support def _where_meta(pred: Number | TensorProxy, a: Number | TensorProxy, b: Number | TensorProxy, /) -> TensorProxy: # Checks types # NOTE pred must be a bool tensor or bool (this is checked later) @@ -2816,13 +2815,6 @@ def _where_meta(pred: Number | TensorProxy, a: Number | TensorProxy, b: Number | utils.check_type(a, (TensorProxy, Number, NumberProxy)) utils.check_type(b, (TensorProxy, Number, NumberProxy)) - if ( - isinstance(pred, (Number, NumberProxy)) - and isinstance(a, (Number, NumberProxy)) - and isinstance(b, (Number, NumberProxy)) - ): - raise NotImplementedError - # Checks pred dtype (bool or bool tensor) if isinstance(pred, (Number, NumberProxy)): utils.check( @@ -2847,11 +2839,19 @@ def _where_meta(pred: Number | TensorProxy, a: Number | TensorProxy, b: Number | numbertype, tensordtype = utils.check_same_dtype(a, b) dtype = tensordtype if tensordtype is not None else numbertype + # Returns a NumberProxy for all-Number inputs + if ( + isinstance(pred, (Number, NumberProxy)) + and isinstance(a, (Number, NumberProxy)) + and isinstance(b, (Number, NumberProxy)) + ): + result_value = pyval(a) if pyval(pred) else pyval(b) + return numberproxy(numbertype, result_value, constraint=utils.resolve_constraints(pred, a, b)) + # Checks shapes utils.check_same_shape(pred, a, b) # Determines output shape - # NOTE Assumes at least one of pred, a, and b is a TensorProxy because of prior check for Number x Number x Number shapes = tuple(x.shape for x in (pred, a, b) if isinstance(x, TensorProxy) and not utils.is_cpu_scalar_tensor(x)) if not shapes: shapes = (pred.shape,) diff --git a/thunder/executors/pythonex.py b/thunder/executors/pythonex.py index 6ccd732123..7c0c051571 100644 --- a/thunder/executors/pythonex.py +++ b/thunder/executors/pythonex.py @@ -381,6 +381,24 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten ex.register_implementation(prims.shape, shape, checker=_always_executable) +def _elementwise_ternary_checker( + a: NumberLike | TensorProxy, b: NumberLike | TensorProxy, c: NumberLike | TensorProxy +) -> bool: + return ( + isinstance(a, (Number, NumberProxy)) + and isinstance(b, (Number, NumberProxy)) + and isinstance(c, (Number, NumberProxy)) + ) + + +def _where_prim_impl(pred, a, b): + return a if pred else b + + +where = ex.register_operator("where", like=prims.where, fn=_where_prim_impl) +ex.register_implementation(prims.where, where, checker=_elementwise_ternary_checker) + + def _sink(*args, **kwargs): return From 0d4d1db6136bf4583a263c80fc2ecca7d4a2ad37 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sun, 30 Nov 2025 08:59:35 -0800 Subject: [PATCH 05/46] Add tests, recompile when number type changes --- thunder/core/jit_ext.py | 1 + thunder/tests/test_elementwise.py | 140 ++++++++++++++++++++---------- 2 files changed, 93 insertions(+), 48 deletions(-) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 6e8f52d45c..c51e2a6b34 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -312,6 +312,7 @@ def proxify(self, value: WrappedValue) -> Any: self.add_constraint((clang.check_number_type_and_value, p, uvalue)) elif co is CACHE_OPTIONS.SYMBOLIC_VALUES: if p is not uvalue: + self.add_constraint((clang.check_instance, p, (type(uvalue),))) value.register_proxy(p) elif co not in (CACHE_OPTIONS.SAME_INPUT, CACHE_OPTIONS.NO_CACHING): raise NotImplementedError(f"Unsupported cache option {co}") diff --git a/thunder/tests/test_elementwise.py b/thunder/tests/test_elementwise.py index 2286f0e971..5e0b50cf42 100644 --- a/thunder/tests/test_elementwise.py +++ b/thunder/tests/test_elementwise.py @@ -3,6 +3,7 @@ import math import operator +import pytest import torch from torch.testing import assert_close, make_tensor @@ -13,23 +14,38 @@ from thunder.tests.framework import instantiate, NOTHING -@instantiate(dtypes=NOTHING, devicetypes=(devices.DeviceType.CPU,)) -def test_elementwise_binary_operations_on_numbers(executor, device, dtype): - # op, allowed a-types, allowed b-types, special handling - elementwise_binary_ops = ( - (operator.add, (bool, int, float), (bool, int, float), None), - (operator.sub, (bool, int, float), (bool, int, float), None), - (operator.mul, (bool, int, float), (bool, int, float), None), - (operator.truediv, (bool, int, float), (bool, int, float), "nonzero_only"), - (operator.floordiv, (bool, int, float), (bool, int, float), "nonzero_only"), - (operator.mod, (bool, int, float), (bool, int, float), "nonzero_only"), - (operator.pow, (bool, int, float, complex), (int,), "pow_exponent"), - (operator.and_, (bool, int), (bool, int), None), - (operator.or_, (bool, int), (bool, int), None), - (operator.xor, (bool, int), (bool, int), None), - (operator.lshift, (bool, int), (int,), "shift_count"), - (operator.rshift, (bool, int), (int,), "shift_count"), - ) +# allowed a-types, allowed b-types, special handling +_elementwise_binary_op_to_test_info = { + operator.add: ((bool, int, float), (bool, int, float), None), + operator.sub: ((bool, int, float), (bool, int, float), None), + operator.mul: ((bool, int, float), (bool, int, float), None), + operator.truediv: ((bool, int, float), (bool, int, float), "nonzero_only"), + operator.floordiv: ((bool, int, float), (bool, int, float), "nonzero_only"), + operator.mod: ((bool, int, float), (bool, int, float), "nonzero_only"), + operator.pow: ((bool, int, float, complex), (int,), "pow_exponent"), + operator.and_: ((bool, int), (bool, int), None), + operator.or_: ((bool, int), (bool, int), None), + operator.xor: ((bool, int), (bool, int), None), + operator.lshift: ((bool, int), (int,), "shift_count"), + operator.rshift: ((bool, int), (int,), "shift_count"), +} + + +@instantiate( + dtypes=NOTHING, + devicetypes=(devices.DeviceType.CPU,), + decorators=( + pytest.mark.parametrize("op", _elementwise_binary_op_to_test_info.keys(), ids=lambda op: op.__name__), + pytest.mark.parametrize("cache_option", ("constant values", "symbolic values")), + ), +) +def test_elementwise_binary_operations_on_numbers(executor, device, dtype, op, cache_option): + a_types, b_types, special = _elementwise_binary_op_to_test_info[op] + + if cache_option == "symbolic values": + # TODO: Support bool + a_types = tuple({int, float}.intersection(a_types)) + b_types = tuple({int, float}.intersection(b_types)) bool_inps = [False, True] int_inps = [-1, 0, 2] @@ -61,36 +77,63 @@ def filter_b_values(vals, special): return shift_inps return vals - for op, a_types, b_types, special in elementwise_binary_ops: + def foo(a, b): + return op(a, b) - def foo(a, b): - return op(a, b) + cfoo = executor.make_callable(foo, cache=cache_option) - cfoo = executor.make_callable(foo) + a_vals = gather_inputs(a_types) + b_vals = filter_b_values(gather_inputs(b_types), special) - a_vals = gather_inputs(a_types) - b_vals = filter_b_values(gather_inputs(b_types), special) + for a in a_vals: + for b in b_vals: + actual = cfoo(a, b) + expected = foo(a, b) + assert_close(actual, expected) - for a in a_vals: - for b in b_vals: - actual = cfoo(a, b) - expected = foo(a, b) + if cache_option == "symbolic values": + foo_a_fixed = partial(foo, a=a) + cfoo_a_fixed = executor.make_callable(foo_a_fixed, cache=cache_option) + actual = cfoo_a_fixed(b=b) + expected = foo_a_fixed(b=b) assert_close(actual, expected) + foo_b_fixed = partial(foo, b=b) + cfoo_b_fixed = executor.make_callable(foo_b_fixed, cache=cache_option) + actual = cfoo_b_fixed(a=a) + expected = foo_b_fixed(a=a) + assert_close(actual, expected) -@instantiate(dtypes=NOTHING, devicetypes=(devices.DeviceType.CPU,)) -def test_elementwise_dunder_operations_on_numbers(executor, device, dtype): - # op, allowed types - elementwise_unary_ops = ( - (builtins.abs, (bool, int, float, complex)), - (math.ceil, (bool, int, float)), - (math.floor, (bool, int, float)), - (operator.inv, (bool, int)), - (operator.neg, (bool, int, float, complex)), - (operator.pos, (bool, int, float, complex)), - (builtins.round, (bool, int, float)), - (math.trunc, (bool, int, float)), - ) + if cache_option == "symbolic values": + assert thunder.cache_misses(cfoo) == len(a_types) * len(b_types) + + +_elementwise_unary_op_to_test_info = { + builtins.abs: (bool, int, float, complex), + math.ceil: (bool, int, float), + math.floor: (bool, int, float), + operator.inv: (bool, int), + operator.neg: (bool, int, float, complex), + operator.pos: (bool, int, float, complex), + builtins.round: (bool, int, float), + math.trunc: (bool, int, float), +} + + +@instantiate( + dtypes=NOTHING, + devicetypes=(devices.DeviceType.CPU,), + decorators=( + pytest.mark.parametrize("op", _elementwise_unary_op_to_test_info.keys(), ids=lambda op: op.__name__), + pytest.mark.parametrize("cache_option", ("constant values", "symbolic values")), + ), +) +def test_elementwise_dunder_operations_on_numbers(executor, device, dtype, op, cache_option): + allowed_types = _elementwise_unary_op_to_test_info[op] + + if cache_option == "symbolic values": + # TODO: Support bool + allowed_types = tuple({int, float}.intersection(allowed_types)) bool_inps = [False, True] int_inps = [-1, 0, 2] @@ -112,18 +155,19 @@ def gather_inputs(allowed_types): return inps - for op, allowed_types in elementwise_unary_ops: + def foo(a): + return op(a) - def foo(a): - return op(a) + cfoo = executor.make_callable(foo, cache=cache_option) - cfoo = executor.make_callable(foo) + for a in gather_inputs(allowed_types): + actual = cfoo(a) + expected = foo(a) - for a in gather_inputs(allowed_types): - actual = cfoo(a) - expected = foo(a) + assert_close(actual, expected) - assert_close(actual, expected) + if cache_option == "symbolic values": + assert thunder.cache_misses(cfoo) == len(allowed_types) # TODO: see issue "Test operator and method variants of operations using From 40f4e7e7cab949b9a0a9440d5a5c1dd552d4a222 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sun, 30 Nov 2025 14:45:11 -0800 Subject: [PATCH 06/46] Add test_where_on_numbers --- thunder/tests/test_elementwise.py | 49 +++++++++++++++++++++++++++++-- 1 file changed, 47 insertions(+), 2 deletions(-) diff --git a/thunder/tests/test_elementwise.py b/thunder/tests/test_elementwise.py index 5e0b50cf42..178acf4dcb 100644 --- a/thunder/tests/test_elementwise.py +++ b/thunder/tests/test_elementwise.py @@ -1,5 +1,6 @@ from functools import partial import builtins +import itertools import math import operator @@ -95,13 +96,11 @@ def foo(a, b): foo_a_fixed = partial(foo, a=a) cfoo_a_fixed = executor.make_callable(foo_a_fixed, cache=cache_option) actual = cfoo_a_fixed(b=b) - expected = foo_a_fixed(b=b) assert_close(actual, expected) foo_b_fixed = partial(foo, b=b) cfoo_b_fixed = executor.make_callable(foo_b_fixed, cache=cache_option) actual = cfoo_b_fixed(a=a) - expected = foo_b_fixed(a=a) assert_close(actual, expected) if cache_option == "symbolic values": @@ -170,6 +169,52 @@ def foo(a): assert thunder.cache_misses(cfoo) == len(allowed_types) +@instantiate( + dtypes=NOTHING, + devicetypes=(devices.DeviceType.CPU,), + decorators=(pytest.mark.parametrize("cache_option", ("constant values", "symbolic values")),), +) +def test_where_on_numbers(executor, device, dtype, cache_option): + bool_inps = [False, True] + int_inps = [-1, 2] + float_inps = [-0.7, 1.1] + + def foo(pred, a, b): + return thunder.core.prims.where(pred, a, b) + + def foo_python(pred, a, b): + return a if pred else b + + cfoo = executor.make_callable(foo, cache=cache_option) + + for inps in [bool_inps, int_inps, float_inps]: + for pred, a, b in itertools.product(bool_inps, inps, inps): + print(pred, a, b) + + actual = cfoo(pred, a, b) + expected = foo_python(pred, a, b) + assert_close(actual, expected) + + if cache_option == "symbolic values": + foo_pred_fixed = partial(foo, pred=pred) + cfoo_pred_fixed = executor.make_callable(foo_pred_fixed, cache=cache_option) + actual = cfoo_pred_fixed(a=a, b=b) + assert_close(actual, expected) + + foo_a_fixed = partial(foo, a=a) + cfoo_a_fixed = executor.make_callable(foo_a_fixed, cache=cache_option) + actual = cfoo_a_fixed(pred=pred, b=b) + assert_close(actual, expected) + + foo_pred_b_fixed = partial(foo, pred=pred, b=b) + cfoo_pred_b_fixed = executor.make_callable(foo_pred_b_fixed, cache=cache_option) + actual = cfoo_pred_b_fixed(a=a) + assert_close(actual, expected) + + if cache_option == "symbolic values": + assert thunder.cache_misses(cfoo) == 3 + + # TODO: see issue "Test operator and method variants of operations using # OpInfos" @instantiate(dtypes=(thunder.float32,)) From 13a2d01e9c19c5a3560371498673faea964db48f Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Mon, 1 Dec 2025 06:39:11 -0800 Subject: [PATCH 07/46] Reduce test time --- thunder/tests/test_elementwise.py | 69 ++++++++++++++++--------------- 1 file changed, 36 insertions(+), 33 deletions(-) diff --git a/thunder/tests/test_elementwise.py b/thunder/tests/test_elementwise.py index 178acf4dcb..8568d55277 100644 --- a/thunder/tests/test_elementwise.py +++ b/thunder/tests/test_elementwise.py @@ -3,6 +3,7 @@ import itertools import math import operator +import random import pytest import torch @@ -44,14 +45,12 @@ def test_elementwise_binary_operations_on_numbers(executor, device, dtype, op, c a_types, b_types, special = _elementwise_binary_op_to_test_info[op] if cache_option == "symbolic values": - # TODO: Support bool a_types = tuple({int, float}.intersection(a_types)) b_types = tuple({int, float}.intersection(b_types)) bool_inps = [False, True] int_inps = [-1, 0, 2] float_inps = [-0.7, 0.0, 0.3, 1.1] - complex_inps = [complex(1, 0.3), complex(-4.1, 0.9)] exponent_inps = [0, 1, 2] shift_inps = [0, 1, 2] @@ -60,7 +59,6 @@ def test_elementwise_binary_operations_on_numbers(executor, device, dtype, op, c bool: bool_inps, int: int_inps, float: float_inps, - complex: complex_inps, } def gather_inputs(allowed_types): @@ -92,20 +90,18 @@ def foo(a, b): expected = foo(a, b) assert_close(actual, expected) - if cache_option == "symbolic values": - foo_a_fixed = partial(foo, a=a) - cfoo_a_fixed = executor.make_callable(foo_a_fixed, cache=cache_option) - actual = cfoo_a_fixed(b=b) - assert_close(actual, expected) - - foo_b_fixed = partial(foo, b=b) - cfoo_b_fixed = executor.make_callable(foo_b_fixed, cache=cache_option) - actual = cfoo_b_fixed(a=a) - assert_close(actual, expected) - if cache_option == "symbolic values": assert thunder.cache_misses(cfoo) == len(a_types) * len(b_types) + fixed_a = random.choice(a_vals) + + foo_a_fixed = partial(foo, a=fixed_a) + cfoo_a_fixed = executor.make_callable(foo_a_fixed, cache=cache_option) + for b in b_vals: + actual = cfoo_a_fixed(b=b) + expected = foo(fixed_a, b) + assert_close(actual, expected) + _elementwise_unary_op_to_test_info = { builtins.abs: (bool, int, float, complex), @@ -131,7 +127,6 @@ def test_elementwise_dunder_operations_on_numbers(executor, device, dtype, op, c allowed_types = _elementwise_unary_op_to_test_info[op] if cache_option == "symbolic values": - # TODO: Support bool allowed_types = tuple({int, float}.intersection(allowed_types)) bool_inps = [False, True] @@ -178,6 +173,7 @@ def test_where_on_numbers(executor, device, dtype, cache_option): bool_inps = [False, True] int_inps = [-1, 2] float_inps = [-0.7, 1.1] + complex_inps = [complex(1, 0.3), complex(-4.1, 0.9)] def foo(pred, a, b): return thunder.core.prims.where(pred, a, b) @@ -187,32 +183,39 @@ def foo_python(pred, a, b): cfoo = executor.make_callable(foo, cache=cache_option) - for inps in [bool_inps, int_inps, float_inps]: + for inps in [bool_inps, int_inps, float_inps, complex_inps]: for pred, a, b in itertools.product(bool_inps, inps, inps): - print(pred, a, b) - actual = cfoo(pred, a, b) expected = foo_python(pred, a, b) assert_close(actual, expected) - if cache_option == "symbolic values": - foo_pred_fixed = partial(foo, pred=pred) - cfoo_pred_fixed = executor.make_callable(foo_pred_fixed, cache=cache_option) - actual = cfoo_pred_fixed(a=a, b=b) - assert_close(actual, expected) + if cache_option == "symbolic values": + assert thunder.cache_misses(cfoo) == 3 - foo_a_fixed = partial(foo, a=a) - cfoo_a_fixed = executor.make_callable(foo_a_fixed, cache=cache_option) - actual = cfoo_a_fixed(pred=pred, b=b) - assert_close(actual, expected) + fixed_pred = random.choice(bool_inps) + foo_pred_fixed = partial(foo, pred=fixed_pred) + cfoo_pred_fixed = executor.make_callable(foo_pred_fixed, cache=cache_option) + for a, b in itertools.product(int_inps, int_inps): + actual = cfoo_pred_fixed(a=a, b=b) + expected = foo_python(fixed_pred, a, b) + assert_close(actual, expected) - foo_pred_b_fixed = partial(foo, pred=pred, b=b) - cfoo_pred_b_fixed = executor.make_callable(foo_pred_b_fixed, cache=cache_option) - actual = cfoo_pred_b_fixed(a=a) - assert_close(actual, expected) + fixed_a = random.choice(float_inps) + foo_a_fixed = partial(foo, a=fixed_a) + cfoo_a_fixed = executor.make_callable(foo_a_fixed, cache=cache_option) + for pred, b in itertools.product(bool_inps, float_inps): + actual = cfoo_a_fixed(pred=fixed_pred, b=b) + expected = foo_python(fixed_pred, fixed_a, b) + assert_close(actual, expected) - if cache_option == "symbolic values": - assert thunder.cache_misses(cfoo) == 3 + fixed_pred = random.choice(bool_inps) + fixed_b = random.choice(bool_inps) + foo_pred_b_fixed = partial(foo, pred=fixed_pred, b=fixed_b) + cfoo_pred_b_fixed = executor.make_callable(foo_pred_b_fixed, cache=cache_option) + for a in bool_inps: + actual = cfoo_pred_b_fixed(a=a) + expected = foo_python(fixed_pred, a, fixed_b) + assert_close(actual, expected) # TODO: see issue "Test operator and method variants of operations using From 6c925fa5056c04ec5b4e9b6d9737aca19e6a19c4 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Mon, 1 Dec 2025 06:41:48 -0800 Subject: [PATCH 08/46] Do not cache ComplexProxy, test on where --- thunder/core/interpreter.py | 2 +- thunder/tests/test_elementwise.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/thunder/core/interpreter.py b/thunder/core/interpreter.py index 235c6689ea..5d50604e70 100644 --- a/thunder/core/interpreter.py +++ b/thunder/core/interpreter.py @@ -398,7 +398,7 @@ def __init__( self._with_provenance_tracking = with_provenance_tracking if with_provenance_tracking: assert isinstance(uncacheable_classes, (list, tuple)) - uncacheable_classes = tuple(set(uncacheable_classes) | {NoneType, int, str, float, bool}) + uncacheable_classes = tuple(set(uncacheable_classes) | {NoneType, int, str, float, bool, complex}) self._uncacheable_classes = uncacheable_classes diff --git a/thunder/tests/test_elementwise.py b/thunder/tests/test_elementwise.py index 8568d55277..2c81aed5a4 100644 --- a/thunder/tests/test_elementwise.py +++ b/thunder/tests/test_elementwise.py @@ -51,6 +51,7 @@ def test_elementwise_binary_operations_on_numbers(executor, device, dtype, op, c bool_inps = [False, True] int_inps = [-1, 0, 2] float_inps = [-0.7, 0.0, 0.3, 1.1] + complex_inps = [complex(1, 0.3), complex(-4.1, 0.9)] exponent_inps = [0, 1, 2] shift_inps = [0, 1, 2] @@ -59,6 +60,7 @@ def test_elementwise_binary_operations_on_numbers(executor, device, dtype, op, c bool: bool_inps, int: int_inps, float: float_inps, + complex: complex_inps, } def gather_inputs(allowed_types): @@ -190,7 +192,7 @@ def foo_python(pred, a, b): assert_close(actual, expected) if cache_option == "symbolic values": - assert thunder.cache_misses(cfoo) == 3 + assert thunder.cache_misses(cfoo) == 4 fixed_pred = random.choice(bool_inps) foo_pred_fixed = partial(foo, pred=fixed_pred) From 135a835a6ee59b94ead1906b910c20ad074afaff Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 15:00:14 -0800 Subject: [PATCH 09/46] Reflect INT_FOR_NUMBER behavior in nvfuserex --- thunder/executors/nvfuserex_impl.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index 8c2d0f6324..4d89116153 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, replace -from functools import partial, lru_cache +from functools import partial, lru_cache, wraps from numbers import Number from typing import Any from collections.abc import Callable, Hashable, Sequence @@ -10,6 +10,7 @@ from typing import cast from looseversion import LooseVersion +from optree import tree_flatten import torch from torch import Tensor @@ -971,6 +972,21 @@ def convert_element_type( register_supported(DTensorPrimIDs.CONVERT_ELEMENT_TYPE, convert_element_type, _convert_element_type_check) +# Helper decorator for primitives that should return ints for number-only inputs, e.g. floor/ceil +def _returns_int_for_numbers(fn: Callable) -> Callable: + @wraps(fn) + def wrapper(*args, fd: FusionDefinition, lc_to_nv_map: dict, **kwargs) -> Any: + result = fn(*args, fd=fd, lc_to_nv_map=lc_to_nv_map, **kwargs) + + if all(isinstance(arg, (Number, NumberProxy)) for arg in tree_flatten((args, kwargs))[0]): + nvint = lcdtype_to_nvdtype(int) + result = fd.ops.cast(result, nvint) + + return result + + return wrapper + + def _bitcast_check(src: TensorProxy, dtype: dtypes.dtype) -> bool: return ( nvfuser_version() > LooseVersion("0.29.0") @@ -1472,6 +1488,7 @@ def bitwise_not(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: register_supported(PrimIDs.BITWISE_NOT, bitwise_not, _elementwise_unary_check) +@_returns_int_for_numbers def ceil(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any: nva = getnv(a, fd, lc_to_nv_map) @@ -1563,6 +1580,7 @@ def expm1(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) register_supported(PrimIDs.EXPM1, expm1, _elementwise_unary_check) +@_returns_int_for_numbers def floor(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any: nva = getnv(a, fd, lc_to_nv_map) @@ -1672,6 +1690,7 @@ def reciprocal(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: d # NOTE nv_round to avoid a name conflict with the builtin round +@_returns_int_for_numbers def nv_round(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any: nva = getnv(a, fd, lc_to_nv_map) @@ -1753,6 +1772,7 @@ def tanh(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) - register_supported(PrimIDs.TANH, tanh, _elementwise_unary_check) +@_returns_int_for_numbers def trunc(a: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any: nva = getnv(a, fd, lc_to_nv_map) From 592b3de7ae1588ccebc9b821c266ce999ddb065a Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 29 Nov 2025 07:11:30 -0800 Subject: [PATCH 10/46] Add py_min, py_max --- thunder/clang/__init__.py | 18 ++++++++++++++++++ thunder/core/prims.py | 16 ++++++++++++++++ thunder/executors/pythonex.py | 4 ++++ 3 files changed, 38 insertions(+) diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py index 90fe53fdf5..c294d1bb5c 100644 --- a/thunder/clang/__init__.py +++ b/thunder/clang/__init__.py @@ -1743,6 +1743,24 @@ def minimum(a, b): return _elementwise_binary_wrapper(a, b, prim=prims.minimum) +@clangop() +def py_max(a: NumberLike | Number, b: NumberLike | Number): + utils.check( + isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), + lambda: "py_max expects number inputs", + ) + return prims.py_max(a, b) + + +@clangop() +def py_min(a: NumberLike | Number, b: NumberLike | Number): + utils.check( + isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), + lambda: "py_min expects number inputs", + ) + return prims.py_min(a, b) + + @clangop(method_name="mul") def mul(a, b): return _elementwise_binary_wrapper(a, b, prim=prims.mul) diff --git a/thunder/core/prims.py b/thunder/core/prims.py index 80a782bdf8..eba42e7ced 100644 --- a/thunder/core/prims.py +++ b/thunder/core/prims.py @@ -232,6 +232,8 @@ class PrimIDs(Enum): LT = auto() MAXIMUM = auto() MINIMUM = auto() + PY_MAX = auto() + PY_MIN = auto() MUL = auto() NE = auto() NEXTAFTER = auto() @@ -2695,6 +2697,20 @@ def _div_numbers(a: Number, b: Number) -> Number: minimum = _make_elementwise_binary_prim(PrimIDs.MINIMUM, "minimum", supported_input_dtypes=comparison_dtypes) +py_max = _make_elementwise_binary_prim( + PrimIDs.PY_MAX, + "py_max", + number_fn=builtins.max, + numbers_only=True, +) + +py_min = _make_elementwise_binary_prim( + PrimIDs.PY_MIN, + "py_min", + number_fn=builtins.min, + numbers_only=True, +) + mul = _make_elementwise_binary_prim( PrimIDs.MUL, "mul", diff --git a/thunder/executors/pythonex.py b/thunder/executors/pythonex.py index 7c0c051571..214d44e5a6 100644 --- a/thunder/executors/pythonex.py +++ b/thunder/executors/pythonex.py @@ -345,6 +345,8 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten gt = ex.register_operator("gt", like=prims.gt, module=operator) le = ex.register_operator("le", like=prims.le, module=operator) lt = ex.register_operator("lt", like=prims.lt, module=operator) +py_max = ex.register_operator("max", like=prims.py_max, module=builtins) +py_min = ex.register_operator("min", like=prims.py_min, module=builtins) mul = ex.register_operator("mul", like=prims.mul, module=operator) ne = ex.register_operator("ne", like=prims.ne, module=operator) # NOTE pythonex_pow to avoid a name conflict with the builtin pow @@ -372,6 +374,8 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten ex.register_implementation(prims.gt, gt, checker=_elementwise_binary_checker) ex.register_implementation(prims.le, le, checker=_elementwise_binary_checker) ex.register_implementation(prims.lt, lt, checker=_elementwise_binary_checker) +ex.register_implementation(prims.py_max, py_max, checker=_elementwise_binary_checker) +ex.register_implementation(prims.py_min, py_min, checker=_elementwise_binary_checker) ex.register_implementation(prims.mul, mul, checker=_elementwise_binary_checker) ex.register_implementation(prims.ne, ne, checker=_elementwise_binary_checker) # NOTE pythonex_pow to avoid a name conflict with the builtin pow From c55ac4aa7d531f399eea655d8a9398b4fc5e19d3 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 29 Nov 2025 07:11:41 -0800 Subject: [PATCH 11/46] Add torch.sym_min/max --- thunder/torch/__init__.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py index 10bca59273..63fdc8c8e9 100644 --- a/thunder/torch/__init__.py +++ b/thunder/torch/__init__.py @@ -2596,6 +2596,24 @@ def minimum(a: TensorProxy, b: TensorProxy) -> TensorProxy: return clang.minimum(a, b) +@torchsymbol(torch.sym_max, id="torch.sym_max") +def sym_max(a: NumberLike, b: NumberLike) -> NumberLike: + utils.check( + isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), + lambda: "torch.sym_max currently supports only number inputs", + ) + return clang.py_max(a, b) + + +@torchsymbol(torch.sym_min, id="torch.sym_min") +def sym_min(a: NumberLike, b: NumberLike) -> NumberLike: + utils.check( + isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), + lambda: "torch.sym_min currently supports only number inputs", + ) + return clang.py_min(a, b) + + # NOTE This is just an alias for proxies to find operation defined for the modulus # operator # TODO Review this alias From 522b3f5a264a587f7ab1fad0ef3228591f67a5ad Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 29 Nov 2025 07:11:48 -0800 Subject: [PATCH 12/46] Add test --- thunder/tests/test_jit_general.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index 3b589cc8fb..e7424b7e31 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -864,6 +864,23 @@ def foo(a, scalar): assert thunder.cache_hits(jfoo) == 1 +def test_symbolic_values_min_max(): + def foo(seq_len, cumulative_len, max_len, sliding_window) -> int: + max_cache_len = torch.sym_min(max_len, seq_len + 128) + offset_raw = cumulative_len - sliding_window + 1 + kv_offset = torch.sym_max(offset_raw, 0) + return max_cache_len, kv_offset + + jfoo = thunder_jit(foo, cache="symbolic values") + + samples = torch.randint(1, 100, (20, 4)) + for sample in samples: + args = [s.item() for s in sample] + assert jfoo(*args) == foo(*args) + + assert thunder.cache_misses(jfoo) == 1 + + def test_post_optimization_transform(): def foo(a, b, c): return a * a + b * c From 3b73a42e745f202881953833759eb29c6c03e5fb Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 29 Nov 2025 07:33:20 -0800 Subject: [PATCH 13/46] Add builtin min/max lookaside --- thunder/core/jit_ext.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index c51e2a6b34..e4a7ab2593 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -653,6 +653,32 @@ def _general_jit_bool_lookaside(wrapped_x: Any) -> bool | INTERPRETER_SIGNALS: _general_jit_lookaside_map[bool] = _general_jit_bool_lookaside +def _general_jit_min_max_lookaside(op_name, symbol, args, kwargs): + if len(args) != 2: + raise TypeError(f"{op_name}() currently supports exactly two positional arguments in thunder.jit") + + if kwargs: + unexpected = ", ".join(map(str, kwargs.keys())) + raise TypeError(f"{op_name}() keyword arguments are not supported in thunder.jit (got: {unexpected})") + + a, b = (unwrap(arg) for arg in args) + reduced = symbol(a, b) + + provenance_inputs = [arg.provenance for arg in args] + + return wrap(reduced, provenance=ProvenanceRecord(PseudoInst.LOOKASIDE, inputs=provenance_inputs)) + + +@register_general_jit_lookaside(max) +def _general_jit_builtin_max_lookaside(*args, **kwargs): + return _general_jit_min_max_lookaside("max", clang.py_max, args, kwargs) + + +@register_general_jit_lookaside(min) +def _general_jit_builtin_min_lookaside(*args, **kwargs): + return _general_jit_min_max_lookaside("min", clang.py_min, args, kwargs) + + def _get_torch_nn_module_named_members_lookaside( model: torch.nn.Module, named_member_method, get_member_method, *unwrapped_args, **unwrapped_kwargs ): From e2578e37e494359aad38c615a670b8f3baec543e Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 29 Nov 2025 07:50:20 -0800 Subject: [PATCH 14/46] Add test --- thunder/tests/test_jit_general.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index e7424b7e31..fe49e57e09 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -864,13 +864,16 @@ def foo(a, scalar): assert thunder.cache_hits(jfoo) == 1 -def test_symbolic_values_min_max(): +@pytest.mark.parametrize("min_op,max_op", [(torch.sym_min, torch.sym_max), (min, max)], ids=("torch", "builtin")) +def test_symbolic_values_min_max(min_op, max_op): def foo(seq_len, cumulative_len, max_len, sliding_window) -> int: - max_cache_len = torch.sym_min(max_len, seq_len + 128) + max_cache_len = min_op(max_len, seq_len + 128) offset_raw = cumulative_len - sliding_window + 1 - kv_offset = torch.sym_max(offset_raw, 0) + kv_offset = max_op(offset_raw, 0) return max_cache_len, kv_offset + samples = torch.randint(1, 100, (4, 20)) + jfoo = thunder_jit(foo, cache="symbolic values") samples = torch.randint(1, 100, (20, 4)) From df72dad1185fcbe0a509ee11afe06d77c3bb8f24 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 08:18:07 -0800 Subject: [PATCH 15/46] Remove py_max/min and merge into prims.maximum Co-authored-by: riccardofelluga --- thunder/clang/__init__.py | 18 ------------------ thunder/core/jit_ext.py | 4 ++-- thunder/core/prims.py | 20 ++++---------------- thunder/executors/pythonex.py | 8 ++++---- thunder/torch/__init__.py | 8 ++++---- 5 files changed, 14 insertions(+), 44 deletions(-) diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py index c294d1bb5c..90fe53fdf5 100644 --- a/thunder/clang/__init__.py +++ b/thunder/clang/__init__.py @@ -1743,24 +1743,6 @@ def minimum(a, b): return _elementwise_binary_wrapper(a, b, prim=prims.minimum) -@clangop() -def py_max(a: NumberLike | Number, b: NumberLike | Number): - utils.check( - isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), - lambda: "py_max expects number inputs", - ) - return prims.py_max(a, b) - - -@clangop() -def py_min(a: NumberLike | Number, b: NumberLike | Number): - utils.check( - isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), - lambda: "py_min expects number inputs", - ) - return prims.py_min(a, b) - - @clangop(method_name="mul") def mul(a, b): return _elementwise_binary_wrapper(a, b, prim=prims.mul) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index e4a7ab2593..8c7eb507c5 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -671,12 +671,12 @@ def _general_jit_min_max_lookaside(op_name, symbol, args, kwargs): @register_general_jit_lookaside(max) def _general_jit_builtin_max_lookaside(*args, **kwargs): - return _general_jit_min_max_lookaside("max", clang.py_max, args, kwargs) + return _general_jit_min_max_lookaside("max", clang.maximum, max, args, kwargs) @register_general_jit_lookaside(min) def _general_jit_builtin_min_lookaside(*args, **kwargs): - return _general_jit_min_max_lookaside("min", clang.py_min, args, kwargs) + return _general_jit_min_max_lookaside("min", clang.minimum, min, args, kwargs) def _get_torch_nn_module_named_members_lookaside( diff --git a/thunder/core/prims.py b/thunder/core/prims.py index eba42e7ced..09281bab6a 100644 --- a/thunder/core/prims.py +++ b/thunder/core/prims.py @@ -232,8 +232,6 @@ class PrimIDs(Enum): LT = auto() MAXIMUM = auto() MINIMUM = auto() - PY_MAX = auto() - PY_MIN = auto() MUL = auto() NE = auto() NEXTAFTER = auto() @@ -2693,22 +2691,12 @@ def _div_numbers(a: Number, b: Number) -> Number: supported_input_dtypes=comparison_dtypes, ) -maximum = _make_elementwise_binary_prim(PrimIDs.MAXIMUM, "maximum", supported_input_dtypes=comparison_dtypes) - -minimum = _make_elementwise_binary_prim(PrimIDs.MINIMUM, "minimum", supported_input_dtypes=comparison_dtypes) - -py_max = _make_elementwise_binary_prim( - PrimIDs.PY_MAX, - "py_max", - number_fn=builtins.max, - numbers_only=True, +maximum = _make_elementwise_binary_prim( + PrimIDs.MAXIMUM, "maximum", supported_input_dtypes=comparison_dtypes, number_fn=builtins.max ) -py_min = _make_elementwise_binary_prim( - PrimIDs.PY_MIN, - "py_min", - number_fn=builtins.min, - numbers_only=True, +minimum = _make_elementwise_binary_prim( + PrimIDs.MINIMUM, "minimum", supported_input_dtypes=comparison_dtypes, number_fn=builtins.min ) mul = _make_elementwise_binary_prim( diff --git a/thunder/executors/pythonex.py b/thunder/executors/pythonex.py index 214d44e5a6..aaa9b3b0c8 100644 --- a/thunder/executors/pythonex.py +++ b/thunder/executors/pythonex.py @@ -345,8 +345,8 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten gt = ex.register_operator("gt", like=prims.gt, module=operator) le = ex.register_operator("le", like=prims.le, module=operator) lt = ex.register_operator("lt", like=prims.lt, module=operator) -py_max = ex.register_operator("max", like=prims.py_max, module=builtins) -py_min = ex.register_operator("min", like=prims.py_min, module=builtins) +maximum = ex.register_operator("max", like=prims.maximum, module=builtins) +minimum = ex.register_operator("min", like=prims.minimum, module=builtins) mul = ex.register_operator("mul", like=prims.mul, module=operator) ne = ex.register_operator("ne", like=prims.ne, module=operator) # NOTE pythonex_pow to avoid a name conflict with the builtin pow @@ -374,8 +374,8 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten ex.register_implementation(prims.gt, gt, checker=_elementwise_binary_checker) ex.register_implementation(prims.le, le, checker=_elementwise_binary_checker) ex.register_implementation(prims.lt, lt, checker=_elementwise_binary_checker) -ex.register_implementation(prims.py_max, py_max, checker=_elementwise_binary_checker) -ex.register_implementation(prims.py_min, py_min, checker=_elementwise_binary_checker) +ex.register_implementation(prims.maximum, maximum, checker=_elementwise_binary_checker) +ex.register_implementation(prims.minimum, minimum, checker=_elementwise_binary_checker) ex.register_implementation(prims.mul, mul, checker=_elementwise_binary_checker) ex.register_implementation(prims.ne, ne, checker=_elementwise_binary_checker) # NOTE pythonex_pow to avoid a name conflict with the builtin pow diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py index 63fdc8c8e9..49d2e85c9b 100644 --- a/thunder/torch/__init__.py +++ b/thunder/torch/__init__.py @@ -2600,18 +2600,18 @@ def minimum(a: TensorProxy, b: TensorProxy) -> TensorProxy: def sym_max(a: NumberLike, b: NumberLike) -> NumberLike: utils.check( isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), - lambda: "torch.sym_max currently supports only number inputs", + lambda: "torch.sym_max supports only number inputs", ) - return clang.py_max(a, b) + return clang.maximum(a, b) @torchsymbol(torch.sym_min, id="torch.sym_min") def sym_min(a: NumberLike, b: NumberLike) -> NumberLike: utils.check( isinstance(a, (Number, NumberProxy)) and isinstance(b, (Number, NumberProxy)), - lambda: "torch.sym_min currently supports only number inputs", + lambda: "torch.sym_min supports only number inputs", ) - return clang.py_min(a, b) + return clang.minimum(a, b) # NOTE This is just an alias for proxies to find operation defined for the modulus From 2806c4019d2232c8a2d4f081aacf1ee2816c4412 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 08:36:36 -0800 Subject: [PATCH 16/46] Fix up --- thunder/core/jit_ext.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 8c7eb507c5..4b96bb1362 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -671,12 +671,12 @@ def _general_jit_min_max_lookaside(op_name, symbol, args, kwargs): @register_general_jit_lookaside(max) def _general_jit_builtin_max_lookaside(*args, **kwargs): - return _general_jit_min_max_lookaside("max", clang.maximum, max, args, kwargs) + return _general_jit_min_max_lookaside("max", clang.maximum, args, kwargs) @register_general_jit_lookaside(min) def _general_jit_builtin_min_lookaside(*args, **kwargs): - return _general_jit_min_max_lookaside("min", clang.minimum, min, args, kwargs) + return _general_jit_min_max_lookaside("min", clang.minimum, args, kwargs) def _get_torch_nn_module_named_members_lookaside( From 37bfd92f45c2e5c1a5d9b756ab09b9c080d9b79f Mon Sep 17 00:00:00 2001 From: beverlylytle Date: Thu, 16 Oct 2025 14:09:54 +0300 Subject: [PATCH 17/46] Apply DCE to subsymbols --- thunder/core/transform_common.py | 34 +++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py index bc14a671f2..cbdc1064e9 100644 --- a/thunder/core/transform_common.py +++ b/thunder/core/transform_common.py @@ -131,7 +131,9 @@ def keep_or_swap(p): if all(map(lambda x: isinstance(x, NumberProxyInterface) and x.name in seen, bsym.flat_outs)): continue output = tree_map(keep_or_swap, bsym.output) - new_bsyms.append(bsym.from_bsym(output=output)) + subsymbols = dce(bsym.subsymbols, output=bsym.output) + new_bsym = bsym.from_bsym(output=output, subsymbols=subsymbols) + new_bsyms.append(new_bsym) return new_bsyms @@ -142,12 +144,21 @@ def keep_or_swap(p): # that only produce non-proxy objects # NOTE needed_proxies is an in/out argument, it takes an initial set of Variables you want to keep, and return # all the needed proxies of the input trace -def dce(trace: Trace, needed_proxies: None | set[Variable] = None) -> Trace: +def dce( + trace_or_bsyms: Trace | list[BoundSymbolInterface], needed_proxies: None | set[Variable] = None, output=None +) -> Trace | list[BoundSymbolInterface]: start_time_ns = time.perf_counter_ns() - producer_map: ProxyDict = producers(trace) + producer_map: ProxyDict = producers(trace_or_bsyms) + + if isinstance(trace_or_bsyms, Trace): + bound_symbols = trace_or_bsyms.bound_symbols + output = trace_or_bsyms.output + else: + bound_symbols = trace_or_bsyms + output = output - flat_trace_outputs, _ = tree_flatten(trace.output) + flat_trace_outputs, _ = tree_flatten(output) if needed_proxies is None: needed_proxies: set[Variable] = set(tuple(variableify(x) for x in flat_trace_outputs if isinstance(x, Proxy))) else: @@ -155,7 +166,7 @@ def dce(trace: Trace, needed_proxies: None | set[Variable] = None) -> Trace: dced = [] bsym: BoundSymbol - for bsym in reversed(trace.bound_symbols): + for bsym in reversed(bound_symbols): # Preserves symbols that should never be collected if has_tags(bsym, {prims.OpTags.DONT_DCE}): needed = True @@ -182,19 +193,24 @@ def dce(trace: Trace, needed_proxies: None | set[Variable] = None) -> Trace: for x in nbsym.flat_proxy_args: needed_proxies.add(variableify(x)) - dcetrace = from_trace(trace) dced_bound_symbols = list(reversed(dced)) # duplicate number proxies happen with the symbolic shapes and are # not covered by the above (due to being in tuples?). dced_bound_symbols = remove_duplicate_number_proxies(dced_bound_symbols) - dcetrace.bound_symbols = dced_bound_symbols + if isinstance(trace_or_bsyms, Trace): + result = from_trace(trace_or_bsyms) + result.bound_symbols = dced_bound_symbols + else: + result = dced_bound_symbols end_time_ns = time.perf_counter_ns() elapsed_time_ns = end_time_ns - start_time_ns elapsed_time_millis = elapsed_time_ns // 1000000 - dcetrace.set_provenance(TraceProvenance(f"Dead Code Elimination (took {elapsed_time_millis} milliseconds)")) - return dcetrace + if isinstance(trace_or_bsyms, Trace): + result.set_provenance(TraceProvenance(f"Dead Code Elimination (took {elapsed_time_millis} milliseconds)")) + + return result # From 5261098a329e4d907f2dd2887239d697a6cfc36f Mon Sep 17 00:00:00 2001 From: beverlylytle Date: Wed, 12 Nov 2025 14:41:33 +0200 Subject: [PATCH 18/46] try in symbol.__call__ instead --- thunder/core/symbol.py | 6 +++++- thunder/core/transform_common.py | 4 +--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/thunder/core/symbol.py b/thunder/core/symbol.py index f5184b5770..30b25cff8d 100644 --- a/thunder/core/symbol.py +++ b/thunder/core/symbol.py @@ -323,6 +323,11 @@ def __call__(self, *args, **kwargs): result = tree_unflatten(flat_results, spec) + # When using symbolic values, there may be duplicate prims.eq and prims.shape subsymbols that can be removed. + from thunder.core.transform_common import dce + + subsymbols = dce(subsymbols, output=result) + trace.pop_scope() cd = get_compile_data() @@ -340,7 +345,6 @@ def tag_tensorproxy_output_as_detached(proxy): return proxy result = tree_map(tag_tensorproxy_output_as_detached, result) - bsym = self.bind(*args, **kwargs, output=result, subsymbols=subsymbols) symbols_list = trace.peek_scope() diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py index cbdc1064e9..2ecdda3d19 100644 --- a/thunder/core/transform_common.py +++ b/thunder/core/transform_common.py @@ -131,9 +131,7 @@ def keep_or_swap(p): if all(map(lambda x: isinstance(x, NumberProxyInterface) and x.name in seen, bsym.flat_outs)): continue output = tree_map(keep_or_swap, bsym.output) - subsymbols = dce(bsym.subsymbols, output=bsym.output) - new_bsym = bsym.from_bsym(output=output, subsymbols=subsymbols) - new_bsyms.append(new_bsym) + new_bsyms.append(bsym.from_bsym(output=output)) return new_bsyms From 681492f98fdfbf047af0edcfc3c5a4e40ed179a7 Mon Sep 17 00:00:00 2001 From: beverlylytle Date: Fri, 14 Nov 2025 16:31:15 +0200 Subject: [PATCH 19/46] come on, ruff, it's a test --- thunder/core/symbol.py | 11 ++++--- thunder/core/transform_common.py | 16 +++++++-- thunder/dynamo/utils.py | 15 +++++---- thunder/executors/nvfuserex_impl.py | 7 ++-- thunder/tests/test_core.py | 51 +++++++++++++++++++++++++---- 5 files changed, 75 insertions(+), 25 deletions(-) diff --git a/thunder/core/symbol.py b/thunder/core/symbol.py index 30b25cff8d..6dd65cfaa5 100644 --- a/thunder/core/symbol.py +++ b/thunder/core/symbol.py @@ -323,11 +323,6 @@ def __call__(self, *args, **kwargs): result = tree_unflatten(flat_results, spec) - # When using symbolic values, there may be duplicate prims.eq and prims.shape subsymbols that can be removed. - from thunder.core.transform_common import dce - - subsymbols = dce(subsymbols, output=result) - trace.pop_scope() cd = get_compile_data() @@ -354,6 +349,12 @@ def tag_tensorproxy_output_as_detached(proxy): exception_type=AssertionError, ) + # When using symbolic values, there may be duplicate prims.eq and prims.shape subsymbols that can be removed. + from thunder.core.transform_common import dce + + subsymbols = dce(subsymbols, output=result) + bsym = bsym.from_bsym(subsymbols=subsymbols) + symbols_list.append(bsym) return result diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py index 2ecdda3d19..928e94eeb0 100644 --- a/thunder/core/transform_common.py +++ b/thunder/core/transform_common.py @@ -143,8 +143,21 @@ def keep_or_swap(p): # NOTE needed_proxies is an in/out argument, it takes an initial set of Variables you want to keep, and return # all the needed proxies of the input trace def dce( - trace_or_bsyms: Trace | list[BoundSymbolInterface], needed_proxies: None | set[Variable] = None, output=None + trace_or_bsyms: Trace | list[BoundSymbolInterface], + needed_proxies: None | set[Variable] = None, + output: Any = None, ) -> Trace | list[BoundSymbolInterface]: + """Runs a Dead Code Elimination (DCE) pass + + Args: + trace_or_bsyms: The trace or list of bound symbols to run the DCE pass on. + needed_proxies: The set of variables to keep. + output: The output of the list of bound symbols. This is only used if the input is a list of bound + symbols, and is required in that case + + Returns: + The trace (if the input is a trace) or list of bound symbols (if the input is a list of bound symbols) after the DCE pass. + """ start_time_ns = time.perf_counter_ns() producer_map: ProxyDict = producers(trace_or_bsyms) @@ -154,7 +167,6 @@ def dce( output = trace_or_bsyms.output else: bound_symbols = trace_or_bsyms - output = output flat_trace_outputs, _ = tree_flatten(output) if needed_proxies is None: diff --git a/thunder/dynamo/utils.py b/thunder/dynamo/utils.py index 8c643cee71..a393d7bb06 100644 --- a/thunder/dynamo/utils.py +++ b/thunder/dynamo/utils.py @@ -259,20 +259,20 @@ def get_backed_value(s): return tuple(map(get_backed_value, vals)) -def get_proxy_inputs_from_node(node: torch.fx.Node) -> tuple[tuple, dict]: +def get_proxy_inputs_from_node(node: torch.fx.Node, tracectx) -> tuple[tuple, dict]: """Creates proxy inputs from a torch.fx.Node for use with Thunder. This function generates proxy inputs for a given torch.fx.Node Args: node (torch.fx.Node): The FX graph node to create proxy inputs for. + tracectx (TraceCtx): The trace context to use to generate proxy inputs. """ import thunder - from thunder.core.trace import TraceCtx from thunder.core.proxies import proxy # We need to be under trace context to generate proxies. - with thunder.core.trace.tracectx(TraceCtx()): + with thunder.core.trace.tracectx(tracectx): def make_input_proxy(arg_node): # This is a Node in the graph representing a Tensor or tuple of Tensors or @@ -380,8 +380,10 @@ def _run_with_cache_info(): cache_info["default_dtype"] = torch.get_default_dtype() cache_info["default_device"] = torch.get_default_device() + tracectx = TraceCtx() + try: - proxy_args, proxy_kwargs = get_proxy_inputs_from_node(node) + proxy_args, proxy_kwargs = get_proxy_inputs_from_node(node, tracectx) except Exception as e: return False, SplitReason( SplitReasonType.EXCEPTION_PROXY_THUNDER_OP, @@ -395,7 +397,7 @@ def _run_with_cache_info(): else thunder_symbol ) # We need to be under trace context to generate proxies. - with thunder.core.trace.tracectx(TraceCtx()): + with thunder.core.trace.tracectx(tracectx): try: function_to_run(*proxy_args, **proxy_kwargs) except Exception as e: @@ -478,6 +480,7 @@ def is_node_supported_by_thunder( """ Determine whether thunder can execute the operation described by this node. """ + from thunder.core.trace import TraceCtx # Docs from the torch.fx.Node - https://pytorch.org/docs/stable/fx.html#torch.fx.Node # Each Node has a function specified by its op property # Below are the details for the ones this function is interested in - @@ -555,7 +558,7 @@ def is_node_supported_by_thunder( if torchctx.has_method(node.target): # `torchctx.get_method` requires args and kwargs to resolve which overload of the method is picked. try: - args, kwargs = get_proxy_inputs_from_node(node) + args, kwargs = get_proxy_inputs_from_node(node, TraceCtx()) except Exception as e: return False, SplitReason( SplitReasonType.EXCEPTION_PROXY_THUNDER_OP, diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index 4d89116153..f1f2ade5a1 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -707,14 +707,11 @@ def has_cuda_input_or_output(self, bsym: BoundSymbol) -> bool: return False def _dce_bsyms(self, input_list, output, bsyms: list[BoundSymbol]) -> list[BoundSymbol]: - trace = TraceCtx(None) - trace.bound_symbols = bsyms - bsyms.append(prims.python_return.bind(output, output=None)) needed_proxies: set[Variable] = set() - trace = dce(trace, needed_proxies) + bsyms = dce(bsyms, needed_proxies, output) # update the input_list by removing the unused inputs input_list[:] = [x for x in input_list if variableify(x) in needed_proxies] - return list(filter(lambda x: x.sym != prims.python_return, trace.bound_symbols)) + return bsyms def fuse(self, region: Region, fusion_counter: int) -> BoundSymbol: sorted_unique_inputs: list[Proxy] = [unvariableify(x) for x in region.inputs] diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index 9e5d9687f6..4fe1c5ba88 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -2170,6 +2170,36 @@ def func(x, y, device): assert [t.name for t in tree_flatten(flatten_cse_trace.output)[0]] == ["t4", "t4", "t6", "t14", "t15", "t16", "t17"] +@instantiate( + dtypes=NOTHING, +) +def test_dce(executor, device, _): + def func(x): + dead_code = x + 1 # noqa: F841 + y = x * x + return y + + x = make_tensor((2, 2), device=device, dtype=torch.float32) + compiled = thunder.jit(func, executors=executor.executors_list()) + compiled(x) + traces = thunder.last_traces(compiled) + + # find last trace before DCE is applied + for i, trace in enumerate(traces): + provenance = trace.get_provenance().pss if trace.get_provenance() else "" + if "Dead Code Elimination" in provenance: + break + i -= 1 + trace = traces[i] + + from thunder.core.transform_common import dce + + dced_trace = dce(trace) + dced_bsyms = dce(trace.bound_symbols) + assert len(dced_trace.bound_symbols) == len(trace.bound_symbols) - 1 + assert len(dced_trace.bound_symbols) == len(dced_bsyms) + + def test_symbol_flat_args(): from thunder.core.symbol import Symbol, BoundSymbol @@ -3295,22 +3325,29 @@ def clean(tr): def test_prims_pack_list(): - def foo(): - pass - - trace = TraceCtx(foo) + def foo(x): + a, b = x + return [a, b] a = torch.randn(2, 2) b = torch.randn(2, 2) + jfoo = thunder.jit(foo) + jfoo((a, b)) + + trace = thunder.last_traces(jfoo)[-1] + + return_bsym = trace.bound_symbols[-1] + trace.bound_symbols = trace.bound_symbols[:-1] + with tracectx(trace): - x = prims.unpack_trivial(a, name="x") - y = prims.unpack_trivial(b, name="y") + x, y = return_bsym.flat_args packed_list = prims.pack_list(x, y) prims.python_return(packed_list) func = trace.python_callable() - actual = func() + print(trace) + actual = func(a, b) expected = [a, b] assert isinstance(actual, list) and actual == expected From 6af4d203ae9c8906faffb0c1b82dfb97fb69a4fd Mon Sep 17 00:00:00 2001 From: beverlylytle Date: Tue, 18 Nov 2025 14:41:51 +0200 Subject: [PATCH 20/46] remove print --- thunder/tests/test_core.py | 1 - 1 file changed, 1 deletion(-) diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index 4fe1c5ba88..eeea8f492d 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -3346,7 +3346,6 @@ def foo(x): prims.python_return(packed_list) func = trace.python_callable() - print(trace) actual = func(a, b) expected = [a, b] From 10d0564a13a702bda85a3459017ec2e9d9a73e2b Mon Sep 17 00:00:00 2001 From: beverlylytle Date: Fri, 21 Nov 2025 11:49:21 +0200 Subject: [PATCH 21/46] respond to comments --- thunder/core/symbol.py | 5 ++-- thunder/core/transform_common.py | 45 +++++++++++++---------------- thunder/executors/nvfuserex_impl.py | 4 +-- 3 files changed, 25 insertions(+), 29 deletions(-) diff --git a/thunder/core/symbol.py b/thunder/core/symbol.py index 6dd65cfaa5..02d4533b54 100644 --- a/thunder/core/symbol.py +++ b/thunder/core/symbol.py @@ -340,6 +340,7 @@ def tag_tensorproxy_output_as_detached(proxy): return proxy result = tree_map(tag_tensorproxy_output_as_detached, result) + bsym = self.bind(*args, **kwargs, output=result, subsymbols=subsymbols) symbols_list = trace.peek_scope() @@ -350,9 +351,9 @@ def tag_tensorproxy_output_as_detached(proxy): ) # When using symbolic values, there may be duplicate prims.eq and prims.shape subsymbols that can be removed. - from thunder.core.transform_common import dce + from thunder.core.transform_common import dce_bsyms - subsymbols = dce(subsymbols, output=result) + subsymbols = dce_bsyms(subsymbols, result) bsym = bsym.from_bsym(subsymbols=subsymbols) symbols_list.append(bsym) diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py index 928e94eeb0..32916cfded 100644 --- a/thunder/core/transform_common.py +++ b/thunder/core/transform_common.py @@ -142,31 +142,22 @@ def keep_or_swap(p): # that only produce non-proxy objects # NOTE needed_proxies is an in/out argument, it takes an initial set of Variables you want to keep, and return # all the needed proxies of the input trace -def dce( - trace_or_bsyms: Trace | list[BoundSymbolInterface], +def dce_bsyms( + bsyms: list[BoundSymbolInterface], + output: Any, needed_proxies: None | set[Variable] = None, - output: Any = None, ) -> Trace | list[BoundSymbolInterface]: """Runs a Dead Code Elimination (DCE) pass Args: - trace_or_bsyms: The trace or list of bound symbols to run the DCE pass on. + bsyms: The list of bound symbols to run the DCE pass on. needed_proxies: The set of variables to keep. - output: The output of the list of bound symbols. This is only used if the input is a list of bound - symbols, and is required in that case + output: The output of the list of bound symbols. Returns: - The trace (if the input is a trace) or list of bound symbols (if the input is a list of bound symbols) after the DCE pass. + The list of bound symbols after the DCE pass. """ - start_time_ns = time.perf_counter_ns() - - producer_map: ProxyDict = producers(trace_or_bsyms) - - if isinstance(trace_or_bsyms, Trace): - bound_symbols = trace_or_bsyms.bound_symbols - output = trace_or_bsyms.output - else: - bound_symbols = trace_or_bsyms + producer_map: ProxyDict = producers(bsyms) flat_trace_outputs, _ = tree_flatten(output) if needed_proxies is None: @@ -176,7 +167,7 @@ def dce( dced = [] bsym: BoundSymbol - for bsym in reversed(bound_symbols): + for bsym in reversed(bsyms): # Preserves symbols that should never be collected if has_tags(bsym, {prims.OpTags.DONT_DCE}): needed = True @@ -208,18 +199,22 @@ def dce( # not covered by the above (due to being in tuples?). dced_bound_symbols = remove_duplicate_number_proxies(dced_bound_symbols) - if isinstance(trace_or_bsyms, Trace): - result = from_trace(trace_or_bsyms) - result.bound_symbols = dced_bound_symbols - else: - result = dced_bound_symbols + return dced_bound_symbols + + +def dce(trace: Trace, needed_proxies: set[Variable]) -> Trace: + start_time_ns = time.perf_counter_ns() + + bsyms = trace.bound_symbol + dced_bsyms = dce_bsyms(bsyms, trace.output, needed_proxies) + result = from_trace(trace) + result.bound_symbols = dced_bsyms + end_time_ns = time.perf_counter_ns() elapsed_time_ns = end_time_ns - start_time_ns elapsed_time_millis = elapsed_time_ns // 1000000 - if isinstance(trace_or_bsyms, Trace): - result.set_provenance(TraceProvenance(f"Dead Code Elimination (took {elapsed_time_millis} milliseconds)")) - + result.set_provenance(TraceProvenance(f"Dead Code Elimination (took {elapsed_time_millis} milliseconds)")) return result diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index f1f2ade5a1..3782a0f3e5 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -43,7 +43,7 @@ from thunder.core.trace import TraceCtx, from_trace, TraceProvenance from thunder.core.symbol import BoundSymbol, BoundSymbolRHS, Symbol, has_tags from thunder.core.devices import Device, DeviceType, cpu -from thunder.core.transform_common import dce, cse_single_bsym, replace_redundant_inputs +from thunder.core.transform_common import dce, dce_bsyms, cse_single_bsym, replace_redundant_inputs from thunder.core.profile import annotate_for_profile from thunder.core.compile_data import get_compile_option from thunder.torch.experimental.dtensor_torch_and_prims import DTensorPrimIDs @@ -708,7 +708,7 @@ def has_cuda_input_or_output(self, bsym: BoundSymbol) -> bool: def _dce_bsyms(self, input_list, output, bsyms: list[BoundSymbol]) -> list[BoundSymbol]: needed_proxies: set[Variable] = set() - bsyms = dce(bsyms, needed_proxies, output) + bsyms = dce_bsyms(bsyms, output, needed_proxies) # update the input_list by removing the unused inputs input_list[:] = [x for x in input_list if variableify(x) in needed_proxies] return bsyms From 2b7d2b38740f60a7144f8c015065f449011eb637 Mon Sep 17 00:00:00 2001 From: beverlylytle Date: Fri, 21 Nov 2025 12:01:41 +0200 Subject: [PATCH 22/46] where's my coffee --- thunder/core/rematerialization.py | 2 +- thunder/core/transform_common.py | 4 ++-- thunder/tests/test_core.py | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/thunder/core/rematerialization.py b/thunder/core/rematerialization.py index a3b3b83504..44a8a9cd67 100644 --- a/thunder/core/rematerialization.py +++ b/thunder/core/rematerialization.py @@ -458,7 +458,7 @@ def rematerialize(trace: TraceCtx) -> TraceCtx: computed_cuts_for_producers[producer] += cut rematerialized_trace = from_trace(trace) - rematerialized_trace.bound_symbols = tuple(new_bsyms.get(bsym, bsym) for bsym in trace.bound_symbols) + rematerialized_trace.bound_symbols = list(new_bsyms.get(bsym, bsym) for bsym in trace.bound_symbols) end_time_ns = time.perf_counter_ns() elapsed_time_ns = end_time_ns - start_time_ns diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py index 32916cfded..b2f6ca8140 100644 --- a/thunder/core/transform_common.py +++ b/thunder/core/transform_common.py @@ -202,10 +202,10 @@ def dce_bsyms( return dced_bound_symbols -def dce(trace: Trace, needed_proxies: set[Variable]) -> Trace: +def dce(trace: Trace, needed_proxies: set[Variable] = None) -> Trace: start_time_ns = time.perf_counter_ns() - bsyms = trace.bound_symbol + bsyms = trace.bound_symbols dced_bsyms = dce_bsyms(bsyms, trace.output, needed_proxies) result = from_trace(trace) result.bound_symbols = dced_bsyms diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index eeea8f492d..d2cfb938ce 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -2192,10 +2192,10 @@ def func(x): i -= 1 trace = traces[i] - from thunder.core.transform_common import dce + from thunder.core.transform_common import dce, dce_bsyms dced_trace = dce(trace) - dced_bsyms = dce(trace.bound_symbols) + dced_bsyms = dce_bsyms(trace.bound_symbols, trace.output) assert len(dced_trace.bound_symbols) == len(trace.bound_symbols) - 1 assert len(dced_trace.bound_symbols) == len(dced_bsyms) From d9e4f9c86378e4488a88dbced1be416fdace4692 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Wed, 10 Dec 2025 14:23:58 -0800 Subject: [PATCH 23/46] MAKE SYMBOLIC VALUES DEFAULT --- thunder/core/options.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/core/options.py b/thunder/core/options.py index bc64145d4b..28eadab407 100644 --- a/thunder/core/options.py +++ b/thunder/core/options.py @@ -68,7 +68,7 @@ def _string_to_cache_option(s: str, /) -> None | CACHE_OPTIONS: def resolve_cache_option(x: Any, /) -> CACHE_OPTIONS: co: None | CACHE_OPTIONS if x is None: - co = CACHE_OPTIONS.CONSTANT_VALUES + co = CACHE_OPTIONS.SYMBOLIC_VALUES elif isinstance(x, CACHE_OPTIONS): co = x elif isinstance(x, str): From 8293bd5594d139eb31a196f4c88d4511f5c80906 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 16:37:36 -0800 Subject: [PATCH 24/46] Shift bsym indices --- thunder/tests/test_core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index d2cfb938ce..eb22bffd99 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -1121,8 +1121,8 @@ def bar(a, shape): top_down_bsyms = toposort_bsym_dag(roots, TOPOSORT_ORDER.TOP_DOWN) bottom_up_bsyms = toposort_bsym_dag(leaves, TOPOSORT_ORDER.BOTTOM_UP) - top_down_reshape_bsym = top_down_bsyms[3] - bottom_up_reshape_bsym = bottom_up_bsyms[2] + top_down_reshape_bsym = top_down_bsyms[5] + bottom_up_reshape_bsym = bottom_up_bsyms[4] assert top_down_reshape_bsym.sym.id == "torch.reshape" assert bottom_up_reshape_bsym.sym.id == "torch.reshape" From 2e0015fe43efe5f25355ddd76ce299ae939486ea Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 16:37:50 -0800 Subject: [PATCH 25/46] Add torch.set_autocast_enabled --- thunder/torch/default_torch_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/thunder/torch/default_torch_ops.py b/thunder/torch/default_torch_ops.py index 195625aa3a..24217c6ab4 100644 --- a/thunder/torch/default_torch_ops.py +++ b/thunder/torch/default_torch_ops.py @@ -2,6 +2,7 @@ torch_auto_registered_ops = { torch: [ + torch.set_autocast_enabled, torch._native_multi_head_attention, torch.lobpcg, torch.unravel_index, From 211a2a4b0d61e15c10b59d091253bbb93a443c1b Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 16:38:09 -0800 Subject: [PATCH 26/46] Add StringProxy.__bool__ --- thunder/core/jit_ext.py | 3 ++- thunder/core/proxies.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 4b96bb1362..2f962c48b2 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -38,6 +38,7 @@ Proxy, ProxyInterface, ProxyTag, + StringProxy, TensorProxy, Variable, is_proxy_name_available, @@ -642,7 +643,7 @@ def _general_jit_hasattr_lookaside(obj: Any, name: str): def _general_jit_bool_lookaside(wrapped_x: Any) -> bool | INTERPRETER_SIGNALS: assert isinstance(wrapped_x, WrappedValue) # It doesn't feel right to insert constraints in bool lookaside, constraints here only applies when the bool value is used in control flow. - if isinstance(wrapped_x.value, NumberProxy): + if isinstance(wrapped_x.value, (NumberProxy, StringProxy)): if wrapped_x.value.is_dynamic(): raise NotImplementedError(f"conversion to bool is not allowed on dynamic proxy={wrapped_x.value}") wrapped_x.value.make_static_constrained() diff --git a/thunder/core/proxies.py b/thunder/core/proxies.py index d2f8e58e2c..25fc08b775 100644 --- a/thunder/core/proxies.py +++ b/thunder/core/proxies.py @@ -455,6 +455,9 @@ def __eq__(self, other): return False return str(self) == str(other) + def __bool__(self) -> bool: + return bool(self.value) + # # Collection proxies From 97a0032cb82ef321c5ded28653b444647931c59f Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 17:13:06 -0800 Subject: [PATCH 27/46] Defer cotangents creation after forward pass --- thunder/core/transforms.py | 34 +++++++++++++++++++--------------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py index 11bf85f4bd..b1e4c110a4 100644 --- a/thunder/core/transforms.py +++ b/thunder/core/transforms.py @@ -3031,10 +3031,24 @@ def vjp_call(primals, cotangents, trace: Trace, **kwargs): primals = (primals,) result, env = augmented_forward_pass(*primals, trace=trace, **kwargs) - check( - len(result) == len(cotangents) if isinstance(result, Sequence) else True, - lambda: f"Expected cotangents to be a sequence of length {len(result)}, got a sequence of length {len(cotangents)}", - ) + + if cotangents is None: + + def ones_like(x): + if isinstance(x, TensorProxy): + return full_like(x, fill_value=1) + elif isinstance(x, NumberProxy): + return type(x.value)(1) + else: + return None + + cotangents = tree_map(lambda v: ones_like(v), result) + else: + check( + len(result) == len(cotangents) if isinstance(result, Sequence) else True, + lambda: f"Expected cotangents to be a sequence of length {len(result)}, got a sequence of length {len(cotangents)}", + ) + return result, backward_pass(env, trace, cotangents) @@ -3075,18 +3089,8 @@ def value_and_grad(func): func (Callable): Function to be differentiated. """ - def ones_like(x): - if isinstance(x, TensorProxy): - return full_like(x, fill_value=1) - elif isinstance(x, NumberProxy): - return type(x.value)(1) - else: - return None - def _value_and_grad(*args, **kwargs): - trace = construct_trace()(func, *args, **kwargs) - cotangents = tree_map(lambda v: ones_like(v), trace.output) - return vjp(func)(args, cotangents, **kwargs) + return vjp(func)(args, None, **kwargs) return _value_and_grad From afb6a5af9de71e06d4975ca5483caa3d3dad32db Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 18:15:06 -0800 Subject: [PATCH 28/46] .numel needs a preceding prims.shape --- thunder/core/proxies.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/thunder/core/proxies.py b/thunder/core/proxies.py index 25fc08b775..96307d30b2 100644 --- a/thunder/core/proxies.py +++ b/thunder/core/proxies.py @@ -1275,7 +1275,7 @@ def _infer_tensor_properties( else: # deferred computation of numel # TODO: similar to how `shape` is handled, this should be CSE or lifted for efficiency - _numel = lambda *args: reduce(operator.mul, _shape, 1) + _numel = lambda self: reduce(operator.mul, self.shape, 1) # TODO Alias rank to ndim? _ndim = len(_shape) @@ -1465,7 +1465,7 @@ def __init__( self._device, self._dtype, self._true_dtype, - self._numel, + _numel, self._ndim, self._requires_grad, self._grad, @@ -1482,6 +1482,11 @@ def __init__( thunder_fsdp_padding_size, ) + if not using_symbolic_values(): + self._numel = _numel + else: + self._numel = lambda self=self: _numel(self) + # NOTE The following properties DO NOT depend on the language context or record # themselves into the trace, so they can be used when working with tensor proxies # outside of a trace or language context From 77d15b49c533234152a4db18a1a0207ec22ad5fa Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 18:24:06 -0800 Subject: [PATCH 29/46] Prologue can't be skipped --- thunder/tests/test_grad.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py index 7e17fb04a4..7009393cb6 100644 --- a/thunder/tests/test_grad.py +++ b/thunder/tests/test_grad.py @@ -621,7 +621,7 @@ def _fdm_jvp(): comp, executor, set_compile_data, - len(sample.kwargs) != 0, + prologue_required=True, ) def _torch_jvp(): From a2621503d894b775608aea8933a752987e65007c Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 18:45:23 -0800 Subject: [PATCH 30/46] Support set_grad_enabled(symbolic) tested by test_set_grad_enabled --- thunder/__init__.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/thunder/__init__.py b/thunder/__init__.py index c0483f8c3a..8793ebec7a 100644 --- a/thunder/__init__.py +++ b/thunder/__init__.py @@ -15,7 +15,7 @@ # imports unused in this file, but referenced as thunder.* elsewhere from thunder.common import trace import thunder.core.devices as devices -from thunder.core.proxies import Proxy +from thunder.core.proxies import NumberProxy, Proxy from thunder.common import ( CompileData, @@ -895,7 +895,12 @@ def fn_(*args, **kwargs) -> Any: result = call_epilogue(cache_entry, result, pro_to_epi) # Reflect the state of is_grad_enabled, as its changes were tracked only inside Thunder - pytorch.set_grad_enabled(cd.is_grad_enabled) + is_grad_enabled = cd.is_grad_enabled + if isinstance(is_grad_enabled, NumberProxy): + # TODO: Verify this assumption + assert is_grad_enabled.is_static_constrained() + is_grad_enabled = is_grad_enabled.value + pytorch.set_grad_enabled(is_grad_enabled) cs.last_computation = cache_entry.computation_fn return result From b9df532f11ce96dd78b0e319856fd872ce936446 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 18:51:06 -0800 Subject: [PATCH 31/46] Support isinstance(symbolic, type) test: thunder/tests/test_core.py::test_integer_isinstance_mimicry --- thunder/core/jit_ext.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 2f962c48b2..0eae9e1b9c 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -469,6 +469,8 @@ def _general_jit_getattr_lookaside(obj: Any, name: str, *maybe_default: Any): @register_general_jit_lookaside(isinstance) def _general_jit_isinstance_lookaside(obj: Any, cls: type | UnionType | tuple[type | UnionType]): + from thunder.core.baseutils import check + uobj = unwrap(obj) ucls = unwrap(cls) if isinstance(uobj, TensorProxy): @@ -480,6 +482,9 @@ def _general_jit_isinstance_lookaside(obj: Any, cls: type | UnionType | tuple[ty ucls = (ucls,) if torch.nn.Parameter in ucls: res = issubclass(obj.python_typ, ucls) + elif isinstance(uobj, NumberProxy): + check(uobj.value is not None, lambda: "isinstance does not support NumberProxy with no value") + res = isinstance(uobj.value, ucls) else: res = isinstance(uobj, ucls) From 06baafb709749c8bddca19e7fa5df8f9aef93036 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 18:51:51 -0800 Subject: [PATCH 32/46] a.numel() instead of a.numel --- thunder/executors/nvfuserex_impl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index 3782a0f3e5..9aa9def5ba 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -1360,7 +1360,7 @@ def squeeze(a: TensorProxy, /, dims: Sequence[int], *, fd: FusionDefinition, lc_ # So for now, it'd be reasonable to disallow 0-size tensors. # Related: https://github.com/Lightning-AI/lightning-thunder/issues/2068 def _take_check(a: TensorProxy, /, index: TensorProxy, dim: int) -> bool: - return are_supported_tensors(a, index) and a.numel > 0 + return are_supported_tensors(a, index) and a.numel() > 0 def take(a: TensorProxy, /, index: TensorProxy, dim: int, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any: From ec8ce20aaa4bbe2e7473e1a5cdc03e541020fafc Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 19:10:21 -0800 Subject: [PATCH 33/46] Remove _take_check numel check --- thunder/executors/nvfuserex_impl.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index 9aa9def5ba..a0677087f2 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -1355,12 +1355,8 @@ def squeeze(a: TensorProxy, /, dims: Sequence[int], *, fd: FusionDefinition, lc_ register_supported(PrimIDs.SQUEEZE, squeeze, _squeeze_check) -# NOTE: Currently `_advanced_indexing` seems to return a `TensorProxy` of wrong shape -# when input is 0-size tensor, leading to a broken nvfuser definition. -# So for now, it'd be reasonable to disallow 0-size tensors. -# Related: https://github.com/Lightning-AI/lightning-thunder/issues/2068 def _take_check(a: TensorProxy, /, index: TensorProxy, dim: int) -> bool: - return are_supported_tensors(a, index) and a.numel() > 0 + return are_supported_tensors(a, index) def take(a: TensorProxy, /, index: TensorProxy, dim: int, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any: From b965ef99d82ba4009ed2ace6be2638c307d0d252 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 2 Dec 2025 20:29:36 -0800 Subject: [PATCH 34/46] FIXME Allow a cycle in swap_map in OpExProcessor grad_transform_on_trace creates a cycle in the proxy names of tensor shapes. Specifically, it translates tA = prims.cat([tB, tC], -1) # "meta f32[..., iX]" (..., iX) = prims.shape(tA) into ta = prims.cat([tb, tc], -1) # "meta f32[..., ix]" (..., iX) = prims.shape(ta) This creates a translation ix -> iX. This leads to a cycle when the transform also produces iX -> ix. Repro: thunder/tests/test_jit_general.py::test_litgpt_variants[meta-long-context-like] --- thunder/core/symbol.py | 12 +++++++++--- thunder/core/trace_interpreter.py | 9 +++++---- thunder/executors/passes.py | 2 +- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/thunder/core/symbol.py b/thunder/core/symbol.py index 02d4533b54..296cf95442 100644 --- a/thunder/core/symbol.py +++ b/thunder/core/symbol.py @@ -453,6 +453,7 @@ def from_bsym_swap_proxies( skip_inputs: bool = False, skip_output: bool = False, skip_subsymbols: bool = False, + allow_cycles: bool = False, ) -> BoundSymbol: """Create a new :class:`BoundSymbol` with its inputs, output, and subsymbols updated with ``swap_map``. @@ -487,9 +488,14 @@ def swap(c): while vfa in swap_map: if swap_map[vfa] is fa: break - baseutils.check( - vfa not in visited, lambda: f"Detected a cycle while swapping; the cycle includes {visited}" - ) + + if vfa in visited: + baseutils.check( + allow_cycles, + lambda: f"Detected a cycle while swapping; the cycle includes {visited}", + ) + break + visited.add(vfa) fa = swap_map[vfa] diff --git a/thunder/core/trace_interpreter.py b/thunder/core/trace_interpreter.py index 92901c0229..febc83d476 100644 --- a/thunder/core/trace_interpreter.py +++ b/thunder/core/trace_interpreter.py @@ -262,11 +262,12 @@ class TraceSubstitutionProcessor: NULL = object() - def __init__(self, trace, *args, **kwargs): + def __init__(self, trace, allow_swap_map_cycles=False, *args, **kwargs): self.env = {} self.trace = trace self.new_trace = from_trace(self.trace) self.have_processed_args = False + self.allow_swap_map_cycles = allow_swap_map_cycles def read(self, x: VariableInterface | Any) -> Any: if isinstance(x, VariableInterface): @@ -398,9 +399,9 @@ def __call__(self): for new_bsym in self.new_bsyms: # TODO: what to do with bsym header? Maybe have a combined from_bsym_swap_proxies and from_bsym? self.new_trace.bound_symbols.append( - new_bsym.from_bsym_swap_proxies(self.swap_map).from_bsym( - source_filename=bsym.source_filename, source_positions=bsym.source_positions - ) + new_bsym.from_bsym_swap_proxies( + self.swap_map, allow_cycles=self.allow_swap_map_cycles + ).from_bsym(source_filename=bsym.source_filename, source_positions=bsym.source_positions) ) result = tree_map(self.do_swap, self.replacement_result) diff --git a/thunder/executors/passes.py b/thunder/executors/passes.py index 96274a588f..8e737db576 100644 --- a/thunder/executors/passes.py +++ b/thunder/executors/passes.py @@ -93,7 +93,7 @@ def process_bsym(self, bsym: BoundSymbol) -> None: start_time_ns = time.perf_counter_ns() - extrace, _ = OpExProcessor(trace)() + extrace, _ = OpExProcessor(trace, allow_swap_map_cycles=True)() end_time_ns = time.perf_counter_ns() elapsed_time_ns = end_time_ns - start_time_ns From ae5806e49087a47af15307b7bbbced374b60c4e9 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 6 Dec 2025 07:58:08 -0800 Subject: [PATCH 35/46] Adjust test --- thunder/tests/test_core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index eb22bffd99..89fa15e8c9 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -306,11 +306,11 @@ def foo(a): b = make_tensor((3, 3), device=device, dtype=tdtype, requires_grad=True) cfoo(b) - assert thunder.cache_misses(cfoo) == 2 + assert thunder.cache_misses(cfoo) == 1 b.grad = make_tensor((3, 3), device=device, dtype=tdtype) cfoo(b) - assert thunder.cache_misses(cfoo) == 2 + assert thunder.cache_misses(cfoo) == 1 @instantiate(dtypes=(thunder.float32,)) From 4b8312f1941ddcbfc4a27f4cd56e6094066ea9ae Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 6 Dec 2025 08:01:56 -0800 Subject: [PATCH 36/46] Adjust test --- thunder/tests/test_core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index 89fa15e8c9..25a7028504 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -3103,8 +3103,8 @@ def fn(): # This should be a cache miss, verify that. d[h] = 2 assert jfn() == 2 # Verify that jfn now returns 2 - assert thunder.cache_hits(jfn) == 1 - assert thunder.cache_misses(jfn) == 2 + assert thunder.cache_hits(jfn) == 2 + assert thunder.cache_misses(jfn) == 1 def test_profiling_decorator(): @@ -3253,7 +3253,7 @@ def fn(x): isinstance(out, thunder.core.proxies.TensorProxy) for out in bsym.flat_outs ): # prims is unpack_sequence and any output is TensorProxy # Verify that we print information about the unpacked TensorProxy. - assert "cpu f32[3]" in str(bsym) + assert "cpu f32[[IntegerProxy name=i0, value=3, static=CONSTRAINT.CONSTRAINABLE]]" in str(bsym) @pytest.mark.parametrize("thunderfx_disable_split_autograd", (True, False)) From a78049d8f829a7fe617c859774a54c9984b9fbcc Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 6 Dec 2025 08:03:38 -0800 Subject: [PATCH 37/46] Adjust test --- thunder/tests/test_core.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index 25a7028504..c9d4bdaf00 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -2776,7 +2776,10 @@ def fn(x): assert tr.bound_symbols[1].sym == ltorch.mul (pystr,) = tr.bound_symbols[1].python(0) - assert pystr == 'result = ltorch.mul(2, x) # result: "cpu f32[2, 2]"' + assert ( + pystr + == 'result = ltorch.mul(2, x) # result: "cpu f32[[IntegerProxy name=i0, value=2, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i1, value=2, static=CONSTRAINT.CONSTRAINABLE]]"' + ) def test_dtype_in_trace(): From b1013ba050e92ae6679509cb32f45c47d2826603 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 6 Dec 2025 08:03:55 -0800 Subject: [PATCH 38/46] Adjust test --- thunder/tests/test_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index c9d4bdaf00..7e787dc2b4 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -331,7 +331,7 @@ def foo(a): b = make_tensor((3, 3), device=device, dtype=tdtype, requires_grad=True) b.grad = make_tensor((3, 3), device=device, dtype=tdtype) cfoo(b) - assert thunder.cache_misses(cfoo) == 2 + assert thunder.cache_misses(cfoo) == 1 @instantiate(dtypes=(thunder.float32,)) From dab3ad2c530f777a34a74876ad646a4e583ce2e7 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 6 Dec 2025 08:06:42 -0800 Subject: [PATCH 39/46] Remove obsolete xfail mark --- thunder/tests/test_core.py | 1 - 1 file changed, 1 deletion(-) diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py index 7e787dc2b4..ca31e25313 100644 --- a/thunder/tests/test_core.py +++ b/thunder/tests/test_core.py @@ -2276,7 +2276,6 @@ def foo(a): assert b.shape == torch.Size(shp) -@pytest.mark.xfail(reason="we improperly use an alias") def test_clone_alias(): def foo(a): b = a.clone() From 18d34b38245788d2e1b4357017fd2d51c7b25620 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 6 Dec 2025 08:22:31 -0800 Subject: [PATCH 40/46] Do not treat str as symbolic values --- thunder/core/jit_ext.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 0eae9e1b9c..1acfe2ea0b 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -312,7 +312,9 @@ def proxify(self, value: WrappedValue) -> Any: else: self.add_constraint((clang.check_number_type_and_value, p, uvalue)) elif co is CACHE_OPTIONS.SYMBOLIC_VALUES: - if p is not uvalue: + if isinstance(uvalue, str): + self.add_constraint((clang.check_string_value, p, uvalue)) + elif p is not uvalue: self.add_constraint((clang.check_instance, p, (type(uvalue),))) value.register_proxy(p) elif co not in (CACHE_OPTIONS.SAME_INPUT, CACHE_OPTIONS.NO_CACHING): From 634de486bfa921485292faf25a853456a6308ea9 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Sat, 6 Dec 2025 08:22:40 -0800 Subject: [PATCH 41/46] Adjust test --- thunder/tests/test_grad.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py index 7009393cb6..4a6f515a0d 100644 --- a/thunder/tests/test_grad.py +++ b/thunder/tests/test_grad.py @@ -1855,7 +1855,7 @@ def foo(a, c): c = 2.0 dynamic_jit = thunder.jit(foo, cache="symbolic values") - static_jit = thunder.jit(foo) + static_jit = thunder.jit(foo, cache="constant values") out = dynamic_jit(a, c) torch.autograd.backward(out, torch.rand_like(out), retain_graph=True) From c668c1cd5fb88fdf5f9abb46560a1a3540f63dd0 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 9 Dec 2025 10:20:14 -0800 Subject: [PATCH 42/46] Skip test --- thunder/tests/test_grad.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py index 4a6f515a0d..ae0a1962e0 100644 --- a/thunder/tests/test_grad.py +++ b/thunder/tests/test_grad.py @@ -1571,6 +1571,8 @@ def test_phantom_grad_vs_torch_consistency(op, device: str, dtype: dtypes.dtype, pytest.skip("Skipping complex operator tests in CI for speed") if torch.device(device).type == "cuda" and dtype is dtypes.bfloat16 and not torch.cuda.is_bf16_supported(): pytest.skip("Your CUDA device does not support bfloat16") + if op == get_opinfo("getitem"): + pytest.xfail("TODO: Support slice input with symbolic values") for sample in op.sample_inputs(device, dtype, requires_grad=True): comp = sample.comp if sample.comp is not None else comp From 829074c621d2b0b0d6f2682de0df8a05d9892576 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 9 Dec 2025 11:17:30 -0800 Subject: [PATCH 43/46] Avoid len(a.shape) in executor checker Fixes thunder/tests/test_grad.py::test_phantom_grad_vs_torch_consistency_var_mean_nvfuser_cuda_thunder.dtypes.float32 --- thunder/executors/nvfuserex_impl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index a0677087f2..dd8d1b671e 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -2211,7 +2211,7 @@ def _var_mean_check( if not is_supported_tensor(a, allow_low_precision_floats=False): return False - if len(a.shape) == 0: + if a.ndim == 0: return False if dtypes.is_complex_dtype(dtypes.to_dtype(a)): From d74784960e7dea6c1be0bb81167b04fae65bb47b Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Thu, 11 Dec 2025 06:29:31 -0800 Subject: [PATCH 44/46] Do not treat bool inputs symbolically --- thunder/core/jit_ext.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 1acfe2ea0b..e7da7c7cd6 100644 --- a/thunder/core/jit_ext.py +++ b/thunder/core/jit_ext.py @@ -304,7 +304,7 @@ def proxify(self, value: WrappedValue) -> Any: assert p.history is not None, f"{p.history}, {value.provenance} {type(p)}" co: CACHE_OPTIONS = get_cache_option() - if co is CACHE_OPTIONS.CONSTANT_VALUES: + if co is CACHE_OPTIONS.CONSTANT_VALUES or isinstance(uvalue, bool): if isinstance(uvalue, str): self.add_constraint((clang.check_string_value, p, uvalue)) elif isinstance(uvalue, slice): From fa642acab8ba82cd5f18bd4785d4ef43433f1279 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Thu, 11 Dec 2025 07:09:02 -0800 Subject: [PATCH 45/46] Adjust test --- thunder/tests/test_jit_general.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index fe49e57e09..9b443d2303 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -172,7 +172,7 @@ def test_cache_basic(): def foo(a, b): return a + b - jfoo = thunder_jit(foo) + jfoo = thunder_jit(foo, cache="constant values") a = torch.randn((2, 2), device="cpu") b = torch.randn((2, 2), device="cpu") From 172320d05e2d0fe3e2b2129aaacea1d83be7cb5e Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Thu, 11 Dec 2025 08:16:08 -0800 Subject: [PATCH 46/46] Update test --- thunder/tests/test_jit_general.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py index 9b443d2303..e01a4acca8 100644 --- a/thunder/tests/test_jit_general.py +++ b/thunder/tests/test_jit_general.py @@ -1425,7 +1425,8 @@ def fn(): assert fn() == thunder_jit(fn)() -def test_failing_prologue_in_last_prologue_traces(): +@pytest.mark.parametrize("cache_option", (CACHE_OPTIONS.CONSTANT_VALUES, CACHE_OPTIONS.SYMBOLIC_VALUES)) +def test_failing_prologue_in_last_prologue_traces(cache_option): # we know that this will fail in the prologue i = 0 @@ -1434,8 +1435,12 @@ def fn(): i += 1 return i - jfn = thunder_jit(fn) - with pytest.raises(RuntimeError, match="Expected 1 to be equal to and have the type of 0"): + jfn = thunder_jit(fn, cache=cache_option) + if cache_option == CACHE_OPTIONS.CONSTANT_VALUES: + message = "Expected 1 to be equal to and have the type of 0" + else: + message = "Expected \\[IntegerProxy name=i1, value=1, static=CONSTRAINT.CONSTRAINABLE\\] to be an instance of one of \\(,\\)" + with pytest.raises(RuntimeError, match=message): jfn() # make sure that we have prologue traces in the last_prologue_traces