From a677329b04480a70e0d1d83563822ef96f05ffa7 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Mon, 7 Sep 2026 18:48:55 +0000 Subject: [PATCH 1/7] Don't make constants tensor inputs in arithmetic ops Signed-off-by: Rostan Tabet --- dali/operators/math/expressions/arithmetic.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dali/operators/math/expressions/arithmetic.cc b/dali/operators/math/expressions/arithmetic.cc index 167c39634dc..5718b089b19 100644 --- a/dali/operators/math/expressions/arithmetic.cc +++ b/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>("integer_constants", "", nullptr, true) + .AddOptionalArg>("integer_constants", "", nullptr) .NumInput(1, 64) // Some arbitrary number that needs to be validated in operator - .AddOptionalArg>("real_constants", "", nullptr, true) + .AddOptionalArg>("real_constants", "", nullptr) .NumOutput(1) .MakeDocHidden() .OutputNDim(0, [](const OpSpec &spec)->std::optional { From be5eb086578d54de7562f008343f04a69c95a654 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Mon, 7 Sep 2026 19:52:53 +0000 Subject: [PATCH 2/7] Treat scalars as constants in ndd arithmetic ops Signed-off-by: Rostan Tabet --- .../dali/experimental/dynamic/_arithmetic.py | 48 ++++++++++++------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py index 87faa299cb7..10d7c7b7e24 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py @@ -14,11 +14,6 @@ import numbers -from typing import Any - - -def _implicitly_convertible(value: Any): - return isinstance(value, (numbers.Real, list, tuple)) def _arithm_op(name: str, *args): @@ -26,21 +21,42 @@ def _arithm_op(name: str, *args): 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 - return _arithmetic_generic_op(*new_args, expression_desc=f"{name}({argsstr})") + # 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:]) + + desc, inputs, integers, reals = [], [], [], [] + for arg in args: + type_ = type(arg) + 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( + *inputs, + expression_desc=f"{name}({' '.join(desc)})", + integer_constants=integers or None, + real_constants=reals or None, + ) From 494d10e484012dc2898bc21a5d34fa01aba56519 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Fri, 11 Sep 2026 09:13:57 +0000 Subject: [PATCH 3/7] Make sure constant integers are fit in in32 in arithmetic ops Signed-off-by: Rostan Tabet --- dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py index 10d7c7b7e24..f6177c025a2 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py @@ -45,6 +45,8 @@ def to_input(arg): desc.append(f"${len(integers)}:bool") integers.append(int(arg)) elif type_ is int: + if (arg >> 31) not in (0, -1): + raise OverflowError(f"Integer constant {arg} is out of range for int32.") desc.append(f"${len(integers)}:int32") integers.append(arg) elif type_ is float: From 42cee6d35d4d93ca4d49c800623bc3f382e11728 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Tue, 8 Sep 2026 14:03:39 +0000 Subject: [PATCH 4/7] Rename call-site terminology to node-site Signed-off-by: Rostan Tabet --- .../experimental/dynamic/_source_analysis.py | 81 +++++++++---------- 1 file changed, 40 insertions(+), 41 deletions(-) diff --git a/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py b/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py index bc66840f69b..b889a0e79d0 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py @@ -185,23 +185,23 @@ class ModuleInfo: scope_of_node: Mapping[cst.CSTNode, Scope] # LibCST ScopeProvider parent_of: Mapping[cst.CSTNode, cst.CSTNode] # built by the fused codegen pass - calls_by_position: Mapping[tuple[int, int, int, int], cst.Call | None] # None if ambiguous - calls_by_line: Mapping[int, tuple[cst.Call, ...]] # 3.10 fallback for call-site identification + nodes_by_position: Mapping[tuple[int, int, int, int], cst.CSTNode | None] # None if ambiguous + nodes_by_line: Mapping[int, tuple[cst.CSTNode, ...]] # 3.10 fallback for site identification - call_cache: dict[tuple[int, int], cst.Call | None] = field(default_factory=dict, repr=False) + node_cache: dict[tuple[int, int], cst.CSTNode | None] = field(default_factory=dict, repr=False) - def call_at(self, frame: types.FrameType) -> cst.Call | None: - """The ``cst.Call`` executing at `frame`'s current instruction, memoized per call site.""" + def node_at(self, frame: types.FrameType) -> cst.CSTNode | None: + """The node executing at `frame`'s current instruction, memoized per site.""" key = (id(frame.f_code), frame.f_lasti) # Use _Unresolved as sentinel value - None is a legitimate cache entry - if (call := self.call_cache.get(key, _Unresolved)) is not _Unresolved: - return call - call = self._resolve_call(frame) - self.call_cache[key] = call - return call - - @NVTXRange("_resolve_call", category="source analysis") - def _resolve_call(self, frame: types.FrameType) -> cst.Call | None: + if (node := self.node_cache.get(key, _Unresolved)) is not _Unresolved: + return node + node = self._resolve_node(frame) + self.node_cache[key] = node + return node + + @NVTXRange("_resolve_node", category="source analysis") + def _resolve_node(self, frame: types.FrameType) -> cst.CSTNode | None: code = frame.f_code if sys.version_info >= (3, 11): # One co_positions tuple per 2-byte code unit. @@ -212,8 +212,8 @@ def _resolve_call(self, frame: types.FrameType) -> cst.Call | None: sc, ec = _byte_to_char_col(lines, sl, sc), _byte_to_char_col(lines, el, ec) if sc is None or ec is None: return None - return self.calls_by_position.get((sl, el, sc, ec)) - candidates = self.calls_by_line.get(frame.f_lineno, ()) + return self.nodes_by_position.get((sl, el, sc, ec)) + candidates = self.nodes_by_line.get(frame.f_lineno, ()) return candidates[0] if len(candidates) == 1 else None def binding(self, name_node: cst.Name) -> Binding | None: @@ -285,7 +285,7 @@ class _FusedCodegenState(PositionProvidingCodegenState): """A single codegen pass that yields everything classification reads off the tree. ``PositionProvider`` already renders every node to compute syntactic positions, so we - piggyback the parent map and call collection onto that same traversal. This replaces + piggyback the parent map and site collection onto that same traversal. This replaces three separate full-tree passes (``PositionProvider``, ``ParentNodeProvider`` and ``matchers.findall``) with one; only ``ScopeProvider`` still needs its own pass. """ @@ -293,7 +293,7 @@ class _FusedCodegenState(PositionProvidingCodegenState): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.parent_of: dict[cst.CSTNode, cst.CSTNode] = {} - self.calls: list[cst.Call] = [] + self.nodes: list[cst.CSTNode] = [] self._node_stack: list[cst.CSTNode] = [] def before_codegen(self, node: cst.CSTNode) -> None: @@ -306,7 +306,7 @@ def after_codegen(self, node: cst.CSTNode) -> None: super().after_codegen(node) self._node_stack.pop() if type(node) is cst.Call: - self.calls.append(node) + self.nodes.append(node) def _get_module_info(code: types.CodeType) -> ModuleInfo | None: @@ -365,22 +365,22 @@ def _get_module_info(code: types.CodeType) -> ModuleInfo | None: ) wrapper.module._codegen(state) - by_position: dict[tuple[int, int, int, int], cst.Call | None] = {} - by_line: dict[int, list[cst.Call]] = {} - with NVTXRange("_get_module_info: Visit calls", category="source analysis"): - for call in state.calls: - r = positions._computed[call] + by_position: dict[tuple[int, int, int, int], cst.CSTNode | None] = {} + by_line: dict[int, list[cst.CSTNode]] = {} + with NVTXRange("_get_module_info: Visit sites", category="source analysis"): + for site in state.nodes: + r = positions._computed[site] span = (r.start.line, r.end.line, r.start.column, r.end.column) by_position[span] = ( - None if span in by_position else call + None if span in by_position else site ) # seen twice -> ambiguous - by_line.setdefault(r.start.line, []).append(call) + by_line.setdefault(r.start.line, []).append(site) info = ModuleInfo( scope_of_node=cast(Mapping[cst.CSTNode, Scope], md), parent_of=state.parent_of, - calls_by_position=by_position, - calls_by_line={ln: tuple(calls) for ln, calls in by_line.items()}, + nodes_by_position=by_position, + nodes_by_line={ln: tuple(nodes) for ln, nodes in by_line.items()}, ) except Exception: info = None @@ -391,31 +391,30 @@ def _get_module_info(code: types.CodeType) -> ModuleInfo | None: return info -_call_cache = {} +_site_cache = {} @dataclass(frozen=True, slots=True) -class CallInfo: - call: Any +class SiteInfo: + node: Any module_info: ModuleInfo meta: dict = field(default_factory=dict) -@NVTXRange("call_info", category="source analysis") -def call_info(frame: types.FrameType) -> CallInfo | None: +@NVTXRange("site_info", category="source analysis") +def site_info(frame: types.FrameType) -> SiteInfo | None: key = (id(frame.f_code), frame.f_lasti) - if (entry := _call_cache.get(key)) is not None: + if (entry := _site_cache.get(key)) is not None: if entry[0]() is frame.f_code: return entry[1] mi = _get_module_info(frame.f_code) if mi is not None: - call = mi.call_at(frame) - call_info = CallInfo(call, mi) + info = SiteInfo(mi.node_at(frame), mi) else: - call_info = None - _call_cache[key] = (weakref.ref(frame.f_code), call_info) - return call_info + info = None + _site_cache[key] = (weakref.ref(frame.f_code), info) + return info @dataclass(slots=True) @@ -441,7 +440,7 @@ def _merge_required_depth(self, child: "_Classifier") -> None: def classify( self, inputs: tuple[Any, ...], raw_kwargs: dict[str, Any] ) -> tuple[list[CaptureRef | Any], dict[str, CaptureRef | Any]] | None: - call = self.module_info.call_at(self.frame) if self.module_info is not None else None + call = self.module_info.node_at(self.frame) if self.module_info is not None else None source_args = _split_call_args(call) if call is not None else None pos_nodes, kw_nodes = source_args or ((), {}) @@ -466,7 +465,7 @@ def classify( def detect_invariant_args( self, inputs: tuple[Any, ...], raw_kwargs: dict[str, Any] ) -> tuple[list[bool], dict[str, bool]] | None: - call = self.module_info.call_at(self.frame) + call = self.module_info.node_at(self.frame) if call is None: return None @@ -614,7 +613,7 @@ def _is_param_invariant(self, name_node: cst.Name, owner_frame: types.FrameType) if mi is None: return False - call = mi.call_at(caller) + call = mi.node_at(caller) if call is None: return False From 92adfd0fe3729b517f8b4902c93e0f05aa763fc2 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Tue, 8 Sep 2026 14:37:57 +0000 Subject: [PATCH 5/7] Mark arithmetic wrappers as transparent Signed-off-by: Rostan Tabet --- .../dali/experimental/dynamic/_arithmetic.py | 38 +++++++++++++++++++ .../dali/experimental/dynamic/_batch.py | 3 +- .../dali/experimental/dynamic/_tensor.py | 3 +- .../nvidia/dali/experimental/dynamic/math.py | 6 +++ 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py index f6177c025a2..8675956cf81 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py @@ -13,8 +13,46 @@ # limitations under the License. +import functools import numbers +from ._call_site import mark_transparent, resolve_callsite_frame + + +@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 diff --git a/dali/python/nvidia/dali/experimental/dynamic/_batch.py b/dali/python/nvidia/dali/experimental/dynamic/_batch.py index 693838390f6..54c288f8009 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_batch.py +++ b/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. diff --git a/dali/python/nvidia/dali/experimental/dynamic/_tensor.py b/dali/python/nvidia/dali/experimental/dynamic/_tensor.py index 53dc1a46444..cfd3e7fc7bb 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_tensor.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_tensor.py @@ -23,7 +23,7 @@ from nvidia.dali._typing import TensorLike from . import _call_site, _eval_mode, _invocation, _stream -from ._arithmetic import _arithm_op +from ._arithmetic import _arithm_op, transparent_arithmetic from ._device import Device, DeviceLike from ._device import device as _device from ._eval_context import EvalContext as _EvalContext @@ -78,6 +78,7 @@ def _try_convert_enums(arr): raise TypeError(f"Unexpected element type {type(item)}") +@transparent_arithmetic class Tensor: """A Tensor object. diff --git a/dali/python/nvidia/dali/experimental/dynamic/math.py b/dali/python/nvidia/dali/experimental/dynamic/math.py index 172e5726745..49f6ad62d28 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/math.py +++ b/dali/python/nvidia/dali/experimental/dynamic/math.py @@ -13,6 +13,7 @@ # limitations under the License. from ._batch import _arithm_op +from ._call_site import mark_transparent as _mark_transparent def sqrt(input): @@ -270,3 +271,8 @@ def clamp(value, lo, hi): :rtype: Tensor or Batch of the type that is calculated based on the type promotion rules. """ return _arithm_op("clamp", value, lo, hi) + + +for _name, _fn in list(globals().items()): + if not _name.startswith("_"): + _mark_transparent(_fn) From 80d7654b1ab75a4a128bd67ad4b7c762ba52dc31 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Wed, 9 Sep 2026 11:01:36 +0000 Subject: [PATCH 6/7] Detect arithmetic ops in source analysis Signed-off-by: Rostan Tabet --- .../dali/experimental/dynamic/_op_builder.py | 28 +--- .../experimental/dynamic/_source_analysis.py | 146 ++++++++++++------ 2 files changed, 98 insertions(+), 76 deletions(-) diff --git a/dali/python/nvidia/dali/experimental/dynamic/_op_builder.py b/dali/python/nvidia/dali/experimental/dynamic/_op_builder.py index d68513f0281..b5605e5fd45 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_op_builder.py +++ b/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 = {} diff --git a/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py b/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py index b889a0e79d0..226231fe1ee 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py @@ -17,6 +17,7 @@ import inspect import itertools import linecache +import opcode import sys import types import weakref @@ -188,32 +189,44 @@ class ModuleInfo: nodes_by_position: Mapping[tuple[int, int, int, int], cst.CSTNode | None] # None if ambiguous nodes_by_line: Mapping[int, tuple[cst.CSTNode, ...]] # 3.10 fallback for site identification - node_cache: dict[tuple[int, int], cst.CSTNode | None] = field(default_factory=dict, repr=False) - - def node_at(self, frame: types.FrameType) -> cst.CSTNode | None: - """The node executing at `frame`'s current instruction, memoized per site.""" - key = (id(frame.f_code), frame.f_lasti) - # Use _Unresolved as sentinel value - None is a legitimate cache entry - if (node := self.node_cache.get(key, _Unresolved)) is not _Unresolved: - return node - node = self._resolve_node(frame) - self.node_cache[key] = node - return node - - @NVTXRange("_resolve_node", category="source analysis") - def _resolve_node(self, frame: types.FrameType) -> cst.CSTNode | None: - code = frame.f_code - if sys.version_info >= (3, 11): - # One co_positions tuple per 2-byte code unit. + if sys.version_info >= (3, 11): + + @NVTXRange("node_at", category="source analysis") + def node_at(self, frame: types.FrameType) -> cst.CSTNode | None: + """The node executing at `frame`'s current instruction.""" + code = frame.f_code + # One co_positions tuple per 2-byte code unit pos = next(itertools.islice(code.co_positions(), frame.f_lasti // 2, None), None) - if pos is not None and all(x is not None for x in pos): - sl, el, sc, ec = cast(tuple[int, int, int, int], pos) - lines = linecache.getlines(code.co_filename) - sc, ec = _byte_to_char_col(lines, sl, sc), _byte_to_char_col(lines, el, ec) - if sc is None or ec is None: - return None - return self.nodes_by_position.get((sl, el, sc, ec)) - candidates = self.nodes_by_line.get(frame.f_lineno, ()) + if pos is None or not all(x is not None for x in pos): + return self._node_on_line(frame.f_lineno, None) + sl, el, sc, ec = cast(tuple[int, int, int, int], pos) + lines = linecache.getlines(code.co_filename) + sc, ec = _byte_to_char_col(lines, sl, sc), _byte_to_char_col(lines, el, ec) + if sc is None or ec is None: + return None + return self.nodes_by_position.get((sl, el, sc, ec)) + + else: + + @NVTXRange("node_at", category="source analysis") + def node_at(self, frame: types.FrameType) -> cst.CSTNode | None: + """The node executing at `frame`'s current instruction.""" + name = opcode.opname[frame.f_code.co_code[frame.f_lasti]] + if name.startswith("CALL_"): + accepted = cst.Call + elif name.startswith("INPLACE_"): + accepted = cst.AugAssign + elif name.startswith("BINARY_") and name != "BINARY_SUBSCR": + accepted = cst.BinaryOperation + else: + return None + return self._node_on_line(frame.f_lineno, accepted) + + def _node_on_line(self, lineno: int, accepted: type[cst.CSTNode] | None) -> cst.CSTNode | None: + """The single node on `lineno`, restricted to the `accepted` type if given.""" + candidates = self.nodes_by_line.get(lineno, ()) + if accepted is not None: + candidates = tuple(n for n in candidates if isinstance(n, accepted)) return candidates[0] if len(candidates) == 1 else None def binding(self, name_node: cst.Name) -> Binding | None: @@ -305,7 +318,8 @@ def before_codegen(self, node: cst.CSTNode) -> None: def after_codegen(self, node: cst.CSTNode) -> None: super().after_codegen(node) self._node_stack.pop() - if type(node) is cst.Call: + # A frame executes either at a call or at an arithmetic operator reached via a dunder + if isinstance(node, (cst.Call, cst.BinaryOperation, cst.AugAssign)): self.nodes.append(node) @@ -396,7 +410,7 @@ def _get_module_info(code: types.CodeType) -> ModuleInfo | None: @dataclass(frozen=True, slots=True) class SiteInfo: - node: Any + node: cst.CSTNode | None module_info: ModuleInfo meta: dict = field(default_factory=dict) @@ -440,8 +454,9 @@ def _merge_required_depth(self, child: "_Classifier") -> None: def classify( self, inputs: tuple[Any, ...], raw_kwargs: dict[str, Any] ) -> tuple[list[CaptureRef | Any], dict[str, CaptureRef | Any]] | None: - call = self.module_info.node_at(self.frame) if self.module_info is not None else None - source_args = _split_call_args(call) if call is not None else None + info = site_info(self.frame) + call = info.node if info is not None else None + source_args = _split_call_args(call) if isinstance(call, cst.Call) else None pos_nodes, kw_nodes = source_args or ((), {}) try: @@ -463,26 +478,29 @@ def classify( @NVTXRange("detect_invariant_args", category="source analysis") def detect_invariant_args( - self, inputs: tuple[Any, ...], raw_kwargs: dict[str, Any] - ) -> tuple[list[bool], dict[str, bool]] | None: - call = self.module_info.node_at(self.frame) - - if call is None: - return None - - split = _split_call_args(call) - if split is None: - return None - pos_nodes, kw_nodes = split + self, + site: cst.CSTNode | None, + inputs: Sequence[Any], + raw_kwargs: Mapping[str, Any], + ) -> tuple[list[bool], set[str]] | None: + match site: + case cst.Call() as call: + split = _split_call_args(call) + if split is None: + return None + pos_nodes, kw_nodes = split + case cst.BinaryOperation(left=l, right=r) | cst.AugAssign(target=l, value=r): + pos_nodes, kw_nodes = (l, r), {} + case _: + return None - classified_inputs: list[CaptureRef | Any] = [] + classified_inputs: list[bool] = [] for i in range(min(len(inputs), len(pos_nodes))): node = pos_nodes[i] classified_inputs.append(self.is_invariant(node, static=True)) - # Go over defaults - they are invariant, because they have to be None + # A missing positional node is a defaulted argument, invariant only if it is None for i in range(len(pos_nodes), len(inputs)): - assert inputs[i] is None # - classified_inputs.append(True) + classified_inputs.append(inputs[i] is None) classified_kwargs = { name for name in raw_kwargs if self.is_invariant(kw_nodes.get(name), static=True) } @@ -609,15 +627,12 @@ def _is_param_invariant(self, name_node: cst.Name, owner_frame: types.FrameType) if caller is None: return False - mi = _get_module_info(caller.f_code) # caller may be in another module - if mi is None: - return False - - call = mi.node_at(caller) - if call is None: + info = site_info(caller) # caller may be in another module + if info is None or not isinstance(info.node, cst.Call): return False - child = _Classifier(mi, caller) + call = info.node + child = _Classifier(info.module_info, caller) result = child._is_arg_invariant(call, name_node.value, owner_frame.f_code) self._merge_required_depth(child) return result @@ -716,3 +731,32 @@ def classify( if classification is None: return None return (*classification, classifier.required_depth) + + +def _classify_site( + frame: types.FrameType | None, + key: str, + inputs: Sequence[Any], + raw_kwargs: Mapping[str, Any], +) -> tuple[list[bool], set[str]] | None: + """Memoized `detect_invariant_args` at `frame`'s site, or None if the site is unresolved.""" + if frame is None or not (info := site_info(frame)): + return None + if key not in info.meta: + classifier = _Classifier(info.module_info, frame) + info.meta[key] = classifier.detect_invariant_args(info.node, inputs, raw_kwargs) + return info.meta[key] + + +def constant_inputs(frame: types.FrameType | None, args: Sequence[Any]) -> Sequence[bool]: + """Per-argument flag telling whether the operand is provably constant at `frame`'s site.""" + classified = _classify_site(frame, "constant_inputs", args, {}) + return classified[0] if classified is not None else (False,) * len(args) + + +def constant_kwargs( + frame: types.FrameType | None, raw_kwargs: Mapping[str, Any] +) -> set[str] | None: + """Names of the keyword arguments provably constant at `frame`'s site.""" + classified = _classify_site(frame, "constant_kwargs", (), raw_kwargs) + return classified[1] if classified is not None else None From 8659680dbbab3b243f9ea14aa8cbb12562e4c1b8 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Wed, 9 Sep 2026 11:02:13 +0000 Subject: [PATCH 7/7] Only promote constants arithmetic ops Signed-off-by: Rostan Tabet --- .../dali/experimental/dynamic/_arithmetic.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py index 8675956cf81..21a0283b295 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py @@ -76,15 +76,24 @@ def to_input(arg): 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 in args: + 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: - if (arg >> 31) not in (0, -1): - raise OverflowError(f"Integer constant {arg} is out of range for int32.") desc.append(f"${len(integers)}:int32") integers.append(arg) elif type_ is float: