Enhance symbolic arithmetics

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,

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
 

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,
     )
 

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,21 +2806,13 @@ 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)
     utils.check_type(pred, (TensorProxy, Number, NumberProxy))
     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,)

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)

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):

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)