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 { diff --git a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py index 87faa299cb7..21a0283b295 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_arithmetic.py @@ -13,12 +13,45 @@ # 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): @@ -26,21 +59,53 @@ 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 + + # 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, + ) diff --git a/dali/python/nvidia/dali/experimental/dynamic/_batch.py b/dali/python/nvidia/dali/experimental/dynamic/_batch.py index d9a910418a9..f82392ed55a 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_batch.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_batch.py @@ -25,7 +25,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, _array_from_python, _is_full_slice @@ -153,6 +153,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/_op_builder.py b/dali/python/nvidia/dali/experimental/dynamic/_op_builder.py index 3bdb1722d89..517ff039d06 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 @@ -524,29 +523,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 bc66840f69b..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 @@ -185,35 +186,47 @@ 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 - - call_cache: dict[tuple[int, int], cst.Call | 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.""" - 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: - code = frame.f_code - if sys.version_info >= (3, 11): - # One co_positions tuple per 2-byte code unit. + 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 + + 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.calls_by_position.get((sl, el, sc, ec)) - candidates = self.calls_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: @@ -285,7 +298,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 +306,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: @@ -305,8 +318,9 @@ 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: - self.calls.append(node) + # 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) def _get_module_info(code: types.CodeType) -> ModuleInfo | None: @@ -365,22 +379,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 +405,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: cst.CSTNode | None 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,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.call_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: @@ -464,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.call_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) } @@ -610,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.call_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 @@ -717,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 diff --git a/dali/python/nvidia/dali/experimental/dynamic/_tensor.py b/dali/python/nvidia/dali/experimental/dynamic/_tensor.py index 7b1b3d21729..dd79efbf2d1 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 @@ -124,6 +124,7 @@ def _array_from_python(data, dtype=None): return arr, converted_dtype_id +@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)