diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py index 333086b906..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, ) @@ -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)): @@ -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, ) @@ -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/interpreter.py b/thunder/core/interpreter.py index 235c6689ea..6bc8270940 100644 --- a/thunder/core/interpreter.py +++ b/thunder/core/interpreter.py @@ -39,6 +39,8 @@ TracebackType, ) +import torch + from thunder.core.baseutils import Singleton, init_colors, extract_callable_name, is_likely_from_collections_namedtuple from thunder.core.codeutils import Positions @@ -397,8 +399,11 @@ def __init__( self._callbacks: dict[INTERPRETER_CALLBACKS, Callable] = callbacks 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}) + if uncacheable_classes is None: + uncacheable_classes = () + uncacheable_classes = tuple( + set(uncacheable_classes) | {NoneType, int, str, float, bool, complex, torch.Tensor} + ) self._uncacheable_classes = uncacheable_classes diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py index 6e8f52d45c..e849e5ddbd 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}") @@ -2164,7 +2165,6 @@ def thunder_general_jit( callbacks=general_jit_callbacks, with_provenance_tracking=True, unwrap_result=False, - uncacheable_classes=(torch.Tensor, int, float, str, NoneType), record_history=compile_data.debug_options.record_interpreter_history, ) diff --git a/thunder/core/prims.py b/thunder/core/prims.py index 90e1f3689b..432250f789 100644 --- a/thunder/core/prims.py +++ b/thunder/core/prims.py @@ -1945,9 +1945,10 @@ 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). -# The ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND enum describes these three behaviors so that +# 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 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 +2317,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 +2387,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, ) @@ -2804,7 +2806,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) @@ -2812,13 +2813,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( @@ -2843,11 +2837,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/core/proxies.py b/thunder/core/proxies.py index d9f2bcdd5e..25fc08b775 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 @@ -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 @@ -739,7 +742,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 +778,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 +800,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 +957,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 +1679,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) 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): 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) diff --git a/thunder/executors/pythonex.py b/thunder/executors/pythonex.py index 0083853f59..7c0c051571 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) @@ -373,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 diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py index 7f06d7eab8..96e32cf484 100644 --- a/thunder/tests/opinfos.py +++ b/thunder/tests/opinfos.py @@ -6687,7 +6687,7 @@ def arange_sample_generator(op, device, dtype, requires_grad, **kwargs): ) for case in partial_cases: - yield SampleInput(*case) + yield SampleInput(*case, dtype=dtype, device=device) arange_opinfo = OpInfo( diff --git a/thunder/tests/test_elementwise.py b/thunder/tests/test_elementwise.py index 2286f0e971..a17f87228c 100644 --- a/thunder/tests/test_elementwise.py +++ b/thunder/tests/test_elementwise.py @@ -1,8 +1,10 @@ from functools import partial import builtins +import itertools import math import operator +import pytest import torch from torch.testing import assert_close, make_tensor @@ -13,23 +15,37 @@ 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": + 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,49 @@ 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) + + 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), +} - for a in a_vals: - for b in b_vals: - actual = cfoo(a, b) - expected = foo(a, b) - assert_close(actual, expected) +@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] -@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": + allowed_types = tuple({int, float}.intersection(allowed_types)) bool_inps = [False, True] int_inps = [-1, 0, 2] @@ -112,19 +141,49 @@ 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) + if cache_option == "symbolic values": + 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] + complex_inps = [complex(1, 0.3), complex(-4.1, 0.9)] + + 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, complex_inps]: + for pred, a, b in itertools.product(bool_inps, inps, inps): + actual = cfoo(pred, a, b) + expected = foo_python(pred, a, b) assert_close(actual, expected) + if cache_option == "symbolic values": + assert thunder.cache_misses(cfoo) == 4 + # TODO: see issue "Test operator and method variants of operations using # OpInfos" diff --git a/thunder/tests/test_interpreter.py b/thunder/tests/test_interpreter.py index 6e36798eda..e7ca65a357 100644 --- a/thunder/tests/test_interpreter.py +++ b/thunder/tests/test_interpreter.py @@ -34,9 +34,7 @@ # This wraps the jit call into a tracking one (using a wrapper function # rather than partial to get a nice test name). def interpret_tracking(*args, **kwargs): - return interpret( - *args, with_provenance_tracking=True, uncacheable_classes=(torch.Tensor, int, float, str, type(None)), **kwargs - ) + return interpret(*args, with_provenance_tracking=True, **kwargs) # This will be called by PyTest and parametrize each test that has diff --git a/thunder/tests/test_nvfuser.py b/thunder/tests/test_nvfuser.py index 4d341cfcb5..6ef5437b12 100644 --- a/thunder/tests/test_nvfuser.py +++ b/thunder/tests/test_nvfuser.py @@ -1061,3 +1061,25 @@ def test_scatter(executor, device: str, dtype: dtypes.dtype): assert len(fusion_bsyms) == 1 outside_fusion_syms = ["unpack_trivial", "python_return"] assert {el.sym.name for el in fw_trace.bound_symbols if not el.sym.is_fusion} == set(outside_fusion_syms) + + +@instantiate( + executors=(nvFuserExecutor,), + dtypes=NOTHING, + decorators=(pytest.mark.xfail(reason="nvFuser does not support symbolic values for arange"),), +) +def test_arange_symbolic_values(executor, device: str, dtype: dtypes.dtype): + from thunder.tests.opinfos import arange_opinfo + + for sample in arange_opinfo.sample_inputs(device, thunder.float32): + compiled_func = executor.make_callable(torch.arange, cache="symbolic values") + out = compiled_func(*sample.args, **sample.kwargs) + expected_out = torch.arange(*sample.args, **sample.kwargs) + torch.testing.assert_close(out, expected_out) + + trace = thunder.last_traces(compiled_func)[-1] + computation_bsyms = [ + bsym for bsym in trace.bound_symbols if bsym.sym not in (prims.python_return, prims.unpack_trivial) + ] + assert len(computation_bsyms) == 1 + assert computation_bsyms[0].sym.name == "nvFusion0"