diff --git a/src/ninetoothed/backends/materializers/triton.py b/src/ninetoothed/backends/materializers/triton.py index 2c8f0f8..ff7e9de 100644 --- a/src/ninetoothed/backends/materializers/triton.py +++ b/src/ninetoothed/backends/materializers/triton.py @@ -8,6 +8,7 @@ import subprocess import sys import tempfile +import threading import types from dataclasses import dataclass, replace from pathlib import Path @@ -157,6 +158,18 @@ def _make_arg_key(args, kwargs): specs=compilation.kernel.tensors, prepare_invocation=_triton_prepare_invocation(launch, kernel), validate_bindings=validate_bindings, + structural_observer=_triton_structural_observer( + compilation.launch_abi, + ), + structural_key=( + lambda args, kwargs, public, bound_public: _triton_structural_key( + compilation, + args, + kwargs, + public, + bound_public=bound_public, + ) + ), ) ) @@ -203,6 +216,19 @@ def _candidate_launch(compilation, launch, kernel, candidate, validate_bindings) binding_overrides=meta_values, prepare_invocation=_triton_prepare_invocation(function, kernel), validate_bindings=validate_bindings, + structural_observer=_triton_structural_observer( + compilation.launch_abi, + binding_overrides=meta_values, + ), + structural_key=( + lambda args, kwargs, public, bound_public: _triton_structural_key( + compilation, + args, + kwargs, + public, + bound_public=bound_public, + ) + ), ) @@ -271,6 +297,7 @@ def invoke(self, values, args, kwargs): class InvocationPlan: call: object requires_values: bool + structurally_rebindable: bool = False def __call__(self, values, args, kwargs): if self.requires_values: @@ -279,16 +306,48 @@ def __call__(self, values, args, kwargs): self.call(args, kwargs) def plan_value(value, values, static_values, call_sources): - for index, candidate in enumerate(values): - if value is candidate: - if call_sources is not None and call_sources[index] is not None: - return CallRef(*call_sources[index]) - - return ( - Literal(static_values[index]) - if static_values is not None - else ValueRef(index) - ) + matches = tuple( + index for index, candidate in enumerate(values) if value is candidate + ) + + if matches: + index = matches[0] + + if call_sources is not None: + sources = tuple(call_sources[match] for match in matches) + + # A direct invocation may only rebind an aliased captured value + # when every occurrence came from the same public argument. If + # one object represented two different call arguments during the + # dry launch, choosing the first CallRef would silently collapse + # distinct objects on a later structurally equivalent call. + if sources[0] is not None and all( + source == sources[0] for source in sources + ): + return CallRef(*sources[0]) + + if len(matches) > 1: + if hasattr(value, "data_ptr") or ( + hasattr(value, "shape") and hasattr(value, "dtype") + ): + return None + + # CPython commonly reuses the same immutable scalar object + # for equal shape, stride, and constexpr values. It is + # safe to detach such an ambiguous scalar as a literal: + # the structural cache key contains every public/runtime + # scalar value, so a later value change cannot hit this + # invocation plan. + if _is_cacheable_runtime_literal(value): + return Literal(value) + + return None + + return ( + Literal(static_values[index]) + if static_values is not None + else ValueRef(index) + ) if hasattr(value, "data_ptr") or ( hasattr(value, "shape") and hasattr(value, "dtype") @@ -387,7 +446,7 @@ def capture(*args, **kwargs): planned_call.kwargs, ) - return InvocationPlan(direct_call, False) + return InvocationPlan(direct_call, False, structurally_rebindable=True) return prepare @@ -471,14 +530,22 @@ def _tuned_runtime_launch( _arm_prepared_runtime_launch, _empty_launch, _first_output, + _PendingStructuralPromotion, _public_values, + _remember_structural_runtime_call, _remember_verified_runtime_call, _runtime_call_identity, + _runtime_owner_refs, + _runtime_owner_refs_match, + _touch_structural_runtime_call, ) active_identity = None active = None prepared_calls = {} + structural_calls = {} + pending_promotions = {} + structural_lock = threading.Lock() def evict(identity, token): nonlocal active, active_identity @@ -497,6 +564,11 @@ def activate(identity, entry): active = entry record(entry[1]) + def deactivate(): + nonlocal active, active_identity + active = None + active_identity = None + def record(selected): handle._selected_tuning_candidate = candidates_by_launch[selected] @@ -523,6 +595,146 @@ def remember_best_effort(identity, selection_key, selected, prepared): except TypeError: return None + def remember_structural(key, selection_key, selected, prepared): + if key is None: + return None + + detached = prepared.detached_for_structural_cache() + + if detached is None: + return None + + entry = (selection_key, selected, detached) + + with structural_lock: + _remember_structural_runtime_call( + structural_calls, + pending_promotions, + key, + entry, + ) + return detached + + def remember_pending(identity, key, args, kwargs): + owner_refs = _runtime_owner_refs(args, kwargs) + + with structural_lock: + if key not in structural_calls: + pending_promotions.pop(key, None) + + return + + if owner_refs is None: + pending_promotions.pop(key, None) + + return + + pending_promotions.pop(key, None) + pending_promotions[key] = _PendingStructuralPromotion( + identity=identity, + owner_refs=owner_refs, + ) + + def remember_structural_best_effort( + key, + selection_key, + selected, + prepared, + ): + try: + return remember_structural(key, selection_key, selected, prepared) + except (AttributeError, TypeError): + return None + + def promote_structural(identity, entry, args, kwargs): + selection_key, selected, prepared = entry + rebinder = getattr(selected, "_ninetoothed_rebind_structural", None) + + if rebinder is None: + return None + + token = object() + + def collected(_reference): + evict(identity, token) + + promoted = rebinder(prepared, args, kwargs, collected, token) + + if promoted is None: + return None + + promoted_entry = (selection_key, selected, promoted) + _remember_verified_runtime_call(prepared_calls, identity, promoted_entry) + activate(identity, promoted_entry) + + return promoted_entry + + def candidate_structural_key( + selected, + args, + kwargs, + public, + *, + alias_signature=None, + ): + observer = getattr(selected, "_ninetoothed_structural_observer", None) + + if observer is not None: + observed = observer(args, kwargs) + + if observed is not None: + return observed + + builder = getattr(selected, "_ninetoothed_structural_key", None) + + if builder is None: + return None + + return builder(args, kwargs, public=public) + + def find_structural( + selection_key, + args, + kwargs, + public, + *, + alias_signature=None, + ): + with structural_lock: + structural_snapshot = tuple(structural_calls.items()) + + for key, entry in reversed(structural_snapshot): + if entry[0] != selection_key: + continue + + selected = entry[1] + current_key = candidate_structural_key( + selected, + args, + kwargs, + public, + alias_signature=alias_signature, + ) + + if current_key != key: + continue + + with structural_lock: + current_entry = _touch_structural_runtime_call( + structural_calls, + pending_promotions, + key, + ) + + if current_entry is None or current_entry[0] != selection_key: + continue + + pending = pending_promotions.pop(key, None) + + return key, current_entry, pending + + return None + def launch(*args, **kwargs): active_snapshot = active active_identity_snapshot = active_identity @@ -571,6 +783,52 @@ def launch(*args, **kwargs): key = tuner._make_arg_key(args, kwargs) alias_signature = _runtime_alias_signature(compilation.launch_abi, public) selection_key = (key, alias_signature) + structural = find_structural( + selection_key, + args, + kwargs, + public, + alias_signature=alias_signature, + ) + + if structural is not None: + structural_key, structural_entry, pending = structural + deactivate() + + # Keep the first structural hit detached. Only a second hit + # with the exact same weakly-held tensor objects earns an + # identity-bound rebind and active fast path. + if ( + pending is not None + and pending.identity == identity + and _runtime_owner_refs_match( + pending.owner_refs, + args, + kwargs, + ) + ): + promoted = promote_structural( + identity, + structural_entry, + args, + kwargs, + ) + else: + promoted = None + + if promoted is None: + remember_pending(identity, structural_key, args, kwargs) + + selected_entry = promoted or structural_entry + + record(selected_entry[1]) + + return selected_entry[1]._ninetoothed_invoke_prepared( + selected_entry[2], + args, + kwargs, + ) + selected = next( ( entry[1] @@ -606,6 +864,19 @@ def launch(*args, **kwargs): selected, prepared, ) + structural_cache_key = candidate_structural_key( + selected, + args, + kwargs, + public, + alias_signature=alias_signature, + ) + remember_structural_best_effort( + structural_cache_key, + selection_key, + selected, + prepared, + ) record(selected) @@ -622,6 +893,19 @@ def launch(*args, **kwargs): return result remember_best_effort(identity, selection_key, selected, prepared) + structural_cache_key = candidate_structural_key( + selected, + args, + kwargs, + public, + alias_signature=alias_signature, + ) + remember_structural_best_effort( + structural_cache_key, + selection_key, + selected, + prepared, + ) return result @@ -640,12 +924,368 @@ def launch(*args, **kwargs): ) record(selected) remember_best_effort(identity, selection_key, selected, prepared) + structural_cache_key = candidate_structural_key( + selected, + args, + kwargs, + public, + alias_signature=alias_signature, + ) + remember_structural_best_effort( + structural_cache_key, + selection_key, + selected, + prepared, + ) return result return launch +def _triton_structural_key(compilation, args, kwargs, public, *, bound_public=None): + from ninetoothed.compiler.runtime import _runtime_structural_key + + return _runtime_structural_key( + compilation.launch_abi, + args, + kwargs, + public, + bound_public=bound_public, + alias_signature=_runtime_alias_signature(compilation.launch_abi, public), + ) + + +def _triton_structural_observer(abi, *, binding_overrides=None, tensor_type=None): + """Build a conservative compact key for ordinary Triton launch ABIs. + + Generated Triton launches repeat each public tensor as a pointer, shape, + and stride argument. For an exact tensor type those derived values are + already covered by the public tensor contract, so the repeated binding + walk in the generic structural key can be omitted. Any ABI or value that + is not proved equivalent falls back to the generic key in the runtime. + """ + from ninetoothed.compiler.runtime import _runtime_literal_key + + try: + if tensor_type is None: + import torch + + tensor_type = torch.Tensor + except ImportError: + return None + + public_args = tuple(abi.public_args) + public_indexes = {name: index for index, name in enumerate(public_args)} + public_count = len(public_args) + overrides = dict(binding_overrides or {}) + bindings = tuple(abi.kernel_args) + supported = {"tensor", "scalar", "constexpr", "shape", "stride", "meta"} + scalar_kinds = {"scalar", "constexpr", "meta"} + dynamic_kinds = {"tensor", "shape", "stride"} + + if ( + not bindings + or len(public_indexes) != public_count + or set(abi.outputs).difference(public_indexes) + or any(name in public_indexes for name in overrides) + ): + return None + + for value in overrides.values(): + if getattr(value, "shape", None) is not None and hasattr(value, "dtype"): + return None + + scalar_sources = [False] * public_count + dynamic_sources = [False] * public_count + required_dims = [-1] * public_count + residual = [] + physical_access = {} + output_names = set(abi.outputs) + + for binding in bindings: + kind = binding.kind + source = binding.source + + if kind not in supported or kind.startswith("jagged_"): + return None + + source_index = public_indexes.get(source) + + if kind in dynamic_kinds: + if source_index is None: + return None + + dynamic_sources[source_index] = True + + if kind in {"shape", "stride"}: + if type(binding.dim) is not int or binding.dim < 0: + return None + + required_dims[source_index] = max( + required_dims[source_index], binding.dim + ) + elif kind in scalar_kinds and source_index is not None: + scalar_sources[source_index] = True + elif kind in {"scalar", "constexpr"}: + if source not in overrides: + return None + + value_key = _runtime_literal_key(overrides[source]) + + if value_key is None: + return None + + residual.append((binding.name, kind, source, binding.dim, value_key)) + elif kind == "meta" and source_index is None: + if source in overrides: + value = overrides[source] + elif binding.value is not None: + value = binding.value + else: + return None + + value_key = _runtime_literal_key(value) + + if value_key is None: + return None + + residual.append((binding.name, kind, source, binding.dim, value_key)) + + if kind not in {"tensor", "scalar", "constexpr"} or source_index is None: + continue + + access = binding.access or ( + "read" + if kind in {"scalar", "constexpr"} + else "read_write" + if source in output_names + else "read" + ) + identity = (source_index, kind) + previous = physical_access.get(identity) + + if previous is not None and previous != access: + access = "read_write" + + physical_access[identity] = access + + physical = tuple( + (source_index, kind, access) + for (source_index, kind), access in physical_access.items() + ) + writers = tuple(item for item in physical if item[2] in {"write", "read_write"}) + readers = tuple(item for item in physical if item[2] in {"read", "read_write"}) + alias_pairs = tuple( + ( + writer[0], + reader[0], + ( + public_args[writer[0]], + writer[1], + public_args[reader[0]], + reader[1], + ), + ) + for writer in writers + for reader in readers + ) + missing = object() + + def observe_tensor(value, *, scalar): + if type(value) is not tensor_type: + return None + + try: + shape = tuple(value.shape) + stride = tuple(value.stride()) + data_ptr = int(value.data_ptr()) + storage_offset = int(value.storage_offset()) + element_size = int(value.element_size()) + dtype = value.dtype + device = value.device + + if len(shape) != len(stride) or device is None or element_size <= 0: + return None + + scalar_key = None + + if scalar: + if shape: + return None + + scalar_key = _runtime_literal_key(value.item()) + + if scalar_key is None: + return None + + state = ( + type(value), + len(shape), + shape, + stride, + dtype, + device, + storage_offset, + data_ptr & 15, + ) + key = ("tensor", state) + + if scalar_key is not None: + key = ("tensor", (*state, scalar_key)) + + lower = upper = 0 + + for size, step in zip(shape, stride): + extent = (size - 1) * step + lower += min(0, extent) + upper += max(0, extent) + + span = ( + (device, 0, 0) + if 0 in shape + else ( + device, + data_ptr + lower * element_size, + data_ptr + (upper + 1) * element_size, + ) + ) + except ( + AttributeError, + RuntimeError, + TypeError, + ValueError, + OverflowError, + ): + return None + + return key, shape, stride, span + + def observe( + args, + kwargs, + *, + public=None, + bound_public=None, + alias_signature=None, + ): + del public, bound_public + + try: + argument_count = len(args) + + if argument_count > public_count: + return None + + if argument_count == public_count and not kwargs: + values = args + else: + values = [missing] * public_count + values[:argument_count] = args + + for name, value in kwargs.items(): + index = public_indexes.get(name) + + if index is None or index < argument_count: + return None + + values[index] = value + + if any(value is missing for value in values): + return None + + public_values = [] + tensor_states = [None] * public_count + + for index, value in enumerate(values): + is_tensor = type(value) is tensor_type + + if not is_tensor and ( + getattr(value, "shape", None) is not None or hasattr(value, "dtype") + ): + return None + + if is_tensor: + observed = observe_tensor( + value, + scalar=scalar_sources[index], + ) + + if observed is None: + return None + + value_key, shape, stride, span = observed + tensor_states[index] = (shape, stride, span) + + if len(shape) <= required_dims[index]: + return None + else: + if dynamic_sources[index]: + return None + + value_key = _runtime_literal_key(value) + + if value_key is None: + return None + + public_values.append((public_args[index], value_key)) + + if alias_signature is None: + aliases = [] + + for writer_index, reader_index, pair in alias_pairs: + writer_state = tensor_states[writer_index] + reader_state = tensor_states[reader_index] + + if writer_state is None or reader_state is None: + continue + + if values[writer_index] is values[reader_index]: + aliases.append(pair) + continue + + first_device, first_start, first_end = writer_state[2] + second_device, second_start, second_end = reader_state[2] + + if first_device != second_device: + continue + + if first_start == first_end or second_start == second_end: + continue + + if first_start < second_end and second_start < first_end: + aliases.append(pair) + + alias_key = tuple(aliases) + else: + alias_key = _runtime_literal_key(alias_signature) + + if alias_key is None: + return None + + key = ( + "runtime-structural-observer-v2", + (argument_count, public_args[:argument_count], tuple(kwargs)), + tuple(public_values), + tuple(residual), + alias_key, + ) + hash(key) + except ( + AttributeError, + KeyError, + RuntimeError, + TypeError, + ValueError, + OverflowError, + ): + return None + + return key + + return observe + + def _triton_specialization_key(compilation, args, kwargs): from ninetoothed.auto_tuner import AutoTuner @@ -710,6 +1350,8 @@ def validate(public): for validator in validators: validator(public) + validate._ninetoothed_observer_safe = True + return validate diff --git a/src/ninetoothed/compiler/runtime.py b/src/ninetoothed/compiler/runtime.py index 6f81bd2..85e08ca 100644 --- a/src/ninetoothed/compiler/runtime.py +++ b/src/ninetoothed/compiler/runtime.py @@ -7,6 +7,7 @@ import shutil import sys import tempfile +import threading import weakref from dataclasses import dataclass, replace from pathlib import Path @@ -788,6 +789,29 @@ def cacheable(self) -> bool: and (self.empty or self.binding_plan is not None) ) + @property + def structurally_cacheable(self) -> bool: + invocation_plan = self.invocation_plan + + return ( + not self.empty + and self.binding_plan is not None + and invocation_plan is not None + and not getattr(invocation_plan, "requires_values", True) + and getattr(invocation_plan, "structurally_rebindable", False) + ) + + def detached_for_structural_cache(self): + if not self.structurally_cacheable: + return None + + return replace( + self, + binding_plan=None, + owner_refs=None, + cache_token=None, + ) + def matches(self, args, kwargs, *, identity_verified=False, call_key=None): if not self.guard.matches( args, @@ -804,6 +828,14 @@ def matches(self, args, kwargs, *, identity_verified=False, call_key=None): ) +@dataclass(frozen=True, kw_only=True) +class _PendingStructuralPromotion: + """Weakly remember a structural hit until its tensor objects persist.""" + + identity: tuple + owner_refs: tuple[weakref.ReferenceType, ...] + + _VERIFIED_RUNTIME_CALL_CACHE_SIZE = 8 @@ -824,6 +856,218 @@ def _remember_verified_runtime_call(prepared_calls, identity, prepared): prepared_calls[identity] = prepared +def _remember_structural_runtime_call(structural_calls, pending_calls, key, prepared): + """Insert a structural entry and evict its matching pending state with it.""" + structural_calls.pop(key, None) + pending_calls.pop(key, None) + structural_calls[key] = prepared + + while len(structural_calls) > _VERIFIED_RUNTIME_CALL_CACHE_SIZE: + evicted = next(iter(structural_calls)) + structural_calls.pop(evicted, None) + pending_calls.pop(evicted, None) + + +def _touch_structural_runtime_call(structural_calls, pending_calls, key): + """Move a structural entry and any pending state to the LRU tail.""" + prepared = structural_calls.pop(key, None) + + if prepared is None: + pending_calls.pop(key, None) + + return None + + structural_calls[key] = prepared + + pending = pending_calls.pop(key, None) + + if pending is not None: + pending_calls[key] = pending + + return prepared + + +def _runtime_owner_refs_match(owner_refs, args, kwargs): + """Check weak owner references against the current tensor objects by identity.""" + if owner_refs is None: + return False + + current_refs = _runtime_owner_refs(args, kwargs) + + if current_refs is None or len(current_refs) != len(owner_refs): + return False + + return all( + previous() is current() for previous, current in zip(owner_refs, current_refs) + ) + + +def _runtime_literal_key(value): + if not _is_cacheable_runtime_literal(value): + return None + + if isinstance(value, tuple): + items = tuple(_runtime_literal_key(item) for item in value) + + if any(item is None for item in items): + return None + + return ("tuple", items) + + return ("literal", type(value), repr(value)) + + +def _runtime_tensor_key(value, *, scalar=False): + try: + shape = tuple(int(size) for size in value.shape) + stride = getattr(value, "stride", None) + + if not callable(stride): + return None + + stride = tuple(int(size) for size in stride()) + data_ptr = getattr(value, "data_ptr", None) + storage_offset = getattr(value, "storage_offset", None) + element_size = getattr(value, "element_size", None) + + if ( + not callable(data_ptr) + or not callable(storage_offset) + or not callable(element_size) + ): + return None + + if len(shape) != len(stride): + return None + + dtype = str(value.dtype).split(".")[-1] + device = getattr(value, "device", None) + + if device is None: + return None + + state = ( + type(value), + len(shape), + shape, + stride, + dtype, + str(device), + int(storage_offset()), + ) + + if scalar: + item = value.item() + scalar_key = _runtime_literal_key(item) + + if scalar_key is None: + return None + + state += (scalar_key,) + + hash(state) + except (AttributeError, RuntimeError, TypeError, ValueError, OverflowError): + return None + + return ("tensor", state) + + +def _runtime_structural_value_key(value, *, scalar=False): + if getattr(value, "shape", None) is not None and hasattr(value, "dtype"): + return _runtime_tensor_key(value, scalar=scalar) + + return _runtime_literal_key(value) + + +def _runtime_structural_key( + abi, + args, + kwargs, + public, + *, + bound_public=None, + alias_signature=None, +): + """Build a tensor-identity-independent key for a rebindable launch plan.""" + if bound_public is None: + bound_public = public + + try: + call_form = ( + len(args), + tuple(abi.public_args[: len(args)]), + tuple(kwargs.keys()), + ) + scalar_sources = { + binding.source + for binding in abi.kernel_args + if binding.kind in {"scalar", "constexpr", "meta"} + } + public_values = tuple( + ( + name, + _runtime_structural_value_key( + public[name], + scalar=name in scalar_sources, + ), + ) + for name in abi.public_args + ) + + if any(value is None for _name, value in public_values): + return None + + if any(binding.kind.startswith("jagged_") for binding in abi.kernel_args): + return None + + runtime_values = [] + + for binding in abi.kernel_args: + value = _binding_value(binding, bound_public) + value_key = _runtime_structural_value_key( + value, + scalar=binding.kind in {"scalar", "constexpr", "meta"}, + ) + + if value_key is None: + return None + + runtime_values.append( + ( + binding.name, + binding.kind, + binding.source, + binding.dim, + value_key, + ) + ) + + alias_key = _runtime_literal_key(alias_signature) + + if alias_key is None: + return None + + key = ( + "runtime-structural-v1", + call_form, + public_values, + tuple(runtime_values), + alias_key, + ) + hash(key) + except ( + AttributeError, + KeyError, + RuntimeError, + TypeError, + ValueError, + OverflowError, + ): + return None + + return key + + def _runtime_owner_refs(args, kwargs, binding_plan=None): refs = [] identities = set() @@ -881,6 +1125,11 @@ def _verified_runtime_launch(launch): active_identity = None active = None prepared_calls = {} + structural_calls = {} + pending_promotions = {} + structural_lock = threading.Lock() + structural_key_builder = getattr(launch, "_ninetoothed_structural_key", None) + structural_observer = getattr(launch, "_ninetoothed_structural_observer", None) def evict(identity, token): nonlocal active, active_identity @@ -898,6 +1147,70 @@ def activate(identity, prepared): active_identity = identity active = prepared + def deactivate(): + nonlocal active, active_identity + active = None + active_identity = None + + def remember_structural(key, prepared): + if key is None: + return None + + detached = prepared.detached_for_structural_cache() + + if detached is None: + return None + + with structural_lock: + _remember_structural_runtime_call( + structural_calls, + pending_promotions, + key, + detached, + ) + return detached + + def remember_pending(identity, key, args, kwargs): + owner_refs = _runtime_owner_refs(args, kwargs) + + with structural_lock: + if key not in structural_calls: + pending_promotions.pop(key, None) + + return + + if owner_refs is None: + pending_promotions.pop(key, None) + + return + + pending_promotions.pop(key, None) + pending_promotions[key] = _PendingStructuralPromotion( + identity=identity, + owner_refs=owner_refs, + ) + + def promote_structural(identity, prepared, args, kwargs): + rebinder = getattr(launch, "_ninetoothed_rebind_structural", None) + + if rebinder is None: + return None + + token = object() + + def collected(_reference): + evict(identity, token) + + promoted = rebinder(prepared, args, kwargs, collected, token) + + if promoted is None: + return None + + _remember_verified_runtime_call(prepared_calls, identity, promoted) + activate(identity, promoted) + + return promoted + def remember(identity, prepared): token = object() @@ -956,6 +1269,58 @@ def verified(*args, **kwargs): return launch._ninetoothed_invoke_prepared(cached, args, kwargs) + structural_key = None + + if structural_observer is not None: + try: + structural_key = structural_observer(args, kwargs) + except Exception: # noqa: BLE001 + structural_key = None + + if structural_key is None and structural_key_builder is not None: + structural_key = structural_key_builder(args, kwargs) + + with structural_lock: + structural = _touch_structural_runtime_call( + structural_calls, + pending_promotions, + structural_key, + ) + pending = ( + pending_promotions.pop(structural_key, None) + if structural is not None + else None + ) + + if structural is not None: + deactivate() + + # A structural plan is safe to invoke through its detached + # CallRefs, but promoting it lets the next call skip the + # structural lookup. Require the same weakly-held tensor + # objects on two consecutive structural hits before doing so. + if ( + pending is not None + and pending.identity == identity + and _runtime_owner_refs_match( + pending.owner_refs, + args, + kwargs, + ) + ): + promoted = promote_structural(identity, structural, args, kwargs) + else: + promoted = None + + if promoted is None: + remember_pending(identity, structural_key, args, kwargs) + + return launch._ninetoothed_invoke_prepared( + promoted or structural, + args, + kwargs, + ) + prepared = launch._ninetoothed_prepare(args, kwargs) cached_prepared = remember( _two_tensor_call_identity(prepared.guard.two_tensor_state) @@ -963,6 +1328,7 @@ def verified(*args, **kwargs): else identity, prepared, ) + remember_structural(structural_key, prepared) return launch._ninetoothed_invoke_prepared( cached_prepared or prepared, @@ -982,9 +1348,18 @@ def _runtime_wrapper( binding_overrides=None, prepare_invocation=None, validate_bindings=None, + structural_key=None, + structural_observer=None, ): overrides = dict(binding_overrides or {}) + if validate_bindings is not None and not getattr( + validate_bindings, + "_ninetoothed_observer_safe", + False, + ): + structural_observer = None + def prepare(args, kwargs, *, public=None): if public is None: public = _public_values(abi, args, kwargs, specs=specs) @@ -1049,6 +1424,43 @@ def prepare(args, kwargs, *, public=None): invocation_plan=invocation_plan, ) + def build_structural_key(args, kwargs, *, public=None): + if structural_key is None: + return None + + if public is None: + public = _public_values(abi, args, kwargs, specs=specs) + + bound_public = dict(public) | overrides + + if validate_bindings is not None: + validate_bindings(bound_public) + + if _empty_launch(abi, public): + return None + + try: + key = structural_key(args, kwargs, public, bound_public) + hash(key) + except Exception: # noqa: BLE001 + return None + + return key + + def build_structural_observer( + args, + kwargs, + *, + public=None, + alias_signature=None, + ): + if structural_observer is None: + return None + + del public, alias_signature + + return structural_observer(args, kwargs) + def invoke(prepared, args, kwargs): if prepared.empty: return _first_output_from_call(abi, args, kwargs) @@ -1085,6 +1497,34 @@ def invoke(prepared, args, kwargs): return result return _first_output_from_call(abi, args, kwargs) + def rebind_structural(prepared, args, kwargs, callback, token): + invocation_plan = prepared.invocation_plan + + if ( + prepared.empty + or invocation_plan is None + or getattr(invocation_plan, "requires_values", True) + or not getattr(invocation_plan, "structurally_rebindable", False) + ): + return None + + owner_refs = _runtime_owner_refs(args, kwargs) + + if owner_refs is None: + return None + + owners = tuple(reference() for reference in owner_refs) + + if any(owner is None for owner in owners): + return None + + return replace( + prepared, + guard=_VerifiedRuntimeCall.from_call(abi, args, kwargs), + owner_refs=tuple(weakref.ref(owner, callback) for owner in owners), + cache_token=token, + ) + def launch(*args, **kwargs): public = _public_values(abi, args, kwargs, specs=specs) bound_public = dict(public) | overrides @@ -1114,6 +1554,19 @@ def launch(*args, **kwargs): launch._ninetoothed_prepare = prepare launch._ninetoothed_invoke_prepared = invoke + launch._ninetoothed_rebind_structural = rebind_structural + launch._ninetoothed_structural_key = ( + build_structural_key + if structural_key is not None and prepare_invocation is not None + else None + ) + launch._ninetoothed_structural_observer = ( + build_structural_observer + if structural_observer is not None + and structural_key is not None + and prepare_invocation is not None + else None + ) return launch diff --git a/tests/test_triton_runtime_auto_tuning.py b/tests/test_triton_runtime_auto_tuning.py index 6f4756d..44b4249 100644 --- a/tests/test_triton_runtime_auto_tuning.py +++ b/tests/test_triton_runtime_auto_tuning.py @@ -70,6 +70,47 @@ def __init__( self._data_ptr = id(self) if data_ptr is None else data_ptr self._storage_offset = 0 + def stride(self, dim=None): + if dim is None: + return self._stride + + return self._stride[dim] + + def data_ptr(self): + return self._data_ptr + + def element_size(self): + return 4 + + def storage_offset(self): + return self._storage_offset + + def numel(self): + result = 1 + + for size in self.shape: + result *= size + return result + + +class _NonWeakrefTensor: + __slots__ = ( + "shape", + "_stride", + "dtype", + "device", + "_data_ptr", + "_storage_offset", + ) + + def __init__(self, shape): + self.shape = tuple(shape) + self._stride = _contiguous_stride(shape) + self.dtype = "float32" + self.device = "cuda:0" + self._data_ptr = id(self) + self._storage_offset = 0 + def stride(self): return self._stride @@ -274,6 +315,849 @@ def _verified_runtime_fixture(*, with_constexpr=False, outputs=()): return runtime._verified_runtime_launch(wrapped), calls +class _RebindableInvocation: + requires_values = False + structurally_rebindable = True + + def __init__(self, calls): + self._calls = calls + + def __call__(self, _values, args, kwargs): + if "output" in kwargs: + value = kwargs["output"] + elif len(args) > 1 and getattr(args[1], "shape", None) is not None: + value = args[1] + else: + value = args[0] + + self._calls.append(weakref.ref(value)) + + +class _NonRebindableInvocation(_RebindableInvocation): + structurally_rebindable = False + + +def _structural_key_builder(abi): + def build(args, kwargs, public, bound_public): + return runtime._runtime_structural_key( + abi, + args, + kwargs, + public, + bound_public=bound_public, + alias_signature=triton_materializer._runtime_alias_signature(abi, public), + ) + + return build + + +def _rebindable_verified_fixture(*, with_output=False, with_constexpr=False, specs=()): + public_args = ["value"] + bindings = [LaunchBinding(name="value", kind="tensor", source="value")] + + if with_output: + public_args.append("output") + bindings.append(LaunchBinding(name="output", kind="tensor", source="output")) + + if with_constexpr: + public_args.append("scale") + bindings.append(LaunchBinding(name="scale", kind="constexpr", source="scale")) + + abi = LaunchABI( + public_args=tuple(public_args), + kernel_args=tuple(bindings), + outputs=("output",) if with_output else (), + ) + prepare_calls = [] + invoked = [] + + def prepare_invocation(_values, _static_values, _call_sources): + prepare_calls.append(object()) + + return _RebindableInvocation(invoked) + + wrapped = runtime._runtime_wrapper( + lambda *_values: None, + abi, + specs=specs, + prepare_invocation=prepare_invocation, + structural_key=_structural_key_builder(abi), + ) + + return runtime._verified_runtime_launch(wrapped), abi, prepare_calls, invoked + + +def _observer_runtime_fixture(*, observer=None, validate_bindings=None): + abi = LaunchABI( + public_args=("value", "output", "scale"), + kernel_args=( + LaunchBinding( + name="value", + kind="tensor", + source="value", + access="read", + ), + LaunchBinding(name="value_size", kind="shape", source="value", dim=0), + LaunchBinding( + name="value_stride", + kind="stride", + source="value", + dim=0, + ), + LaunchBinding( + name="output", + kind="tensor", + source="output", + access="write", + ), + LaunchBinding( + name="output_size", + kind="shape", + source="output", + dim=0, + ), + LaunchBinding( + name="output_stride", + kind="stride", + source="output", + dim=0, + ), + LaunchBinding(name="scale", kind="constexpr", source="scale"), + ), + outputs=("output",), + ) + prepare_calls = [] + + def prepare_invocation(_values, _static_values, _call_sources): + prepare_calls.append(object()) + + return _RebindableInvocation([]) + + if observer is None: + observer = triton_materializer._triton_structural_observer( + abi, + tensor_type=_FakeTensor, + ) + + wrapped = runtime._runtime_wrapper( + lambda *_values: None, + abi, + prepare_invocation=prepare_invocation, + validate_bindings=validate_bindings, + structural_key=_structural_key_builder(abi), + structural_observer=observer, + ) + + return runtime._verified_runtime_launch(wrapped), abi, prepare_calls + + +def _observer_values(*, size=4, stride=1, dtype="float32", device="cuda:0"): + value = _FakeTensor( + (size,), + stride=(stride,), + dtype=dtype, + device=device, + data_ptr=1024, + ) + output = _FakeTensor( + (size,), + stride=(stride,), + dtype=dtype, + device=device, + data_ptr=2048, + ) + + return value, output + + +def test_triton_observer_key_equivalence_covers_layout_scalar_form_and_alias(): + observer = triton_materializer._triton_structural_observer( + _observer_runtime_fixture()[1], + tensor_type=_FakeTensor, + ) + first = _observer_values() + second = _observer_values() + + assert observer((*first, 2), {}) == observer((*second, 2), {}) + + for changed in ( + ("shape", _observer_values(size=8)), + ("stride", _observer_values(stride=2)), + ("dtype", _observer_values(dtype="float16")), + ("device", _observer_values(device="cuda:1")), + ): + assert observer((*first, 2), {}) != observer((*changed[1], 2), {}) + + zero_sized = _observer_values(size=0) + assert observer((*zero_sized, 2), {}) is not None + + offset_value, offset_output = _observer_values() + offset_value._storage_offset = 1 + assert observer((*first, 2), {}) != observer((offset_value, offset_output, 2), {}) + assert observer((*first, 2), {}) != observer((*first, 3), {}) + assert observer( + ( + *first[:1], + first[1], + ), + {"scale": 2}, + ) != observer((*first, 2), {}) + assert observer((first[0], first[0], 2), {}) != observer((*first, 2), {}) + + +def test_verified_observer_hit_avoids_public_revalidation(monkeypatch): + validated = 0 + + def validate(_public): + nonlocal validated + validated += 1 + + validate._ninetoothed_observer_safe = True + original_public = runtime._public_values + public_calls = 0 + observer_calls = 0 + + def public(*args, **kwargs): + nonlocal public_calls + public_calls += 1 + + return original_public(*args, **kwargs) + + original_observer = triton_materializer._triton_structural_observer( + _observer_runtime_fixture()[1], + tensor_type=_FakeTensor, + ) + + def observed(*args, **kwargs): + nonlocal observer_calls + observer_calls += 1 + + return original_observer(*args, **kwargs) + + # Rebuild with the counting observer so the wrapper captures it before + # monkeypatching the public-value helper. + launch, _abi, prepare_calls = _observer_runtime_fixture( + observer=observed, + validate_bindings=validate, + ) + monkeypatch.setattr(runtime, "_public_values", public) + first = _observer_values() + replacement = _observer_values() + + launch(*first, 2) + launch(*replacement, 2) + + assert len(prepare_calls) == 1 + assert observer_calls == 2 + assert public_calls == 1 + assert validated == 1 + + +def test_triton_observer_fails_closed_for_unknown_types_and_jagged_bindings(): + abi = _observer_runtime_fixture()[1] + observer = triton_materializer._triton_structural_observer( + abi, + tensor_type=_FakeTensor, + ) + + class TensorSubclass(_FakeTensor): + pass + + value, output = _observer_values() + subclass = TensorSubclass((4,), data_ptr=4096) + assert observer((subclass, output, 2), {}) is None + + jagged = LaunchABI( + public_args=("value",), + kernel_args=( + LaunchBinding(name="value", kind="jagged_values", source="value"), + ), + ) + assert ( + triton_materializer._triton_structural_observer( + jagged, + tensor_type=_FakeTensor, + ) + is None + ) + + +def test_runtime_custom_validator_disables_observer_and_uses_structural_fallback(): + validated = [] + observed = [] + + def validate(public): + validated.append(public["scale"]) + + def observer(*args, **kwargs): + observed.append((args, kwargs)) + + return ("must-not-be-used",) + + launch, _abi, prepare_calls = _observer_runtime_fixture( + observer=observer, + validate_bindings=validate, + ) + first = _observer_values() + replacement = _observer_values() + + launch(*first, 2) + launch(*replacement, 2) + + assert observed == [] + assert len(prepare_calls) == 1 + assert validated + + +def test_verified_structural_cache_reuses_plan_with_current_output(): + launch, _abi, prepare_calls, invoked = _rebindable_verified_fixture( + with_output=True + ) + value = _FakeTensor((4,), data_ptr=1024) + output = _FakeTensor((4,), data_ptr=2048) + replacement_value = _FakeTensor((4,), data_ptr=4096) + replacement_output = _FakeTensor((4,), data_ptr=8192) + + assert launch(value, output=output) is output + assert launch(replacement_value, output=replacement_output) is replacement_output + assert launch(replacement_value, output=replacement_output) is replacement_output + + assert len(prepare_calls) == 1 + assert [reference() for reference in invoked] == [ + output, + replacement_output, + replacement_output, + ] + + nonlocals = inspect.getclosurevars(launch).nonlocals + structural = nonlocals["structural_calls"] + prepared = next(iter(structural.values())) + assert prepared.binding_plan is None + assert prepared.owner_refs is None + assert prepared.cache_token is None + assert nonlocals["active_identity"] == runtime._runtime_call_identity( + (replacement_value,), {"output": replacement_output} + ) + assert nonlocals["active"].matches( + (replacement_value,), + {"output": replacement_output}, + identity_verified=True, + ) + assert nonlocals["active"].owner_refs is not None + + replacement_value_ref = weakref.ref(replacement_value) + replacement_output_ref = weakref.ref(replacement_output) + del replacement_value, replacement_output + gc.collect() + + assert replacement_value_ref() is None + assert replacement_output_ref() is None + assert inspect.getclosurevars(launch).nonlocals["active"] is None + + +def test_verified_structural_hit_promotes_only_after_persistent_replacement(): + launch, _abi, prepare_calls, invoked = _rebindable_verified_fixture( + with_output=True + ) + value = _FakeTensor((4,), data_ptr=1024) + output = _FakeTensor((4,), data_ptr=2048) + replacement_value = _FakeTensor((4,), data_ptr=4096) + replacement_output = _FakeTensor((4,), data_ptr=8192) + + launch(value, output=output) + launch(replacement_value, output=replacement_output) + + nonlocals = inspect.getclosurevars(launch).nonlocals + assert nonlocals["active"] is None + pending = nonlocals["pending_promotions"] + assert len(pending) == 1 + pending_entry = next(iter(pending.values())) + assert pending_entry.identity == runtime._runtime_call_identity( + (replacement_value,), {"output": replacement_output} + ) + assert [reference() for reference in pending_entry.owner_refs] == [ + replacement_value, + replacement_output, + ] + + launch(replacement_value, output=replacement_output) + assert len(prepare_calls) == 1 + assert len(invoked) == 3 + nonlocals = inspect.getclosurevars(launch).nonlocals + assert nonlocals["active"] is not None + assert nonlocals["active_identity"] == runtime._runtime_call_identity( + (replacement_value,), {"output": replacement_output} + ) + assert nonlocals["pending_promotions"] == {} + + launch(replacement_value, output=replacement_output) + assert len(prepare_calls) == 1 + assert len(invoked) == 4 + + +def test_verified_structural_hit_requires_every_owner_to_persist(): + launch, _abi, prepare_calls, _invoked = _rebindable_verified_fixture( + with_output=True + ) + value = _FakeTensor((4,), data_ptr=1024) + output = _FakeTensor((4,), data_ptr=2048) + replacement_value = _FakeTensor((4,), data_ptr=4096) + first_output = _FakeTensor((4,), data_ptr=8192) + second_output = _FakeTensor((4,), data_ptr=16384) + + launch(value, output=output) + launch(replacement_value, output=first_output) + launch(replacement_value, output=second_output) + + nonlocals = inspect.getclosurevars(launch).nonlocals + assert nonlocals["active"] is None + assert len(prepare_calls) == 1 + pending = next(iter(nonlocals["pending_promotions"].values())) + assert [reference() for reference in pending.owner_refs] == [ + replacement_value, + second_output, + ] + + launch(replacement_value, output=second_output) + nonlocals = inspect.getclosurevars(launch).nonlocals + assert nonlocals["active"] is not None + + +def test_verified_structural_hit_does_not_promote_fresh_object_after_gc(): + launch, _abi, _prepare_calls, _invoked = _rebindable_verified_fixture() + value = _FakeTensor((4,), data_ptr=1024) + launch(value) + + temporary = _FakeTensor((4,), data_ptr=4096) + temporary_ref = weakref.ref(temporary) + launch(temporary) + pending_ref = next( + iter(inspect.getclosurevars(launch).nonlocals["pending_promotions"].values()) + ).owner_refs[0] + del temporary + gc.collect() + + assert temporary_ref() is None + assert pending_ref() is None + + fresh = _FakeTensor((4,), data_ptr=8192) + launch(fresh) + nonlocals = inspect.getclosurevars(launch).nonlocals + assert nonlocals["active"] is None + assert nonlocals["pending_promotions"] + assert next(iter(nonlocals["pending_promotions"].values())).owner_refs[0]() is fresh + + +def test_verified_structural_pending_owner_check_is_aba_safe(): + launch, _abi, _prepare_calls, _invoked = _rebindable_verified_fixture() + value = _FakeTensor((4,), data_ptr=1024) + launch(value) + first = _FakeTensor((4,), data_ptr=4096) + launch(first) + + pending_ref = next( + iter(inspect.getclosurevars(launch).nonlocals["pending_promotions"].values()) + ).owner_refs[0] + first_ref = weakref.ref(first) + del first + gc.collect() + + replacement = _FakeTensor((4,), data_ptr=8192) + assert first_ref() is None + assert pending_ref() is None + assert not runtime._runtime_owner_refs_match( + (pending_ref,), + (replacement,), + {}, + ) + + launch(replacement) + assert inspect.getclosurevars(launch).nonlocals["active"] is None + + +def test_verified_structural_pending_is_bounded_with_structural_lru(): + launch, _abi, _prepare_calls, _invoked = _rebindable_verified_fixture() + evicted_key = None + + for size in range(1, 10): + launch(_FakeTensor((size,), data_ptr=1000 + size)) + launch(_FakeTensor((size,), data_ptr=2000 + size)) + + if size == 1: + evicted_key = next( + iter(inspect.getclosurevars(launch).nonlocals["structural_calls"]) + ) + + nonlocals = inspect.getclosurevars(launch).nonlocals + assert len(nonlocals["structural_calls"]) == 8 + assert len(nonlocals["pending_promotions"]) == 8 + assert tuple(nonlocals["pending_promotions"]) == tuple( + nonlocals["structural_calls"] + ) + + # Simulate a structural hit racing with LRU eviction: a late pending + # insertion for the evicted entry must not create an orphan ninth item. + late = _FakeTensor((1,), data_ptr=9999) + nonlocals["remember_pending"]( + runtime._runtime_call_identity((late,), {}), + evicted_key, + (late,), + {}, + ) + assert evicted_key not in nonlocals["pending_promotions"] + assert len(nonlocals["pending_promotions"]) == 8 + + +def test_verified_structural_hit_fails_safe_for_nonweakref_owner(): + class RebindableNoopInvocation: + requires_values = False + structurally_rebindable = True + + def __call__(self, _values, _args, _kwargs): + return None + + abi = LaunchABI( + public_args=("value", "observer"), + kernel_args=(LaunchBinding(name="value", kind="tensor", source="value"),), + outputs=(), + ) + prepare_calls = [] + + def prepare_invocation(_values, _static_values, _call_sources): + prepare_calls.append(object()) + + return RebindableNoopInvocation() + + wrapped = runtime._runtime_wrapper( + lambda *_values: None, + abi, + prepare_invocation=prepare_invocation, + structural_key=_structural_key_builder(abi), + ) + launch = runtime._verified_runtime_launch(wrapped) + first = _FakeTensor((4,)) + first_observer = _NonWeakrefTensor((4,)) + replacement = _FakeTensor((4,)) + replacement_observer = _NonWeakrefTensor((4,)) + + launch(first, first_observer) + launch(replacement, replacement_observer) + + nonlocals = inspect.getclosurevars(launch).nonlocals + assert len(prepare_calls) == 1 + assert len(nonlocals["structural_calls"]) == 1 + assert nonlocals["active"] is None + assert nonlocals["pending_promotions"] == {} + + +def test_verified_structural_cache_rejects_call_form_changes(): + launch, _abi, prepare_calls, _invoked = _rebindable_verified_fixture( + with_output=True + ) + value = _FakeTensor((4,)) + output = _FakeTensor((4,)) + replacement_value = _FakeTensor((4,)) + replacement_output = _FakeTensor((4,)) + + launch(value, output=output) + launch(replacement_value, replacement_output) + + assert len(prepare_calls) == 2 + + +def test_verified_structural_cache_rejects_alias_changes(): + launch, _abi, prepare_calls, _invoked = _rebindable_verified_fixture( + with_output=True + ) + value = _FakeTensor((4,), data_ptr=1024) + output = _FakeTensor((4,), data_ptr=2048) + replacement = _FakeTensor((4,), data_ptr=4096) + + launch(value, output=output) + launch(replacement, output=replacement) + + assert len(prepare_calls) == 2 + + +@pytest.mark.parametrize( + ("field", "value"), + ( + ("shape", (8,)), + ("_stride", (2,)), + ("dtype", "float16"), + ("device", "cuda:1"), + ), +) +def test_verified_structural_cache_rejects_tensor_contract_changes(field, value): + launch, _abi, prepare_calls, _invoked = _rebindable_verified_fixture() + first = _FakeTensor((4,)) + replacement = _FakeTensor((4,)) + + launch(first) + setattr(replacement, field, value) + launch(replacement) + + assert len(prepare_calls) == 2 + + +def test_verified_structural_cache_rejects_scalar_changes(): + launch, _abi, prepare_calls, _invoked = _rebindable_verified_fixture( + with_constexpr=True + ) + first = _FakeTensor((4,)) + replacement = _FakeTensor((4,)) + + launch(first, 2) + launch(replacement, 3) + + assert len(prepare_calls) == 2 + + +def test_verified_structural_cache_rejects_storage_offset_changes(): + launch, _abi, prepare_calls, _invoked = _rebindable_verified_fixture() + first = _FakeTensor((4,)) + replacement = _FakeTensor((4,)) + replacement._storage_offset = 1 + + launch(first) + launch(replacement) + + assert len(prepare_calls) == 2 + + +def test_verified_structural_cache_rejects_zero_dimensional_scalar_changes(): + launch, _abi, prepare_calls, _invoked = _rebindable_verified_fixture( + with_constexpr=True + ) + value = _FakeTensor((4,)) + scale = torch.tensor(2) + + launch(value, scale) + scale.fill_(3) + launch(value, scale) + + assert len(prepare_calls) == 2 + + +def test_verified_structural_cache_requires_rebindable_invocation(): + abi = LaunchABI( + public_args=("value",), + kernel_args=(LaunchBinding(name="value", kind="tensor", source="value"),), + ) + prepare_calls = [] + + def prepare_invocation(_values, _static_values, _call_sources): + prepare_calls.append(object()) + + return _NonRebindableInvocation([]) + + wrapped = runtime._runtime_wrapper( + lambda *_values: None, + abi, + prepare_invocation=prepare_invocation, + structural_key=_structural_key_builder(abi), + ) + launch = runtime._verified_runtime_launch(wrapped) + + launch(_FakeTensor((4,))) + launch(_FakeTensor((4,))) + + assert len(prepare_calls) == 2 + + +def test_verified_structural_cache_does_not_retain_tensors(): + launch, _abi, _prepare_calls, invoked = _rebindable_verified_fixture( + with_output=True + ) + value = _FakeTensor((4,)) + output = _FakeTensor((4,)) + value_ref = weakref.ref(value) + output_ref = weakref.ref(output) + + launch(value, output=output) + del value, output + gc.collect() + + assert value_ref() is None + assert output_ref() is None + assert invoked[0]() is None + + structural = inspect.getclosurevars(launch).nonlocals["structural_calls"] + prepared = next(iter(structural.values())) + assert prepared.binding_plan is None + assert prepared.owner_refs is None + + +def test_verified_structural_cache_preserves_invalid_argument_errors(): + specs = ( + SimpleNamespace( + name="value", + ndim=1, + dtype=None, + attrs={"source_ndim": 1}, + ), + ) + launch, _abi, _prepare_calls, _invoked = _rebindable_verified_fixture(specs=specs) + value = _FakeTensor((4,)) + + with pytest.raises(TypeError, match="Unknown kernel arguments"): + launch(value, unknown=True) + + with pytest.raises(TypeError, match="rank 2"): + launch(_FakeTensor((2, 2))) + + +def test_tuned_structural_cache_reuses_selected_plan_with_current_output(): + abi = LaunchABI( + public_args=("value", "output"), + kernel_args=( + LaunchBinding(name="value", kind="tensor", source="value"), + LaunchBinding(name="output", kind="tensor", source="output"), + ), + outputs=("output",), + ) + prepare_calls = [] + invoked = [] + + def candidate(): + def prepare_invocation(_values, _static_values, _call_sources): + prepare_calls.append(object()) + + return _RebindableInvocation(invoked) + + return runtime._runtime_wrapper( + lambda *_values: None, + abi, + prepare_invocation=prepare_invocation, + structural_key=_structural_key_builder(abi), + ) + + candidates = (candidate(), candidate()) + compilation = SimpleNamespace( + launch_abi=abi, + kernel=SimpleNamespace(tensors=()), + ) + tuner = _FakeTuner( + candidates, + lambda args, kwargs: triton_materializer._triton_specialization_key( + compilation, args, kwargs + ), + ) + handle = SimpleNamespace(_selected_tuning_candidate=None) + launch = triton_materializer._tuned_runtime_launch( + tuner, + dict(zip(candidates, ({"id": "first"}, {"id": "second"}))), + handle, + compilation, + ) + value = _FakeTensor((4,), data_ptr=1024) + output = _FakeTensor((4,), data_ptr=2048) + replacement_value = _FakeTensor((4,), data_ptr=4096) + replacement_output = _FakeTensor((4,), data_ptr=8192) + + assert launch(value, output=output) is output + handle._selected_tuning_candidate = None + assert launch(replacement_value, output=replacement_output) is replacement_output + + nonlocals = inspect.getclosurevars(launch).nonlocals + assert nonlocals["active"] is None + assert handle._selected_tuning_candidate == {"id": "first"} + pending_promotions = inspect.getclosurevars(nonlocals["find_structural"]).nonlocals[ + "pending_promotions" + ] + pending = next(iter(pending_promotions.values())) + assert [reference() for reference in pending.owner_refs] == [ + replacement_value, + replacement_output, + ] + + assert launch(replacement_value, output=replacement_output) is replacement_output + + nonlocals = inspect.getclosurevars(launch).nonlocals + assert len(prepare_calls) == 1 + assert [reference() for reference in invoked] == [ + output, + replacement_output, + replacement_output, + ] + assert handle._selected_tuning_candidate == {"id": "first"} + assert nonlocals["active_identity"] == runtime._runtime_call_identity( + (replacement_value,), {"output": replacement_output} + ) + promoted = nonlocals["active"][2] + assert promoted.matches( + (replacement_value,), + {"output": replacement_output}, + identity_verified=True, + ) + assert promoted.owner_refs is not None + assert pending_promotions == {} + + +def test_tuned_structural_cache_validates_candidate_binding_overrides(): + abi = LaunchABI( + public_args=("value", "output", "block"), + kernel_args=( + LaunchBinding(name="value", kind="tensor", source="value"), + LaunchBinding(name="output", kind="tensor", source="output"), + LaunchBinding(name="block", kind="meta", source="block"), + ), + outputs=("output",), + ) + prepare_calls = [] + invoked = [] + validated_blocks = [] + + def validate(bound_public): + validated_blocks.append(bound_public["block"]) + + if bound_public["block"] != 64: + raise ValueError("Candidate binding override was not applied.") + + def prepare_invocation(_values, _static_values, _call_sources): + prepare_calls.append(object()) + + return _RebindableInvocation(invoked) + + candidate = runtime._runtime_wrapper( + lambda *_values: None, + abi, + binding_overrides={"block": 64}, + prepare_invocation=prepare_invocation, + validate_bindings=validate, + structural_key=_structural_key_builder(abi), + ) + compilation = SimpleNamespace( + launch_abi=abi, + kernel=SimpleNamespace(tensors=()), + ) + tuner = _FakeTuner( + (candidate,), + lambda args, kwargs: triton_materializer._triton_specialization_key( + compilation, args, kwargs + ), + ) + handle = SimpleNamespace(_selected_tuning_candidate=None) + launch = triton_materializer._tuned_runtime_launch( + tuner, + {candidate: {"id": "only"}}, + handle, + compilation, + validate_bindings=validate, + ) + value = _FakeTensor((4,), data_ptr=1024) + output = _FakeTensor((4,), data_ptr=2048) + replacement_value = _FakeTensor((4,), data_ptr=4096) + replacement_output = _FakeTensor((4,), data_ptr=8192) + + assert launch(value, output, 32) is output + assert launch(replacement_value, replacement_output, 32) is replacement_output + + assert len(prepare_calls) == 1 + assert set(validated_blocks) == {64} + assert [reference() for reference in invoked] == [output, replacement_output] + + def test_verified_runtime_launch_rebinds_zero_dimensional_value(monkeypatch): binding_calls = 0 original_bound = runtime._bound_values @@ -537,12 +1421,84 @@ def generated_launch(value, size, _ninetoothed_num_warps=4): assert unowned_invocation is not None assert not unowned_invocation.requires_values + assert unowned_invocation.structurally_rebindable assert value_ref() is None with pytest.raises(ValueError, match="positive"): prepare(("pointer", 0)) +def test_triton_prepared_invocation_rejects_ambiguous_call_refs(monkeypatch): + kernel_calls = [] + + class FakeKernel: + def __getitem__(self, _grid): + def launch(*args, **kwargs): + kernel_calls.append((args, kwargs)) + + return launch + + kernel = FakeKernel() + monkeypatch.setitem(globals(), "_prepared_test_kernel", kernel) + + def generated_launch(first, second, output): + _prepared_test_kernel[(1,)](first, second, output) + + prepare = triton_materializer._triton_prepare_invocation( + generated_launch, + kernel, + ) + shared = _FakeTensor((4,)) + output = _FakeTensor((4,)) + + assert prepare is not None + assert ( + prepare( + (shared, shared, output), + (None, None, None), + (("positional", 0), ("positional", 1), ("positional", 2)), + ) + is None + ) + assert kernel_calls == [] + + +def test_triton_prepared_invocation_literalizes_ambiguous_call_refs_for_scalars( + monkeypatch, +): + kernel_calls = [] + + class FakeKernel: + def __getitem__(self, _grid): + def launch(*args, **kwargs): + kernel_calls.append((args, kwargs)) + + return launch + + kernel = FakeKernel() + monkeypatch.setitem(globals(), "_prepared_test_kernel", kernel) + + def generated_launch(value, first_size, second_size): + _prepared_test_kernel[(1,)](value, first_size, second_size) + + prepare = triton_materializer._triton_prepare_invocation( + generated_launch, + kernel, + ) + value = _FakeTensor((4,)) + invocation = prepare( + (value, 64, 64), + (None, 64, 64), + (("positional", 0), ("positional", 1), ("positional", 2)), + ) + replacement = _FakeTensor((4,)) + + assert invocation is not None + assert invocation.structurally_rebindable + invocation((), (replacement, 64, 64), {}) + assert kernel_calls == [((replacement, 64, 64), {})] + + def test_triton_direct_winner_reuses_verified_binding_and_restores_aba(monkeypatch): public_calls = 0 binding_calls = 0