Extend static analysis for arithmetic operators

dali/operators/math/expressions/arithmetic.cc

@@ -142,9 +142,9 @@ Examples::
   add(&0 mul(&1 $0:int8))
   add(&0 rand()))code",
             DALIDataType::DALI_STRING, false)
-    .AddOptionalArg<std::vector<int32_t>>("integer_constants", "", nullptr, true)
+    .AddOptionalArg<std::vector<int32_t>>("integer_constants", "", nullptr)
     .NumInput(1, 64)  // Some arbitrary number that needs to be validated in operator
-    .AddOptionalArg<std::vector<float>>("real_constants", "", nullptr, true)
+    .AddOptionalArg<std::vector<float>>("real_constants", "", nullptr)
     .NumOutput(1)
     .MakeDocHidden()
     .OutputNDim(0, [](const OpSpec &spec)->std::optional<int> {

dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py

@@ -13,34 +13,99 @@
 # limitations under the License.
 
 
+import functools
 import numbers
-from typing import Any
 
+from ._call_site import mark_transparent, resolve_callsite_frame
 
-def _implicitly_convertible(value: Any):
-    return isinstance(value, (numbers.Real, list, tuple))
+
+@functools.cache
+def _arithmetic_dunders():
+    # Comparisons are omitted for now because we'd need to handle chains and reverse orders
+    binary_ops = (
+        "add",
+        "sub",
+        "mul",
+        "truediv",
+        "floordiv",
+        "mod",
+        "pow",
+        "lshift",
+        "rshift",
+        "and",
+        "or",
+        "xor",
+        "matmul",
+        "divmod",
+    )
+    unary_ops = ("neg", "pos", "abs", "invert")
+
+    binary_dunders = (f"__{prefix}{stem}__" for stem in binary_ops for prefix in ("", "r"))
+    unary_dunders = (f"__{stem}__" for stem in unary_ops)
+
+    return (*binary_dunders, *unary_dunders)
+
+
+def transparent_arithmetic(cls: type) -> type:
+    """Annotate a class' arithmetic dunders with ``mark_transparent``"""
+    for dunder in _arithmetic_dunders():
+        if func := vars(cls).get(dunder):
+            mark_transparent(func)
+    return cls
 
 
 def _arithm_op(name: str, *args):
     from . import _arithmetic_generic_op
     from ._batch import Batch
     from ._tensor import Tensor, as_tensor
 
-    # scalar arguments are turned into tensors
-    argsstr = " ".join(f"&{i}" for i in range(len(args)))
-    gpu = any(arg.device.device_type == "gpu" for arg in args if isinstance(arg, (Tensor, Batch)))
+    tensor_args = [arg for arg in args if isinstance(arg, (Tensor, Batch))]
+    gpu = any(arg.device.device_type == "gpu" for arg in tensor_args)
 
-    new_args = []
-    for arg in args:
+    def to_input(arg):
         if not isinstance(arg, (Tensor, Batch)):
-            if gpu and _implicitly_convertible(arg):
-                arg = as_tensor(arg, device="gpu")
-            else:
-                arg = as_tensor(arg)
+            device = "gpu" if gpu and isinstance(arg, (numbers.Real, list, tuple)) else None
+            arg = as_tensor(arg, device=device)
 
         if (arg.device.device_type == "gpu") != gpu:
             raise ValueError("Cannot mix GPU and CPU inputs.")
 
-        new_args.append(arg)
+        return arg
+
+    # only reachable from math functions called with only scalars, e.g. ndd.math.max(2, 3)
+    if not tensor_args and args:
+        args = (to_input(args[0]), *args[1:])
+
+    if any(type(arg) in (bool, int, float) for arg in args):
+        from ._source_analysis import constant_inputs
+
+        constants = constant_inputs(resolve_callsite_frame(depth_hint=3), args)
+    else:
+        constants = (False,) * len(args)
+
+    desc, inputs, integers, reals = [], [], [], []
+    for arg, constant in zip(args, constants, strict=True):
+        type_ = type(arg)
+        if type_ is int and (arg >> 31) not in (0, -1):
+            raise OverflowError(f"Integer {arg} is out of range for int32.")
+
+        type_ = type_ if constant else None
+        if type_ is bool:
+            desc.append(f"${len(integers)}:bool")
+            integers.append(int(arg))
+        elif type_ is int:
+            desc.append(f"${len(integers)}:int32")
+            integers.append(arg)
+        elif type_ is float:
+            desc.append(f"${len(reals)}:float32")
+            reals.append(arg)
+        else:
+            desc.append(f"&{len(inputs)}")
+            inputs.append(to_input(arg))
 
-    return _arithmetic_generic_op(*new_args, expression_desc=f"{name}({argsstr})")
+    return _arithmetic_generic_op(
+        *inputs,
+        expression_desc=f"{name}({' '.join(desc)})",
+        integer_constants=integers or None,
+        real_constants=reals or None,
+    )

dali/python/nvidia/dali/experimental/dynamic/_batch.py

@@ -26,7 +26,7 @@
     from . import _invocation
 from . import _eval_mode, _stream as _stream_module
 from ._eval_context import EvalContext as _EvalContext
-from ._arithmetic import _arithm_op
+from ._arithmetic import _arithm_op, transparent_arithmetic
 from ._device import Device, DeviceLike
 from ._device import device as _device
 from ._tensor import Tensor, _is_full_slice, _try_convert_enums
@@ -154,6 +154,7 @@ def as_batch(self, copy: bool = False):
         return batch(self) if copy else as_batch(self)  # type: ignore
 
 
+@transparent_arithmetic
 class Batch:
     """A Batch object.
 

dali/python/nvidia/dali/experimental/dynamic/_op_builder.py

@@ -29,8 +29,7 @@
 from ._capture import _capture_intercept
 from ._eval_mode import EvalMode
 from ._nvtx import NVTXRange
-from ._source_analysis import _Classifier
-from ._source_analysis import call_info as _call_info
+from ._source_analysis import constant_kwargs
 from ._tensor import Tensor
 from ._tensor import tensor as to_tensor
 from .capture._invariant import unwrap_invariant, unwrap_invariants
@@ -522,29 +521,8 @@ def fn_call(
         if _caller_frame is None:
             _caller_frame = resolve_callsite_frame(depth_hint=3)
 
-        constant_args = None
-        if _caller_frame is not None:
-            info = _call_info(_caller_frame)
-            if info is not None:
-                arg_classification = None
-                if "constant_args" in info.meta:
-                    constant_args = info.meta["constant_args"]
-                else:
-                    # TODO(michalz): use (inputs, raw_kwargs) when we have a way to utilize
-                    #                constant inputs
-                    arg_classification = _Classifier(
-                        info.module_info, _caller_frame
-                    ).detect_invariant_args([], raw_kwargs)
-                    if arg_classification is not None:
-                        # For future use
-                        # info.meta["constant_inputs"] = arg_classification[0]
-                        info.meta["constant_args"] = arg_classification[1]
-                        constant_args = arg_classification[1]
-                    else:
-                        # For future use
-                        # info.meta["constant_inputs"] = None
-                        info.meta["constant_args"] = None
-                        constant_args = None
+        # TODO(michalz): utilize constant inputs
+        constant_args = constant_kwargs(_caller_frame, raw_kwargs)
 
         init_args = {}
         call_args = {}