Copy
Source code for tensorplay._stax.stax
"""The TensorPlay native graph compiler backend.
This is a compiler backend, not a second public compiler frontend:
``tensorplay.compile`` owns
capture, guards, specialization, and graph-break policy; this module owns
lowering and executable generation for the canonical graph.
small lazy adapter. Native Stax code is loaded only when the backend is
actually selected, keeping import-time overhead out of the frontend.
"""
from __future__ import annotations
import operator
import numbers
import re
from typing import Any
from ..graph.passes import POINTWISE_FUSED_OP_NAMES
from ..graph import GraphModule, Node
from ..library import CustomOpDef as _CustomOpDef
def _nodes(value: Any):
if isinstance(value, Node):
yield value
elif isinstance(value, (tuple, list)):
for item in value:
yield from _nodes(item)
elif isinstance(value, dict):
for item in value.values():
yield from _nodes(item)
elif isinstance(value, slice):
yield from _nodes(value.start)
yield from _nodes(value.stop)
yield from _nodes(value.step)
def _native_value_leaves(value: Any, values: dict[Node, Any]) -> list[Any]:
if isinstance(value, Node):
native = values.get(value)
if isinstance(native, tuple):
return list(native)
return [] if native is None else [native]
if isinstance(value, tuple | list):
result: list[Any] = []
for item in value:
result.extend(_native_value_leaves(item, values))
return result
if isinstance(value, dict):
result = []
for item in value.values():
result.extend(_native_value_leaves(item, values))
return result
if isinstance(value, slice):
result = []
for item in (value.start, value.stop, value.step):
result.extend(_native_value_leaves(item, values))
return result
return []
def _native_output_spec(value: Any, values: dict[Node, Any]) -> Any:
if isinstance(value, Node):
native = values.get(value)
if isinstance(native, tuple):
custom = value.meta.get("custom")
template = custom.get("nested_output_template") if isinstance(custom, dict) else None
if template is None:
raise RuntimeError("native multi-output node has no output template")
return ("nested", template)
if native is None:
raise RuntimeError(f"native graph has no value for output {value.name!r}")
return ("leaf",)
if isinstance(value, tuple):
return ("tuple", tuple(_native_output_spec(item, values) for item in value))
if isinstance(value, list):
return ("list", tuple(_native_output_spec(item, values) for item in value))
if isinstance(value, dict):
return (
"dict",
tuple((key, _native_output_spec(item, values)) for key, item in value.items()),
)
raise RuntimeError("native graph outputs must be tensor values")
def _consume_template(template: Any, outputs: list[Any], index: int) -> tuple[Any, int]:
kind = template[0]
if kind in {"leaf", "tensor"}:
if index >= len(outputs):
raise RuntimeError("native graph produced too few outputs")
return outputs[index], index + 1
if kind == "nested":
return _consume_template(template[1], outputs, index)
if kind == "tuple":
result = []
for item in template[1]:
value, index = _consume_template(item, outputs, index)
result.append(value)
return tuple(result), index
if kind == "list":
result = []
for item in template[1]:
value, index = _consume_template(item, outputs, index)
result.append(value)
return result, index
if kind == "dict":
result = {}
for key, item in template[1]:
value, index = _consume_template(item, outputs, index)
result[key] = value
return result, index
raise RuntimeError(f"unknown native output template {kind!r}")
def _target_name(target: Any) -> str:
if target is operator.add:
return "add"
if target is operator.sub:
return "sub"
if target is operator.mul:
return "mul"
if target is operator.truediv:
return "div"
if target is operator.pow:
return "pow"
if target is operator.matmul:
return "matmul"
if target is operator.neg:
return "neg"
if target is operator.pos:
return "pos"
return getattr(target, "__name__", str(target))
_NATIVE_OPS = {
"add",
"sub",
"mul",
"div",
"pow",
"matmul",
"t",
"linear",
"neg",
"pos",
"abs",
"sin",
"cos",
"exp",
"log",
"sigmoid",
"sqrt",
"square",
"tanh",
"relu",
"mm",
# Tensor kernels used by the ResNet inference graph. These are kept in
# the native graph instead of falling back to the generated Python
# executor; the latter still calls every functional wrapper through the
# interpreter and is not a compiled path in any meaningful sense.
"conv2d",
"add_relu",
"batch_norm",
"max_pool2d",
"adaptive_avg_pool2d",
"flatten",
}
# The fused-op name set is shared by the graph pass and this lowering.
_CPU_FUSED_OPS = POINTWISE_FUSED_OP_NAMES
# Opcodes of the CPU fused interpreter (p10 StaxPointwiseKernels switch).
# This is a strict subset of _CPU_FUSED_OPS: graphs whose programs contain
# Triton-only opcodes fall back per lowering instead of reaching the CPU
# program runner.
_CPU_FUSED_OPCODES = {
"add": 1,
"sub": 2,
"mul": 3,
"div": 4,
"pow": 5,
"neg": 6,
"pos": 7,
"abs": 8,
"sin": 9,
"cos": 10,
"exp": 11,
"log": 12,
"sigmoid": 13,
"sqrt": 14,
"square": 15,
"tanh": 16,
"relu": 17,
"relu_grad": 18,
"abs_grad": 19,
}
# Triton-only opcode extension. Emitted programs keep the shared triple
# format; the extra opcodes never reach the CPU interpreter because the CPU
# program builder runs against _CPU_FUSED_OPCODES.
#
# ``where`` is ternary, so it is encoded as a two-instruction pair:
# ``where`` carries (cond, a) and the immediately following ``where_rest``
# carries (cond, b). The code generator pairs them by adjacency; ``where``
# itself emits no source. ``cast`` stores the target float-dtype id in its
# rhs operand slot (it has a single value operand).
_TRITON_EXTRA_OPCODES = {
"lt": 20,
"le": 21,
"gt": 22,
"ge": 23,
"eq": 24,
"ne": 25,
"where": 26,
"where_rest": 27,
"minimum": 28,
"maximum": 29,
"clamp_min": 30,
"clamp_max": 31,
"rsqrt": 32,
"exp2": 33,
"erf": 34,
"cast": 35,
}
_TRITON_OPCODES = dict(_CPU_FUSED_OPCODES, **_TRITON_EXTRA_OPCODES)
# Cast targets accepted by the fused program, keyed by ``str(dtype)``.
# Non-float casts (bool/int) stay uncompiled for now: the program's value
# space is float and the store path types outputs off the sample dtype.
_CAST_DTYPE_IDS = {
"tensorplay.float16": 1,
"tensorplay.bfloat16": 2,
"tensorplay.float32": 3,
"tensorplay.float64": 4,
}
# Backward-compatibility alias: the autograd gate keeps covering exactly the
# ops whose elementwise VJP rules exist (the CPU interpreter's surface minus
# pow). Triton-only opcodes stay out, so training graphs using them fall
# back instead of producing a wrong gradient.
_CPU_FUSED_AUTOGRAD_OPS = frozenset(_CPU_FUSED_OPCODES) - {"pow"}
def _is_scalar(value: Any) -> bool:
return isinstance(value, (bool, int, float))
def _set_scalar_attr(native_node: Any, value: Any, position: int) -> None:
if isinstance(value, bool) or isinstance(value, int):
native_node.set_int_attr("scalar_value", int(value))
elif isinstance(value, numbers.Real):
native_node.set_float_attr("scalar_value", float(value))
else:
raise TypeError(f"unsupported Stax scalar constant: {type(value)!r}")
native_node.set_int_attr("scalar_position", position)
def _int_list(value: Any) -> list[int] | None:
"""Return a constant integer list accepted by a native Stax node."""
if not isinstance(value, (tuple, list)):
return None
if any(isinstance(item, bool) or not isinstance(item, int) for item in value):
return None
return [int(item) for item in value]
def _set_int_list_attr(native_node: Any, key: str, value: Any) -> bool:
values = _int_list(value)
if values is None:
return False
native_node.set_ints_attr(key, values)
return True
def _normalize_pointwise_grad_output(grad_output: Any, reference: Any) -> Any:
"""Match the output shape expected by a fused elementwise backward.
TensorPlay's current reduction backward may hand a scalar tangent to a
custom Function for ``output.sum().backward()``. The compiled backward
contract supplies the expanded tangent, so normalize that boundary here
before entering either the p10 or Triton backward kernel.
"""
if (
grad_output.numel() == 1 and reference.numel() != 1
) or not grad_output.is_contiguous():
import tensorplay
return tensorplay.ones_like(reference, requires_grad=False) * grad_output
return grad_output
def _attach_fast_call(lowering: Any, exec_fn: Any = None) -> None:
"""Install the C steady-state trampoline for a compiled lowering.
``exec_fn`` selects the steady-state execution entry; the default is the
native ``Graph.execute`` bound method. Lowerings with a direct kernel
entry pass it here so the trampoline skips the graph walk.
"""
import tensorplay
installer = tensorplay._C._stax.install_call_trampoline
tail = [
lowering.graph_module._get_attr(target)
for target in lowering.attribute_targets
]
tail.extend(lowering.constant_values)
lowering._fast_call = installer(
lowering,
exec_fn if exec_fn is not None else lowering.graph.execute,
tail,
tensorplay.Tensor,
len(lowering.placeholders),
int(getattr(lowering, "_output_count", 1)),
getattr(lowering, "_gradient_plan", None) is not None,
int(getattr(lowering, "_native_direct", 0) or 0),
)
def _metadata_fingerprint(value: Any) -> Any:
"""Metadata snapshot for tensors outside the version-counter contract.
Only reached when ``_version`` is unavailable and the tensor is not an
inference tensor; every component is normalized so the snapshot stays
comparable across calls.
"""
shape = getattr(value, "shape", ())
if callable(shape):
shape = shape()
try:
shape = tuple(int(item) for item in shape)
except (TypeError, ValueError):
shape = repr(shape)
stride = getattr(value, "stride", ())
if callable(stride):
stride = stride()
try:
stride = tuple(int(item) for item in stride)
except (TypeError, ValueError):
stride = repr(stride)
dtype = getattr(value, "dtype", None)
if callable(dtype):
dtype = dtype()
device = getattr(value, "device", None)
if callable(device):
device = device()
return ("metadata", shape, stride, dtype, device)
class _NativeLowering:
def __init__(
self,
graph_module: GraphModule,
graph: Any,
attribute_targets: list[str],
constant_values: list[Any] | None = None,
output_count: int = 1,
native_values: dict[Node, Any] | None = None,
output_spec: Any = None,
public_output_count: int | None = None,
) -> None:
self.graph_module = graph_module
self.graph = graph
self.placeholders = graph_module.graph.placeholders
self.attribute_targets = attribute_targets
self.constant_values = list(constant_values or [])
self._output_count = output_count
self._public_output_count = (
output_count if public_output_count is None else public_output_count
)
self._output_spec = output_spec
self.native_values = dict(native_values or {})
self._tensorplay_codegen = "stax-native"
# (id, _version) memo of the last resolved input vector; attributes
# and constants appended by _bind_inputs are process-stable.
self._bind_fp: Any = None
self._last_bound_inputs: list[Any] | None = None
_attach_fast_call(self)
@staticmethod
def _input_route_fingerprint(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
import tensorplay
def fp(value: Any) -> Any:
if isinstance(value, tensorplay.Tensor):
try:
version = value._version
except RuntimeError:
# Inference tensors carry no version counter and are
# immutable, so the identity alone keys the entry: no
# metadata snapshot, no divert on the next call.
if getattr(value, "is_inference", lambda: False)():
return ("t", id(value), None)
version = _metadata_fingerprint(value)
return (
"t",
id(value),
version,
)
return ("o", id(value))
items = [fp(item) for item in args]
items.extend((k, fp(v)) for k, v in sorted(kwargs.items()))
return tuple(items)
def _bind_inputs_fresh(self, *args: Any, **kwargs: Any) -> list[Any]:
bound = self.graph_module.signature.bind_partial(*args, **kwargs)
bound.apply_defaults()
inputs = [
bound.arguments[node.target if isinstance(node.target, str) else node.name]
for node in self.placeholders
]
inputs.extend(
self.graph_module._get_attr(target) for target in self.attribute_targets
)
inputs.extend(self.constant_values)
return inputs
def _bind_inputs(self, *args: Any, **kwargs: Any) -> list[Any]:
# Attribute targets and constants are process-stable; user inputs are
# covered by the (id, _version) fingerprint, so an unchanged call
# reuses the previously resolved input list without signature binding.
fp = self._input_route_fingerprint(args, kwargs)
if fp == self._bind_fp:
return list(self._last_bound_inputs)
inputs = self._bind_inputs_fresh(*args, **kwargs)
self._bind_fp = fp
self._last_bound_inputs = inputs
return inputs
def __call__(self, *args: Any, **kwargs: Any) -> Any:
inputs = self._bind_inputs(*args, **kwargs)
outputs = self.graph.execute(inputs)
public_outputs = outputs[: self._public_output_count]
if self._output_spec is not None:
value, consumed = _consume_template(self._output_spec, public_outputs, 0)
if consumed != len(public_outputs):
raise RuntimeError("native graph produced an unexpected output count")
return value
if len(public_outputs) == 1:
return public_outputs[0]
return tuple(public_outputs)
class _CpuFusedPointwiseLowering(_NativeLowering):
"""Executable wrapper for Stax's vectorized CPU expression kernel."""
def __init__(
self,
graph_module: GraphModule,
graph: Any,
attribute_targets: list[str],
expected_shape: tuple[int, ...],
expected_dtype: Any,
expected_device: Any,
gradient_plan: tuple[list[int], list[float], tuple[int, ...]] | None = None,
strict_native: bool = False,
native_runner: Any = None,
native_direct: int = 0,
expected_layouts: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...]
| None = None,
) -> None:
super().__init__(graph_module, graph, attribute_targets)
self._expected_shape = expected_shape
self._expected_dtype = expected_dtype
self._expected_device = expected_device
# Per-input (shape, strides) pinned at lowering time; set only for
# broadcast/strided specializations whose generated addressing is
# valid for exactly these layouts.
self._expected_layouts = expected_layouts
self._gradient_plan = gradient_plan
self._strict_native = strict_native
self._tensorplay_codegen = "stax-fused-cpu"
self._autograd_function: Any | None = None
# Runtime-generated C kernel for the native route: straight-line
# compiler-scheduled code replacing the program interpreter when the
# system compiler is available. None keeps the graph-execute path.
self._native_runner = native_runner
# Address of the kernel's pointer-level entry (0 = absent); the C
# steady-state trampoline reads it for the direct launch path.
self._native_direct = int(native_direct) if native_direct else 0
# Route memo (id, _version, requires_grad) per input: eligibility and
# autograd routing are pure functions of these, so steady-state calls
# skip the per-input shape/dtype/device/contiguity probes entirely.
# In-place mutation bumps _version; fresh tensors have fresh ids.
self._route_fp: tuple[Any, ...] | None = None
self._route: str | None = None
_attach_fast_call(self, exec_fn=self._native_runner)
if gradient_plan is not None:
from ..autograd import Function
lowering = self
class _FusedPointwiseAutograd(Function):
@staticmethod
def forward(ctx: Any, *forward_inputs: Any) -> Any:
ctx.save_for_backward(*forward_inputs)
return lowering._execute_inputs(list(forward_inputs))
@staticmethod
def backward(ctx: Any, *grad_outputs: Any) -> tuple[Any, ...]:
grad_output = grad_outputs[0] if grad_outputs else None
if grad_output is None:
return (None,) * len(ctx.saved_tensors)
gradients = lowering._execute_backward(
ctx.saved_tensors,
grad_output,
)
return gradients
self._autograd_function = _FusedPointwiseAutograd
@staticmethod
def _eligible_inputs(
inputs: list[Any],
expected_shape: tuple[int, ...],
expected_dtype: Any,
expected_device: Any,
expected_layouts: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...]
| None = None,
) -> bool:
try:
import tensorplay
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return False
if expected_layouts is not None:
if len(inputs) != len(expected_layouts):
return False
# Broadcast/strided specializations pin each input's exact
# (shape, strides): the generated addressing was proven for
# that layout, so any deviation must re-lower.
# The expected device is captured from the region's own sample
# inputs, so comparing against it is what pins the device.
return bool(inputs) and all(
isinstance(value, tensor_type)
and value.dtype == expected_dtype
and value.device == expected_device
and tuple(int(item) for item in value.shape) == layout[0]
and tuple(int(item) for item in value.stride()) == layout[1]
for value, layout in zip(inputs, expected_layouts)
)
return bool(inputs) and all(
isinstance(value, tensor_type)
and value.dtype == expected_dtype
and value.device == expected_device
and tuple(int(item) for item in value.shape) == expected_shape
and value.is_contiguous()
for value in inputs
)
def _execute_inputs(self, inputs: list[Any]) -> Any:
if self._native_runner is not None:
return self._native_runner(inputs)
outputs = self.graph.execute(inputs)
if len(outputs) != 1:
return tuple(outputs)
return outputs[0]
def _execute_backward(
self,
inputs: tuple[Any, ...],
grad_output: Any,
) -> tuple[Any, ...]:
gradients = []
if self._gradient_plan is None:
raise RuntimeError("Stax fused pointwise backward plan is missing")
import tensorplay
grad_output = _normalize_pointwise_grad_output(grad_output, inputs[0])
program, constants, output_refs = self._gradient_plan
gradients = tensorplay._C._stax.execute_fused_pointwise_multi(
[*inputs, grad_output],
program,
constants,
output_refs,
)
return tuple(gradients)
@staticmethod
def _input_route_fingerprint(value: Any) -> Any:
import tensorplay
if isinstance(value, tensorplay.Tensor):
try:
version = value._version
except RuntimeError:
# Inference tensors are immutable: the identity alone keys
# the entry, no metadata snapshot needed.
if getattr(value, "is_inference", lambda: False)():
return ("t", id(value), None)
version = _metadata_fingerprint(value)
return (
"t",
id(value),
version,
bool(getattr(value, "requires_grad", False)),
)
return ("o", id(value))
def _resolve_route(self, inputs: list[Any]) -> str:
if not self._eligible_inputs(
inputs,
self._expected_shape,
self._expected_dtype,
self._expected_device,
getattr(self, "_expected_layouts", None),
):
return "fallback"
if self._gradient_plan is not None and any(
value.requires_grad for value in inputs
):
return "autograd"
return "native"
def __call__(self, *args: Any, **kwargs: Any) -> Any:
if not kwargs and len(args) == len(self.placeholders):
inputs = list(args)
else:
bound = self.graph_module.signature.bind_partial(*args, **kwargs)
bound.apply_defaults()
inputs = [
bound.arguments[node.target if isinstance(node.target, str) else node.name]
for node in self.placeholders
]
fp = tuple(self._input_route_fingerprint(value) for value in inputs)
if fp != self._route_fp:
self._route = self._resolve_route(inputs)
self._route_fp = fp
route = self._route
if route == "fallback":
raise RuntimeError(
"Stax fused CPU lowering received inputs outside its "
"compiled specialization"
)
if route == "autograd":
if self._autograd_function is None:
raise RuntimeError("Stax fused pointwise autograd function is missing")
return self._autograd_function.apply(*inputs)
return self._execute_inputs(inputs)
_ARITHMETIC_OPS = frozenset({"add", "sub", "mul", "div", "pow"})
_COMPARISON_OPS = frozenset({"lt", "le", "gt", "ge", "eq", "ne"})
_ORDER_OPS = frozenset({"minimum", "maximum", "clamp_min", "clamp_max"})
_UNARY_OPS = frozenset(
{
"neg",
"pos",
"abs",
"sin",
"cos",
"exp",
"log",
"sigmoid",
"sqrt",
"square",
"tanh",
"relu",
"rsqrt",
"exp2",
"erf",
}
)
_CAST_METHOD_DTYPES = {
"float": "tensorplay.float32",
"half": "tensorplay.float16",
"double": "tensorplay.float64",
}
def _build_pointwise_program(
graph_module: GraphModule,
*,
skip_node: Node | None = None,
output_override: Node | None = None,
allow_empty: bool = False,
opcodes: dict[str, int] | None = None,
nodes: list[Node] | None = None,
extra_refs: dict[Node, int] | None = None,
input_slots: int | None = None,
constants: list[float] | None = None,
) -> tuple[list[Node], list[int], list[float], list[tuple[str, int, int, int]], int] | None:
"""Encode one canonical pointwise graph as Stax's postfix program.
``skip_node`` excludes one node from the program (used by the Triton
reduction-epilogue path, which folds a full-reduction ``sum`` tail into
the kernel instead of lowering it as an op), with ``output_override``
naming the program's result node.
``nodes`` narrows the walk to one ordered slice of the region, with
``extra_refs`` naming values the slice may read but does not compute and
``input_slots`` widening the reference space those extra names live in.
A caller passing ``constants`` accumulates into it, so several slices of
one region share a single constant pool and a single reference space.
``opcodes`` selects the target opcode table: the CPU fused interpreter
supports the base table only, while the Triton code generator accepts
the extended surface (comparisons, ``where``, order relations, casts).
Values are typed numerically — comparisons yield booleans, everything
else yields floats, and the program output must be a float value (the
store path types outputs off the sample dtype).
"""
table = _TRITON_OPCODES if opcodes is None else opcodes
external_nodes = list(graph_module.graph.placeholders)
base = len(external_nodes) if input_slots is None else int(input_slots)
refs: dict[Node, int] = {
node: index for index, node in enumerate(external_nodes)
}
refs.update(extra_refs or {})
ref_types: dict[int, str] = {index: "num" for index in range(base)}
program: list[int] = []
if constants is None:
constants = []
instructions: list[tuple[str, int, int, int]] = []
temp_count = 0
def constant_ref(value: Any) -> int:
if not _is_scalar(value):
raise TypeError("Stax CPU pointwise constants must be scalar")
constants.append(float(value))
return -len(constants)
def value_ref(value: Any) -> int:
if isinstance(value, Node):
if value not in refs:
raise ValueError("pointwise graph references an unavailable value")
return refs[value]
return constant_ref(value)
def value_type(ref: int) -> str:
return ref_types.get(ref, "num")
def emit(op_name: str, lhs: int, rhs: int = -1) -> int | None:
nonlocal temp_count
code = table.get(op_name)
if code is None:
return None
program.extend((code, lhs, rhs))
result = base + temp_count
temp_count += 1
instructions.append((op_name, lhs, rhs, result))
return result
for node in graph_module.graph.nodes if nodes is None else nodes:
if node is skip_node:
continue
if node.op in {"placeholder", "output"}:
continue
if node.op not in {"call_function", "call_method"}:
return None
op_name = _target_name(node.target)
if op_name not in _CPU_FUSED_OPS:
return None
kwargs = node.kwargs or {}
result_type = "num"
if op_name in {"add", "sub"} and (
len(node.args) == 3
or ("alpha" in (node.kwargs or {}) and len(node.args) == 2)
):
if len(node.args) == 3 and "alpha" not in kwargs:
lhs, rhs, alpha = node.args
else:
if len(node.args) != 2:
return None
lhs, rhs = node.args
alpha = kwargs.get("alpha", 1)
if not _is_scalar(alpha):
return None
lhs_ref = value_ref(lhs)
rhs_ref = value_ref(rhs)
if alpha != 1:
scaled = emit("mul", rhs_ref, constant_ref(alpha))
if scaled is None:
return None
ref_types[scaled] = "num"
rhs_ref = scaled
node_ref = emit(op_name, lhs_ref, rhs_ref)
if node_ref is None:
return None
refs[node] = node_ref
ref_types[node_ref] = "num"
continue
if op_name in _ARITHMETIC_OPS or op_name in _COMPARISON_OPS or (
op_name in _ORDER_OPS
):
if kwargs or len(node.args) != 2:
return None
node_ref = emit(
op_name, value_ref(node.args[0]), value_ref(node.args[1])
)
if node_ref is None:
return None
if op_name in _COMPARISON_OPS:
result_type = "bool"
refs[node] = node_ref
ref_types[node_ref] = result_type
continue
if op_name in _UNARY_OPS:
if kwargs or len(node.args) != 1:
return None
node_ref = emit(op_name, value_ref(node.args[0]))
if node_ref is None:
return None
refs[node] = node_ref
ref_types[node_ref] = "num"
continue
if op_name == "where":
if kwargs or len(node.args) != 3:
return None
cond_ref = value_ref(node.args[0])
a_ref = value_ref(node.args[1])
b_ref = value_ref(node.args[2])
# v1 contract: a boolean condition selects between float values.
# Numeric conditions and boolean branches stay uncompiled.
if value_type(cond_ref) != "bool":
return None
if value_type(a_ref) != "num" or value_type(b_ref) != "num":
return None
then_ref = emit("where", cond_ref, a_ref)
if then_ref is None:
return None
ref_types[then_ref] = "num"
node_ref = emit("where_rest", cond_ref, b_ref)
if node_ref is None:
return None
refs[node] = node_ref
ref_types[node_ref] = "num"
continue
if op_name in _CAST_METHOD_DTYPES or op_name == "to":
if op_name == "to":
if set(kwargs) - {"dtype"} or len(node.args) > 2:
return None
dtype_value = (
node.args[1] if len(node.args) > 1 else kwargs.get("dtype")
)
if dtype_value is None:
return None
dtype_key = str(dtype_value)
else:
if kwargs or len(node.args) != 1:
return None
dtype_key = _CAST_METHOD_DTYPES[op_name]
dtype_id = _CAST_DTYPE_IDS.get(dtype_key)
if dtype_id is None:
return None
node_ref = emit("cast", value_ref(node.args[0]), dtype_id)
if node_ref is None:
return None
refs[node] = node_ref
ref_types[node_ref] = "num"
continue
return None
output_values = (
[output_override]
if output_override is not None
else [
value
for output in graph_module.graph.outputs
for value in _nodes(output.args)
]
)
if (not program and not allow_empty) or len(output_values) != 1 or (
output_values[0] not in refs
):
return None
if value_type(refs[output_values[0]]) != "num":
# A boolean program output would need a typed store path.
return None
return external_nodes, program, constants, instructions, refs[output_values[0]]
def _broadcast_shape(shapes: tuple[tuple[int, ...], ...]) -> tuple[int, ...] | None:
"""Broadcast several shapes to one result shape (``None`` on mismatch)."""
rank = max((len(s) for s in shapes), default=0)
result: list[int] = []
for dim in range(rank):
extent = 1
for shape in shapes:
idx = dim - (rank - len(shape))
d = shape[idx] if idx >= 0 else 1
if d != 1:
if extent != 1 and extent != d:
return None
extent = d
result.append(extent)
return tuple(result)
def _lower_cpu_fused_pointwise(
graph_module: GraphModule,
example_inputs: list[Any],
*,
strict_native: bool = False,
dynamic: bool = False,
) -> _CpuFusedPointwiseLowering | None:
"""Build one CPU expression program for a pointwise graph.
The specialization requires matching contiguous float32 CPU tensors. For
grad-enabled pointwise graphs, Stax also emits a vectorized reverse-mode
program and attaches it through TensorPlay's Function contract. General
broadcasting, views, and unsupported derivatives stay on the native p10
path.
Two program surfaces are attempted in order: the base opcode table,
which every execution route (compiled kernel, program interpreter,
fused backward) can run, and the extended surface (comparisons,
``where``, order relations, casts), which only the runtime-generated C
kernel can execute — extended programs therefore require a successful
native build and a grad-free graph.
"""
if dynamic:
return None
try:
import tensorplay
native_module = getattr(tensorplay._C, "_stax", None)
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return None
if native_module is None or not hasattr(native_module.Graph, "execute"):
return None
if not example_inputs or any(not isinstance(value, tensor_type) for value in example_inputs):
return None
first = example_inputs[0]
if (
not first.device.is_cpu()
or first.dtype != tensorplay.float32
or not first.is_contiguous()
):
return None
if any(
value.device != first.device or value.dtype != first.dtype
for value in example_inputs[1:]
):
return None
# Broadcast/strided acceptance: inputs may differ in shape or layout as
# long as the emitter can prove every address it generates contiguous
# within a vector. Everything else keeps the generic fallback.
input_shapes = tuple(
tuple(int(item) for item in value.shape) for value in example_inputs
)
input_strides = tuple(
tuple(int(item) for item in value.stride()) for value in example_inputs
)
output_shape = input_shapes[0]
if _broadcast_shape(input_shapes) != output_shape:
return None
try:
from .codegen.cpp import analyze_input_modes
input_modes = analyze_input_modes(
input_shapes, input_strides, output_shape, lane_count=16
)
except (TypeError, ValueError):
return None
if input_modes is None:
return None
# Legacy surface (every input flat) keeps the program-interpreter
# fallback; anything else requires the compiled kernel, whose generated
# addressing is only valid for these exact layouts, and a grad-free
# graph (the fused backward program assumes flat inputs).
layouts_only = any(mode != "flat" for mode, _ in input_modes)
if layouts_only and any(
value.requires_grad for value in example_inputs
):
return None
extended = False
try:
pointwise = _build_pointwise_program(
graph_module, opcodes=_CPU_FUSED_OPCODES
)
if pointwise is None:
pointwise = _build_pointwise_program(
graph_module, opcodes=_TRITON_OPCODES
)
extended = pointwise is not None
if extended and any(
value.requires_grad for value in example_inputs
):
return None
except (TypeError, ValueError, RuntimeError):
return None
if pointwise is None:
return None
external_nodes, program, constants, instructions, output_ref = pointwise
if len(external_nodes) != len(example_inputs):
return None
native_runner: Any = None
native_direct = 0
try:
from .codegen.cpp import build_cpu_native_kernel
built = build_cpu_native_kernel(
instructions,
constants,
len(external_nodes),
output_ref,
shape=first.shape,
device=first.device,
input_shapes=input_shapes,
input_strides=input_strides,
)
if isinstance(built, tuple):
native_runner, native_direct = built
else:
native_runner = built
except Exception:
native_runner = None
native_direct = 0
if extended and native_runner is None:
return None
# Broadcast/strided layouts have no interpreter-compatible program: the
# compiled kernel is the only route that can address them.
if layouts_only and native_runner is None:
return None
try:
graph = native_module.Graph()
native_values: dict[Node, Any] = {
node: graph.add_input() for node in external_nodes
}
output_values = [
value for output in graph_module.graph.outputs for value in _nodes(output.args)
]
if len(output_values) != 1:
return None
fused = graph.create_node("fused_pointwise", output_values[0].name)
for node in external_nodes:
fused.add_input(native_values[node])
fused.set_int_attr("input_count", len(external_nodes))
fused.set_ints_attr("program", program)
fused.set_floats_attr("constants", constants)
graph.register_output(fused.add_output())
except (TypeError, ValueError, RuntimeError):
return None
gradient_plan: tuple[list[int], list[float], tuple[int, ...]] | None = None
if not extended and not layouts_only and any(
value.requires_grad for value in example_inputs
):
if any(op_name not in _CPU_FUSED_AUTOGRAD_OPS for op_name, *_ in instructions):
return None
try:
gradient_plan = _build_fused_gradient_graphs(
len(external_nodes),
instructions,
program,
constants,
len(program) // 3,
output_ref,
)
except (TypeError, ValueError, RuntimeError):
return None
if gradient_plan is None:
return None
return _CpuFusedPointwiseLowering(
graph_module,
graph,
[],
tuple(int(item) for item in first.shape),
first.dtype,
first.device,
gradient_plan,
strict_native,
native_runner,
native_direct,
expected_layouts=(
tuple(zip(input_shapes, input_strides)) if layouts_only else None
),
)
# Reduction spellings the fused CPU reduction path recognizes. ``max``/``min``
# are accepted only in their whole-tensor form: with a dimension they return a
# value/index pair, which is a different lowering contract.
_REDUCTION_METHODS = frozenset(
{"sum", "mean", "prod", "max", "min", "amax", "amin"}
)
_DIMLESS_ONLY_REDUCTIONS = frozenset({"max", "min"})
_DIM_ONLY_REDUCTIONS = frozenset({"amax", "amin"})
def _reduction_dtype_ok(value: Any) -> bool:
"""Whether a reduction's ``dtype`` argument keeps the float32 contract."""
if value is None:
return True
import tensorplay
if value is tensorplay.float32:
return True
# The captured call carries the sentinel that means "keep the input
# dtype"; the route already pinned float32 inputs.
return str(value).rsplit(".", 1)[-1].lower() == "undefined"
def _parse_reduction(node: Node, rank: int) -> Any:
"""Read one reduction node into a :class:`ReduceSpec`, or ``None``.
Positional and keyword spellings both resolve here: the trailing
positional arguments of a reduction are ``dim``, ``keepdim``, ``dtype``
in that order.
"""
from .codegen.cpp_reduction import ReduceSpec
if node.op not in {"call_function", "call_method"}:
return None
name = _target_name(node.target)
if name not in _REDUCTION_METHODS:
return None
args = list(node.args)
if not args or not isinstance(args[0], Node):
return None
rest = args[1:]
kwargs = dict(node.kwargs or {})
if len(rest) > 3:
return None
dim: Any = None
keepdim: Any = False
if len(rest) >= 1:
dim = rest[0]
if len(rest) >= 2:
keepdim = rest[1]
if len(rest) >= 3 and not _reduction_dtype_ok(rest[2]):
return None
if "dim" in kwargs:
if dim is not None:
return None
dim = kwargs.pop("dim")
if "keepdim" in kwargs:
if len(rest) >= 2:
return None
keepdim = kwargs.pop("keepdim")
if "dtype" in kwargs and not _reduction_dtype_ok(kwargs.pop("dtype")):
return None
if kwargs:
return None
if not isinstance(keepdim, bool):
return None
if dim is None:
if name in _DIM_ONLY_REDUCTIONS:
return None
return ReduceSpec(name, tuple(range(rank)), False)
if name in _DIMLESS_ONLY_REDUCTIONS:
return None
if isinstance(dim, bool):
return None
if isinstance(dim, int):
dims: tuple[int, ...] = (int(dim),)
elif isinstance(dim, (tuple, list)) and dim and all(
isinstance(item, int) and not isinstance(item, bool) for item in dim
):
dims = tuple(int(item) for item in dim)
else:
return None
return ReduceSpec(name, dims, keepdim)
class _CpuFusedReductionLowering:
"""Executable wrapper for Stax's fused CPU reduction kernel.
The kernel owns the whole region: it evaluates the pointwise expression
and the reduction in one pass, allocates its own output, and returns the
wrapped tensor, so the steady-state call never builds an intermediate.
"""
def __init__(
self,
graph_module: GraphModule,
expected_dtype: Any,
expected_device: Any,
expected_layouts: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...],
native_runner: Any,
native_direct: int,
out_shape: tuple[int, ...],
strict_native: bool = False,
) -> None:
self.graph_module = graph_module
# No interpreter-executable graph backs this region: the compiled
# kernel is the only route, and the route check below guarantees the
# inputs it was specialized for.
self.graph = None
self.placeholders = graph_module.graph.placeholders
self.attribute_targets: list[str] = []
self.constant_values: list[Any] = []
self.native_values: dict[Node, Any] = {}
self._output_count = 1
self._public_output_count = 1
self._output_spec = None
self._tensorplay_codegen = "stax-fused-cpu-reduce"
self._expected_dtype = expected_dtype
self._expected_device = expected_device
self._expected_layouts = expected_layouts
self._out_shape = out_shape
self._strict_native = strict_native
self._native_runner = native_runner
self._native_direct = int(native_direct) if native_direct else 0
self._route_fp: tuple[Any, ...] | None = None
self._route: str | None = None
_attach_fast_call(self, exec_fn=native_runner)
def _resolve_route(self, inputs: list[Any]) -> str:
if not _CpuFusedPointwiseLowering._eligible_inputs(
inputs,
(),
self._expected_dtype,
self._expected_device,
self._expected_layouts,
):
return "fallback"
if any(getattr(value, "requires_grad", False) for value in inputs):
# Reverse mode over a fused reduction is not part of this
# lowering's contract; a grad-carrying call re-lowers.
return "fallback"
return "native"
def __call__(self, *args: Any, **kwargs: Any) -> Any:
if not kwargs and len(args) == len(self.placeholders):
inputs = list(args)
else:
bound = self.graph_module.signature.bind_partial(*args, **kwargs)
bound.apply_defaults()
inputs = [
bound.arguments[
node.target if isinstance(node.target, str) else node.name
]
for node in self.placeholders
]
fp = tuple(
_CpuFusedPointwiseLowering._input_route_fingerprint(value)
for value in inputs
)
if fp != self._route_fp:
self._route = self._resolve_route(inputs)
self._route_fp = fp
if self._route == "fallback":
raise RuntimeError(
"Stax fused CPU reduction received inputs outside its "
"compiled specialization"
)
return self._native_runner(inputs)
def _lower_cpu_fused_reduction(
graph_module: GraphModule,
example_inputs: list[Any],
*,
strict_native: bool = False,
dynamic: bool = False,
) -> _CpuFusedReductionLowering | None:
"""Build one fused CPU kernel for a pointwise region ending in a reduction.
The region's single output must be a reduction whose operand is used
nowhere else, so folding it into the reduction loop cannot change what any
other node observes. Everything upstream of the reduction is encoded as
the same expression program the pointwise path uses.
"""
if dynamic:
return None
try:
import tensorplay
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return None
if not example_inputs or any(
not isinstance(value, tensor_type) for value in example_inputs
):
return None
first = example_inputs[0]
if (
not first.device.is_cpu()
or first.dtype != tensorplay.float32
or not first.is_contiguous()
):
return None
if any(
value.device != first.device or value.dtype != first.dtype
for value in example_inputs[1:]
):
return None
if any(value.requires_grad for value in example_inputs):
return None
input_shapes = tuple(
tuple(int(item) for item in value.shape) for value in example_inputs
)
input_strides = tuple(
tuple(int(item) for item in value.stride()) for value in example_inputs
)
in_shape = _broadcast_shape(input_shapes)
if in_shape is None or not in_shape:
return None
output_values = [
value
for output in graph_module.graph.outputs
for value in _nodes(output.args)
]
if len(output_values) != 1 or not isinstance(output_values[0], Node):
return None
reduce_node = output_values[0]
spec = _parse_reduction(reduce_node, len(in_shape))
if spec is None:
return None
source = reduce_node.args[0]
if not isinstance(source, Node) or len(source.users) != 1:
return None
try:
pointwise = _build_pointwise_program(
graph_module,
skip_node=reduce_node,
output_override=source,
allow_empty=True,
opcodes=_TRITON_OPCODES,
)
except (TypeError, ValueError, RuntimeError):
return None
if pointwise is None:
return None
external_nodes, _program, constants, instructions, output_ref = pointwise
if len(external_nodes) != len(example_inputs):
return None
try:
from .codegen.cpp_reduction import build_cpu_reduction_kernel
built = build_cpu_reduction_kernel(
instructions,
constants,
len(external_nodes),
output_ref,
spec,
in_shape=in_shape,
device=first.device,
input_shapes=input_shapes,
input_strides=input_strides,
)
except Exception:
built = None
if built is None:
return None
runner, direct, out_shape = built
return _CpuFusedReductionLowering(
graph_module,
first.dtype,
first.device,
tuple(zip(input_shapes, input_strides)),
runner,
direct,
out_shape,
strict_native,
)
def _elem_dependencies(
target: Node, kinds: dict[Node, str], stop: set[Node] | None = None
) -> list[Node]:
"""Order the elementwise nodes one value depends on, producers first.
Placeholders, row values, and anything in ``stop`` are read through
references rather than recomputed, so they end the walk. A node reached
from two stages and not staged appears in both of their orders: that
stage re-evaluates it while the row is still in cache, which is cheaper
than the memory traffic of materializing it.
"""
if target.op == "placeholder" or (stop is not None and target in stop):
return []
order: list[Node] = []
seen: set[Node] = set()
stack: list[tuple[Node, bool]] = [(target, False)]
while stack:
node, expanded = stack.pop()
if expanded:
order.append(node)
continue
if node in seen:
continue
seen.add(node)
stack.append((node, True))
operands = list(_nodes(node.args)) + list(_nodes(node.kwargs or {}))
for operand in operands:
if stop is not None and operand in stop:
continue
if kinds.get(operand) == "elem" and operand.op != "placeholder":
if operand not in seen:
stack.append((operand, False))
return order
def _plan_row_fusion(
graph_module: GraphModule,
in_shape: tuple[int, ...],
input_shapes: tuple[tuple[int, ...], ...],
*,
stage: bool = True,
) -> Any:
"""Split a region into row stages, or return ``None``.
A node is *elementwise* when it carries one value per input element and
*row-valued* when it carries one value per row: reductions over the
trailing axis turn the former into the latter, and every later stage may
read row values as broadcasts. The region qualifies when at least one
such reduction exists and every node lands in one of the two classes.
"""
from .codegen.cpp_rowfusion import ROW_OPS, RowFusion, RowStep, _ROW_UNARY
rank = len(in_shape)
if rank < 2 or any(int(extent) <= 0 for extent in in_shape):
return None
extent = int(in_shape[-1])
rows = 1
for size in in_shape[:-1]:
rows *= int(size)
placeholders = list(graph_module.graph.placeholders)
if not placeholders or len(placeholders) != len(input_shapes):
return None
# Every input has to span the reduced axis: a value that is constant along
# it would be row-valued, and the classification below assumes it is not.
for shape in input_shapes:
if len(shape) > rank:
return None
aligned = (1,) * (rank - len(shape)) + tuple(int(size) for size in shape)
if aligned[-1] != extent:
return None
output_values = [
value
for output in graph_module.graph.outputs
for value in _nodes(output.args)
]
if len(output_values) != 1 or not isinstance(output_values[0], Node):
return None
target = output_values[0]
if target.op == "placeholder":
return None
kinds: dict[Node, str] = {node: "elem" for node in placeholders}
row_shapes: dict[Node, tuple[int, ...]] = {}
slots: dict[Node, int] = {}
staged: list[tuple[str, Node, Any]] = []
for node in graph_module.graph.nodes:
if node.op in {"placeholder", "output"}:
continue
if node.op not in {"call_function", "call_method"}:
return None
spec = _parse_reduction(node, rank)
if spec is not None:
normalized = spec.normalized(rank)
if normalized is None or normalized.dims != (rank - 1,):
return None
source = node.args[0]
if not isinstance(source, Node) or kinds.get(source) != "elem":
return None
kinds[node] = "row"
slots[node] = len(slots)
row_shapes[node] = (
tuple(in_shape[:-1]) + (1,)
if normalized.keepdim
else tuple(in_shape[:-1])
)
staged.append(("reduce", node, normalized))
continue
operands = list(_nodes(node.args)) + list(_nodes(node.kwargs or {}))
if any(operand not in kinds for operand in operands):
return None
if not operands or any(kinds[operand] == "elem" for operand in operands):
kinds[node] = "elem"
continue
name = _target_name(node.target)
if name not in ROW_OPS or (node.kwargs or {}):
return None
arity = 1 if name in _ROW_UNARY else 2
if len(node.args) != arity:
return None
if any(
not isinstance(arg, Node) and not _is_scalar(arg)
for arg in node.args
):
return None
shapes = {
row_shapes[arg] for arg in node.args if isinstance(arg, Node)
}
if len(shapes) != 1:
return None
kinds[node] = "row"
slots[node] = len(slots)
row_shapes[node] = next(iter(shapes))
staged.append(("rowop", node, name))
if not slots or not any(entry[0] == "reduce" for entry in staged):
return None
# A row value an elementwise stage reads has to broadcast along the
# reduced axis, which is what the kept trailing dimension expresses.
keep = tuple(in_shape[:-1]) + (1,)
for node in slots:
if any(kinds.get(user) == "elem" for user in node.users):
if row_shapes[node] != keep:
return None
input_count = len(placeholders)
extra_refs = {node: input_count + slot for node, slot in slots.items()}
# A value a reduction pass computes and a later pass needs again is worth
# keeping: the pass already holds it in a register, so staging it costs
# one store and saves the later pass the whole expression behind it.
reduce_sources = [
node.args[0] for kind, node, _payload in staged if kind == "reduce"
]
later: list[set[Node]] = []
seen_later: set[Node] = set()
for source in reversed(reduce_sources[1:]):
seen_later |= set(_elem_dependencies(source, kinds))
later.append(set(seen_later))
later.reverse()
if kinds[target] == "elem":
output_deps = set(_elem_dependencies(target, kinds))
else:
output_deps = set()
stages: dict[Node, int] = {}
for index, source in enumerate(reduce_sources) if stage else ():
if source.op == "placeholder" or source in stages:
continue
reused = output_deps | (later[index] if index < len(later) else set())
if source in reused:
stages[source] = len(stages)
total_inputs = input_count + len(slots) + len(stages)
stage_refs = {
node: input_count + len(slots) + slot for node, slot in stages.items()
}
constants: list[float] = []
def elem_program(node: Node, available: dict[Node, int]) -> Any:
try:
return _build_pointwise_program(
graph_module,
output_override=node,
allow_empty=True,
opcodes=_TRITON_OPCODES,
nodes=_elem_dependencies(node, kinds, stop=set(available)),
extra_refs={**extra_refs, **available},
input_slots=total_inputs,
constants=constants,
)
except (TypeError, ValueError, RuntimeError):
return None
def row_operand(value: Any) -> int | None:
if isinstance(value, Node):
return extra_refs.get(value)
if not _is_scalar(value):
return None
constants.append(float(value))
return -len(constants)
steps: list[Any] = []
available: dict[Node, int] = {}
for entry_kind, node, payload in staged:
if entry_kind == "reduce":
source = node.args[0]
built = elem_program(source, available)
if built is None:
return None
instructions, output_ref = built[3], built[4]
stage = stage_refs.get(source)
steps.append(
RowStep(
kind="reduce",
slot=slots[node],
op=payload.op,
instructions=tuple(instructions),
output_ref=output_ref,
stage=-1 if stage is None else stage - input_count - len(slots),
)
)
if stage is not None:
available[source] = stage
continue
lhs = row_operand(node.args[0])
rhs = -1 if payload in _ROW_UNARY else row_operand(node.args[1])
if lhs is None or rhs is None:
return None
steps.append(
RowStep(kind="rowop", slot=slots[node], op=payload, lhs=lhs, rhs=rhs)
)
if kinds[target] == "elem":
built = elem_program(target, available)
if built is None:
return None
out_instructions, out_ref = built[3], built[4]
if not out_instructions:
return None
output_kind = "elem"
out_shape = tuple(int(size) for size in in_shape)
else:
output_kind = "row"
out_instructions = []
out_ref = extra_refs[target]
out_shape = row_shapes[target]
return RowFusion(
input_count=input_count,
row_slots=len(slots),
stage_slots=len(stages),
constants=tuple(constants),
steps=tuple(steps),
output_kind=output_kind,
out_instructions=tuple(out_instructions),
out_ref=out_ref,
reduce_extent=extent,
rows=rows,
in_shape=tuple(int(size) for size in in_shape),
out_shape=tuple(int(size) for size in out_shape),
)
class _CpuRowFusionLowering(_CpuFusedReductionLowering):
"""Executable wrapper for a row-staged CPU kernel.
The route contract is the reduction kernel's: one compiled specialization
that owns the whole region, allocates its output, and declines anything
outside the layouts it was built for.
"""
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
self._tensorplay_codegen = "stax-fused-cpu-rowfuse"
def __call__(self, *args: Any, **kwargs: Any) -> Any:
try:
return super().__call__(*args, **kwargs)
except RuntimeError as error:
if "fused CPU reduction" in str(error):
raise RuntimeError(
"Stax row-staged CPU kernel received inputs outside its "
"compiled specialization"
) from None
raise
def _expand_row_normalizations(graph_module: GraphModule) -> GraphModule | None:
"""Rewrite the softmax family into primitives on a copy of the region.
The composites this expands are single fused kernels of their own, so the
expansion is only worth having when it is fused back into one kernel.
Working on a copy is what makes that conditional: a region the row-staged
planner then declines keeps the operators -- and the kernels -- it had.
"""
from ..graph.passes import DecomposeRowNormalizations, row_normalization_names
known = row_normalization_names()
present = False
for node in graph_module.graph.nodes:
if node.op == "call_method":
name = node.target if isinstance(node.target, str) else None
elif node.op == "call_function":
name = getattr(node.target, "__name__", None)
else:
continue
if name in known:
present = True
break
if not present:
return None
try:
clone = _copy_region_graph(graph_module.graph)
expanded = GraphModule(
graph_module.root, clone, graph_module.signature
)
result = DecomposeRowNormalizations()(expanded)
except Exception:
return None
return result.graph_module if result.modified else None
def _copy_region_graph(graph: Any) -> Any:
"""Duplicate a graph's nodes without duplicating what they carry.
Node metadata holds the traced tensor values, so a deep copy of the
region would clone every intermediate; node-level copies keep those
references shared, which is all a rewrite needs.
"""
from ..graph import Graph, map_arg
clone = Graph()
mapping: dict[Node, Node] = {}
for node in graph.nodes:
if node.op == "output":
continue
mapping[node] = clone.node_copy(node, lambda value: mapping[value])
for node in graph.nodes:
if node.op == "output":
clone.output(
map_arg(node.args[0], lambda value: mapping[value]), node.type
)
return clone
def _row_fusion_plan(
graph_module: GraphModule,
example_inputs: list[Any],
on_device: str,
) -> Any:
"""Guard a region and plan it as row stages; ``None`` when it is not one.
The plan itself carries no device: the same stages compile to a CPU loop
nest or to one CUDA program per row, so both lowerings share this front
end and differ only in the generator they hand the plan to.
"""
try:
import tensorplay
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return None
if not example_inputs or any(
not isinstance(value, tensor_type) for value in example_inputs
):
return None
first = example_inputs[0]
on = first.device.is_cuda() if on_device == "cuda" else first.device.is_cpu()
if not on or first.dtype != tensorplay.float32 or not first.is_contiguous():
return None
if any(
value.device != first.device
or value.dtype != first.dtype
or not value.is_contiguous()
for value in example_inputs[1:]
):
return None
if any(value.requires_grad for value in example_inputs):
return None
input_shapes = tuple(
tuple(int(item) for item in value.shape) for value in example_inputs
)
input_strides = tuple(
tuple(int(item) for item in value.stride()) for value in example_inputs
)
in_shape = _broadcast_shape(input_shapes)
if in_shape is None or not in_shape:
return None
# Keeping a value alive across a reduction is a win where the buffer
# sits in cache and a loss where it sits in a register file: one device
# stages, the other recomputes.
stage = on_device != "cuda"
module = graph_module
fusion = _plan_row_fusion(module, in_shape, input_shapes, stage=stage)
if fusion is None:
expanded = _expand_row_normalizations(graph_module)
if expanded is None:
return None
fusion = _plan_row_fusion(
expanded, in_shape, input_shapes, stage=stage
)
if fusion is None:
return None
module = expanded
return module, fusion, input_shapes, input_strides, first
class _CudaRowFusionLowering(_CpuFusedReductionLowering):
"""Executable wrapper for a row-staged CUDA kernel.
One program per output row, the row resident for the whole region: the
inputs are read once no matter how many reductions the region contains.
"""
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
self._tensorplay_codegen = "stax-fused-cuda-rowfuse"
def __call__(self, *args: Any, **kwargs: Any) -> Any:
try:
return super().__call__(*args, **kwargs)
except RuntimeError as error:
if "fused CPU reduction" in str(error):
raise RuntimeError(
"Stax row-staged CUDA kernel received inputs outside its "
"compiled specialization"
) from None
raise
def _lower_cuda_row_fusion(
graph_module: GraphModule,
example_inputs: list[Any],
*,
strict_native: bool = False,
dynamic: bool = False,
) -> Any:
"""Build one CUDA kernel for a region with reductions in the middle.
Splitting such a region at every reduction gives one kernel per stage,
each streaming the input again and writing an intermediate the next one
reads back. Keeping the row resident across the stages removes all of
that traffic and leaves one launch.
"""
if dynamic:
return None
planned = _row_fusion_plan(graph_module, example_inputs, "cuda")
if planned is None:
return None
module, fusion, input_shapes, input_strides, first = planned
try:
from .codegen.triton_rowfusion import build_cuda_row_fusion_kernel
launch = build_cuda_row_fusion_kernel(
fusion,
input_shapes=input_shapes,
input_strides=input_strides,
)
except Exception:
launch = None
if launch is None:
return None
return _CudaRowFusionLowering(
module,
first.dtype,
first.device,
tuple(zip(input_shapes, input_strides)),
launch,
0,
fusion.out_shape,
strict_native,
)
def _lower_cpu_row_fusion(
graph_module: GraphModule,
example_inputs: list[Any],
*,
strict_native: bool = False,
dynamic: bool = False,
) -> Any:
"""Build one CPU kernel for a region with reductions in the middle.
The reduction results feed elementwise work over the same axis they were
reduced along, so the region cannot be expressed as a pointwise program
or as a pointwise program ending in a reduction. Staging it per row keeps
every intermediate in cache instead of writing it out and reading it back.
"""
if dynamic:
return None
planned = _row_fusion_plan(graph_module, example_inputs, "cpu")
if planned is None:
return None
module, fusion, input_shapes, input_strides, first = planned
try:
from .codegen.cpp_rowfusion import build_cpu_row_fusion_kernel
built = build_cpu_row_fusion_kernel(
fusion,
device=first.device,
input_shapes=input_shapes,
input_strides=input_strides,
)
except Exception:
built = None
if built is None:
return None
runner, direct = built
return _CpuRowFusionLowering(
module,
first.dtype,
first.device,
tuple(zip(input_shapes, input_strides)),
runner,
direct,
fusion.out_shape,
strict_native,
)
# ---------------------------------------------------------------------------
# Mixed regions: generated kernels between operators that run as they are
def _traced_value(graph_module: GraphModule, node: Node) -> Any:
"""The tensor a node carried when the region was captured, or ``None``."""
if node.op == "get_attr":
try:
return graph_module._get_attr(node.target)
except (AttributeError, KeyError, RuntimeError):
return None
return node.meta.get("val")
def _tensor_layout(value: Any) -> tuple[tuple[int, ...], tuple[int, ...]] | None:
try:
shape = tuple(int(item) for item in value.shape)
stride = tuple(int(item) for item in value.stride())
except (AttributeError, TypeError, ValueError):
return None
return shape, stride
def _segment_externals(nodes: tuple[Node, ...]) -> list[Node]:
"""Values a segment reads from outside itself, in first-use order."""
inside = set(nodes)
externals: list[Node] = []
seen: set[Node] = set()
for node in nodes:
operands = list(_nodes(node.args)) + list(_nodes(node.kwargs or {}))
for operand in operands:
if operand in inside or operand in seen:
continue
seen.add(operand)
externals.append(operand)
return externals
def _build_segment_kernel(
graph_module: GraphModule,
segment: Any,
device: Any,
) -> tuple[list[Node], Any] | None:
"""Compile one fusible segment; return its inputs and callable runner."""
externals = _segment_externals(segment.nodes)
if not externals or len(externals) > 16:
return None
layouts = []
for node in externals:
layout = _tensor_layout(_traced_value(graph_module, node))
if layout is None:
return None
layouts.append(layout)
input_shapes = tuple(shape for shape, _stride in layouts)
input_strides = tuple(stride for _shape, stride in layouts)
refs = {node: index for index, node in enumerate(externals)}
reduce_node = segment.tail if segment.kind == "pw+red" else None
body = [node for node in segment.nodes if node is not reduce_node]
source = segment.producer if reduce_node is not None else segment.nodes[-1]
if source is None:
return None
try:
program = _build_pointwise_program(
graph_module,
output_override=source,
allow_empty=reduce_node is not None,
opcodes=_TRITON_OPCODES,
nodes=body,
extra_refs=refs,
input_slots=len(externals),
)
except (TypeError, ValueError, RuntimeError):
return None
if program is None:
return None
_external, _encoded, constants, instructions, output_ref = program
if reduce_node is None:
layout = _tensor_layout(_traced_value(graph_module, source))
if layout is None:
return None
out_shape = layout[0]
try:
from .codegen.cpp import build_cpu_native_kernel
built = build_cpu_native_kernel(
instructions,
constants,
len(externals),
output_ref,
shape=out_shape,
device=device,
input_shapes=input_shapes,
input_strides=input_strides,
)
except Exception:
return None
if built is None:
return None
runner = built[0] if isinstance(built, tuple) else built
return externals, runner
layout = _tensor_layout(_traced_value(graph_module, source))
if layout is None:
return None
try:
from .codegen.cpp_reduction import build_cpu_reduction_kernel
built = build_cpu_reduction_kernel(
instructions,
constants,
len(externals),
output_ref,
segment.reduction,
in_shape=layout[0],
device=device,
input_shapes=input_shapes,
input_strides=input_strides,
)
except Exception:
return None
if built is None:
return None
return externals, built[0]
def _kernel_step(runner: Any, sources: tuple[int, ...]):
"""Close over one generated kernel and the slots holding its inputs."""
def run(values: list[Any]) -> Any:
return runner([values[slot] for slot in sources])
return run
class _CpuSegmentedLowering:
"""Runs a mixed region: generated kernels between untouched operators.
A region that mixes fusible work with operators the generators do not
cover used to lose the compiled route entirely. Here the fusible runs
become one kernel each and the rest of the region runs exactly as it
was captured, so a pointwise chain between two matrix products costs a
single pass over its data instead of one pass per operator.
"""
def __init__(
self,
graph_module: GraphModule,
steps: list[tuple],
slot_count: int,
constants: dict[int, Any],
output_slot: int,
expected_dtype: Any,
expected_device: Any,
expected_layouts: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...],
strict_native: bool = False,
) -> None:
self.graph_module = graph_module
self.graph = None
self.placeholders = graph_module.graph.placeholders
self.attribute_targets: list[str] = []
self.constant_values: list[Any] = []
self.native_values: dict[Node, Any] = {}
self._output_count = 1
self._public_output_count = 1
self._output_spec = None
self._tensorplay_codegen = "stax-fused-cpu-segments"
self._steps = steps
self._template: list[Any] = [None] * slot_count
for slot, value in constants.items():
self._template[slot] = value
self._output_slot = output_slot
self._expected_dtype = expected_dtype
self._expected_device = expected_device
self._expected_layouts = expected_layouts
self._strict_native = strict_native
self._route_fp: tuple[Any, ...] | None = None
self._route: str | None = None
_attach_fast_call(self, exec_fn=self._execute)
def _execute(self, inputs: list[Any]) -> Any:
# The captured constants never move, so the value table starts as a
# copy of a template that already holds them.
values = self._template.copy()
values[: len(inputs)] = inputs
for step, target, release in self._steps:
values[target] = step(values)
# Dropping an intermediate as soon as its last reader has run
# hands the buffer straight back to the allocator, so the next
# operator writes into memory that is still warm.
for slot in release:
values[slot] = None
return values[self._output_slot]
def _resolve_route(self, inputs: list[Any]) -> str:
if not _CpuFusedPointwiseLowering._eligible_inputs(
inputs,
(),
self._expected_dtype,
self._expected_device,
self._expected_layouts,
):
return "fallback"
if any(getattr(value, "requires_grad", False) for value in inputs):
return "fallback"
return "native"
def __call__(self, *args: Any, **kwargs: Any) -> Any:
if not kwargs and len(args) == len(self.placeholders):
inputs = list(args)
else:
bound = self.graph_module.signature.bind_partial(*args, **kwargs)
bound.apply_defaults()
inputs = [
bound.arguments[
node.target if isinstance(node.target, str) else node.name
]
for node in self.placeholders
]
fp = tuple(
_CpuFusedPointwiseLowering._input_route_fingerprint(value)
for value in inputs
)
if fp != self._route_fp:
self._route = self._resolve_route(inputs)
self._route_fp = fp
if self._route == "fallback":
raise RuntimeError(
"Stax segmented CPU region received inputs outside its "
"compiled specialization"
)
return self._execute(inputs)
def _lower_cpu_segmented(
graph_module: GraphModule,
example_inputs: list[Any],
*,
strict_native: bool = False,
dynamic: bool = False,
) -> _CpuSegmentedLowering | None:
"""Compile the fusible runs of a region the whole-region paths declined.
The scheduler partitions the region; every pointwise run and every run
ending in a reduction becomes one generated kernel, and each remaining
operator stays a single call between them. A store-time epilogue is the
one schedule this path cannot wire yet, and a region containing one keeps
its existing route rather than losing the fusion the epilogue expresses.
"""
if dynamic:
return None
try:
import tensorplay
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return None
if not example_inputs or any(
not isinstance(value, tensor_type) for value in example_inputs
):
return None
first = example_inputs[0]
if not first.device.is_cpu() or first.dtype != tensorplay.float32:
return None
if any(
value.device != first.device or value.dtype != first.dtype
for value in example_inputs[1:]
):
return None
if any(value.requires_grad for value in example_inputs):
return None
if len(example_inputs) != len(graph_module.graph.placeholders):
return None
output_values = [
value
for output in graph_module.graph.outputs
for value in _nodes(output.args)
]
if len(output_values) != 1 or not isinstance(output_values[0], Node):
return None
final = output_values[0]
def is_pointwise(node: Node) -> bool:
if node.op not in {"call_function", "call_method"}:
return False
return _target_name(node.target) in _CPU_FUSED_OPS
def classify_reduction(node: Node) -> Any:
if node.op not in {"call_function", "call_method"}:
return None
if not node.args or not isinstance(node.args[0], Node):
return None
layout = _tensor_layout(_traced_value(graph_module, node.args[0]))
if layout is None:
return None
return _parse_reduction(node, len(layout[0]))
from .scheduler import segment_graph
segments = segment_graph(
graph_module,
is_pointwise=is_pointwise,
classify_reduction=classify_reduction,
)
if segments is None:
return None
if any(segment.epilogue for segment in segments):
return None
compiled_kinds = {"pw", "pw+red"}
if not any(segment.kind in compiled_kinds for segment in segments):
return None
if not any(segment.kind == "extern" for segment in segments):
# A region that is fusible end to end belongs to the whole-region
# paths; reaching here means they declined it for another reason.
return None
slots: dict[Node, int] = {
node: index for index, node in enumerate(graph_module.graph.placeholders)
}
constants: dict[int, Any] = {}
steps: list[tuple] = []
next_slot = len(slots)
def slot_for(node: Node) -> int | None:
if node in slots:
return slots[node]
if node.op != "get_attr":
return None
value = _traced_value(graph_module, node)
if not isinstance(value, tensor_type):
return None
nonlocal next_slot
slots[node] = next_slot
constants[next_slot] = value
next_slot += 1
return slots[node]
def extern_step(node: Node):
"""Close over one operator's call, resolving its operands by slot.
The plan is built once: an argument is either a slot to read or a
value to pass through, so the steady-state call walks a flat list
instead of rebuilding the captured argument structure.
"""
target = node.target
op = node.op
kwargs_template = dict(node.kwargs or {})
table = slots
simple = all(
not isinstance(item, (list, tuple, dict, slice))
for item in (*node.args, *kwargs_template.values())
)
if not simple:
from ..graph import map_arg
def run_general(values: list[Any]) -> Any:
resolve = lambda item: values[table[item]] # noqa: E731
args = map_arg(node.args, resolve)
kwargs = map_arg(kwargs_template, resolve)
if op == "call_function":
return target(*args, **kwargs)
return getattr(args[0], target)(*args[1:], **kwargs)
return run_general
plan = tuple(
(table[item], None) if isinstance(item, Node) else (-1, item)
for item in node.args
)
keys = tuple(kwargs_template)
key_plan = tuple(
(table[item], None) if isinstance(item, Node) else (-1, item)
for item in kwargs_template.values()
)
if not keys and op == "call_function":
def run_positional(values: list[Any]) -> Any:
return target(
*[
values[slot] if value is None else value
for slot, value in plan
]
)
return run_positional
def run(values: list[Any]) -> Any:
args = [values[slot] if value is None else value for slot, value in plan]
kwargs = {
key: (values[slot] if value is None else value)
for key, (slot, value) in zip(keys, key_plan)
}
if op == "call_function":
return target(*args, **kwargs)
return getattr(args[0], target)(*args[1:], **kwargs)
return run
def add_extern(node: Node) -> bool:
nonlocal next_slot
if node.op == "get_attr":
return slot_for(node) is not None
if node.op not in {"call_function", "call_method"}:
return False
operands = list(_nodes(node.args)) + list(_nodes(node.kwargs or {}))
for operand in operands:
if slot_for(operand) is None:
return False
sources = tuple(slots[operand] for operand in operands)
slots[node] = next_slot
next_slot += 1
steps.append((extern_step(node), slots[node], sources))
return True
compiled_count = 0
for segment in segments:
if segment.kind != "extern":
built = _build_segment_kernel(graph_module, segment, first.device)
if built is not None:
externals, runner = built
sources = []
for node in externals:
source = slot_for(node)
if source is None:
return None
sources.append(source)
export = segment.export_node
if export is None:
return None
slots[export] = next_slot
next_slot += 1
steps.append(
(
_kernel_step(runner, tuple(sources)),
slots[export],
tuple(sources),
)
)
compiled_count += 1
continue
# One run the generators cannot express does not cost the region
# its other kernels: those operators run individually instead.
for node in segment.nodes:
if not add_extern(node):
return None
if compiled_count == 0:
return None
if final not in slots:
return None
# Liveness: a value is dropped right after the step that reads it last,
# so long regions do not hold every intermediate alive to the end.
last_use: dict[int, int] = {}
for index, (_step, _target, sources) in enumerate(steps):
for slot in sources:
last_use[slot] = index
output_slot = slots[final]
plan = [
(
step,
target,
tuple(
slot
for slot, index in last_use.items()
if index == position and slot != output_slot
),
)
for position, (step, target, _sources) in enumerate(steps)
]
input_shapes = tuple(
tuple(int(item) for item in value.shape) for value in example_inputs
)
input_strides = tuple(
tuple(int(item) for item in value.stride()) for value in example_inputs
)
return _CpuSegmentedLowering(
graph_module,
plan,
next_slot,
constants,
output_slot,
first.dtype,
first.device,
tuple(zip(input_shapes, input_strides)),
strict_native,
)
def _build_fused_gradient_graphs(
input_count: int,
instructions: list[tuple[str, int, int, int]],
forward_program: list[int],
forward_constants: list[float],
forward_temp_count: int,
output_ref: int,
) -> tuple[list[int], list[float], tuple[int, ...]] | None:
"""Create one shared fused reverse-mode program for all inputs.
The forward intermediates are emitted once. Each input derivative then
extends that same program and records one final temporary, allowing the
native kernel to evaluate all gradients in one vector loop.
"""
def remap_forward_ref(ref: int) -> int:
if ref >= input_count:
return ref + 1 # reserve the final external input for grad_output
return ref
remapped_forward_program: list[int] = []
for offset in range(0, len(forward_program), 3):
remapped_forward_program.extend(
(
forward_program[offset],
remap_forward_ref(forward_program[offset + 1]),
remap_forward_ref(forward_program[offset + 2]),
)
)
program = list(remapped_forward_program)
constants = list(forward_constants)
temp_count = forward_temp_count
zero_ref = -(len(constants) + 1)
constants.append(0.0)
one_ref = -(len(constants) + 1)
constants.append(1.0)
two_ref = -(len(constants) + 1)
constants.append(2.0)
grad_output_ref = input_count
output_refs: list[int] = []
def emit(op_name: str, lhs: int, rhs: int = -1) -> int:
nonlocal temp_count
if op_name not in _CPU_FUSED_OPCODES:
raise ValueError(f"unsupported fused derivative op: {op_name}")
program.extend((_CPU_FUSED_OPCODES[op_name], lhs, rhs))
result = input_count + 1 + temp_count
temp_count += 1
return result
def is_zero(ref: int) -> bool:
return ref == zero_ref
def is_one(ref: int) -> bool:
return ref == one_ref
def add_ref(lhs: int, rhs: int) -> int:
if is_zero(lhs):
return rhs
if is_zero(rhs):
return lhs
return emit("add", lhs, rhs)
def sub_ref(lhs: int, rhs: int) -> int:
if is_zero(rhs):
return lhs
return emit("sub", lhs, rhs)
def mul_ref(lhs: int, rhs: int) -> int:
if is_zero(lhs) or is_zero(rhs):
return zero_ref
if is_one(lhs):
return rhs
if is_one(rhs):
return lhs
return emit("mul", lhs, rhs)
def neg_ref(ref: int) -> int:
if is_zero(ref):
return ref
return emit("neg", ref)
def div_ref(lhs: int, rhs: int) -> int:
if is_zero(lhs):
return zero_ref
if is_one(rhs):
return lhs
return emit("div", lhs, rhs)
adjoints: dict[int, int] = {
remap_forward_ref(output_ref): grad_output_ref,
}
def add_adjoint(ref: int, contribution: int) -> None:
if ref < 0 or is_zero(contribution):
return
adjoints[ref] = add_ref(adjoints.get(ref, zero_ref), contribution)
for op_name, lhs, rhs, result in reversed(instructions):
lhs = remap_forward_ref(lhs)
rhs = remap_forward_ref(rhs)
result = remap_forward_ref(result)
grad = adjoints.get(result, zero_ref)
if is_zero(grad):
continue
if op_name == "add":
add_adjoint(lhs, grad)
add_adjoint(rhs, grad)
elif op_name == "sub":
add_adjoint(lhs, grad)
add_adjoint(rhs, neg_ref(grad))
elif op_name == "mul":
add_adjoint(lhs, mul_ref(grad, rhs))
add_adjoint(rhs, mul_ref(grad, lhs))
elif op_name == "div":
add_adjoint(lhs, div_ref(grad, rhs))
denominator = mul_ref(rhs, rhs)
add_adjoint(rhs, neg_ref(div_ref(mul_ref(grad, lhs), denominator)))
elif op_name == "neg":
add_adjoint(lhs, neg_ref(grad))
elif op_name == "pos":
add_adjoint(lhs, grad)
elif op_name == "abs":
add_adjoint(lhs, mul_ref(grad, emit("abs_grad", lhs)))
elif op_name == "sin":
add_adjoint(lhs, mul_ref(grad, emit("cos", lhs)))
elif op_name == "cos":
add_adjoint(lhs, mul_ref(grad, neg_ref(emit("sin", lhs))))
elif op_name == "exp":
add_adjoint(lhs, mul_ref(grad, result))
elif op_name == "log":
add_adjoint(lhs, div_ref(grad, lhs))
elif op_name == "sigmoid":
local = mul_ref(result, sub_ref(one_ref, result))
add_adjoint(lhs, mul_ref(grad, local))
elif op_name == "sqrt":
add_adjoint(lhs, div_ref(grad, mul_ref(two_ref, result)))
elif op_name == "square":
add_adjoint(lhs, mul_ref(grad, mul_ref(two_ref, lhs)))
elif op_name == "tanh":
local = sub_ref(one_ref, mul_ref(result, result))
add_adjoint(lhs, mul_ref(grad, local))
elif op_name == "relu":
add_adjoint(lhs, mul_ref(grad, emit("relu_grad", lhs)))
else:
return None
# Make every output a temporary. This also handles a disconnected input
# (constant zero) and an input that receives grad_output directly.
for input_ref in range(input_count):
output_refs.append(emit("pos", adjoints.get(input_ref, zero_ref)))
# Remove forward values that are not needed by any local derivative. For
# example, d(sin(x))/dx uses cos(x), not the forward sin(x) result; this
# follows the derivative graph rather than blindly replaying all of
# the forward graph.
total_input_count = input_count + 1
instruction_count = len(program) // 3
live = [False] * instruction_count
pending = list(output_refs)
while pending:
ref = pending.pop()
if ref < total_input_count:
continue
instruction = ref - total_input_count
if instruction < 0 or instruction >= instruction_count or live[instruction]:
continue
live[instruction] = True
offset = instruction * 3
pending.extend((program[offset + 1], program[offset + 2]))
compact_refs: dict[int, int] = {}
next_instruction = 0
for instruction, is_live in enumerate(live):
if is_live:
compact_refs[total_input_count + instruction] = (
total_input_count + next_instruction
)
next_instruction += 1
compact_program: list[int] = []
for instruction, is_live in enumerate(live):
if not is_live:
continue
offset = instruction * 3
compact_program.extend(
(
program[offset],
compact_refs.get(program[offset + 1], program[offset + 1]),
compact_refs.get(program[offset + 2], program[offset + 2]),
)
)
output_refs = [compact_refs[ref] for ref in output_refs]
program = compact_program
if len(program) // 3 > 64:
return None
return program, constants, tuple(output_refs)
def _fold_eval_conv_batch_norm(
graph_module: GraphModule,
example_inputs: list[Any],
) -> dict[Node, tuple[Node, Any, Any]]:
"""Precompute inference BatchNorm parameters for Conv2d users.
ResNet inference contains the stable pattern ``conv2d -> batch_norm``.
Folding the running-statistics transform into the convolution removes one
full feature-map kernel and its intermediate write. The optimization is
intentionally restricted to eval-mode BatchNorm with a single Conv2d
user, so training graphs and branched tensors retain the ordinary native
operators.
The returned tensors are compile-time constants owned by the native
lowering. TensorPlay's public compile contract currently has no
parameter-version guard, therefore this pass is only enabled for the
inference lowering path; callers that mutate parameters must recompile.
"""
# Do not fold a graph that is being differentiated with respect to its
# runtime inputs. The folded parameters are inference constants, while
# eval-mode autograd still needs the original parameter edges.
if any(getattr(value, "requires_grad", False) for value in example_inputs):
return {}
try:
import tensorplay
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return {}
# Folding changes parameter dataflow, so it is valid only for the
# no-grad inference specialization that the benchmark requests. A
# grad-enabled compile must retain the ordinary Conv/BN autograd edges.
if tensorplay.is_grad_enabled():
return {}
def tensor_attr(value: Any) -> Any | None:
if not isinstance(value, Node) or value.op != "get_attr":
return None
attribute = graph_module._get_attr(value.target)
return attribute if isinstance(attribute, tensor_type) else None
folded: dict[Node, tuple[Node, Any, Any]] = {}
for batch_norm in graph_module.graph.nodes:
if (
batch_norm.op != "call_function"
or _target_name(batch_norm.target) != "batch_norm"
or len(batch_norm.args) != 8
or batch_norm.kwargs
or batch_norm.args[5] is not False
):
continue
conv = batch_norm.args[0]
if (
not isinstance(conv, Node)
or conv.op != "call_function"
or _target_name(conv.target) != "conv2d"
or len(conv.args) != 7
or conv.kwargs
or conv.users != {batch_norm}
):
continue
running_mean = tensor_attr(batch_norm.args[1])
running_var = tensor_attr(batch_norm.args[2])
bn_weight = tensor_attr(batch_norm.args[3])
bn_bias = tensor_attr(batch_norm.args[4])
conv_weight = tensor_attr(conv.args[1])
conv_bias = tensor_attr(conv.args[2])
eps = batch_norm.args[7]
if (
running_mean is None
or running_var is None
or conv_weight is None
or not isinstance(eps, numbers.Real)
or conv.args[2] is not None and conv_bias is None
or batch_norm.args[3] is not None and bn_weight is None
or batch_norm.args[4] is not None and bn_bias is None
):
continue
try:
with tensorplay.no_grad():
running_mean = running_mean.detach()
running_var = running_var.detach()
conv_weight = conv_weight.detach()
weight_coeff = tensorplay.rsqrt(running_var + float(eps))
scale = (
bn_weight.detach()
if bn_weight is not None
else tensorplay.ones_like(running_var)
) * weight_coeff
folded_weight = conv_weight * scale.reshape((-1, 1, 1, 1))
base_bias = (
conv_bias.detach()
if conv_bias is not None
else tensorplay.zeros_like(running_mean)
)
folded_bias = (
(base_bias - running_mean) * scale
+ (
bn_bias.detach()
if bn_bias is not None
else tensorplay.zeros_like(running_mean)
)
)
except (AttributeError, RuntimeError, TypeError, ValueError):
# Keep the ordinary native Conv+BN path if a backend dtype or
# device cannot materialize the folded constants.
continue
folded[conv] = (batch_norm, folded_weight, folded_bias)
return folded
_NATIVE_OP_SUPPORT: dict[str, bool] = {}
def _native_runs_linear() -> bool:
"""Whether the loaded runtime can execute a fused ``linear`` node.
The lowering and the runtime library are built separately, and a tree can
hold one newer than the other, so a node this module knows how to emit is
not necessarily one the runtime knows how to run. Probed once per
process with a one-node graph; a runtime that cannot run it keeps the
transpose-product-add form, which every runtime can.
"""
known = _NATIVE_OP_SUPPORT.get("linear")
if known is not None:
return known
supported = False
try:
import tensorplay
native = tensorplay._C._stax
graph = native.Graph()
node = graph.create_node("linear", "probe")
node.add_input(graph.add_input())
node.add_input(graph.add_input())
graph.register_output(node.add_output())
graph.execute([tensorplay.zeros((1, 1)), tensorplay.zeros((1, 1))])
supported = True
except Exception: # noqa: BLE001 - any failure means "emit the long form"
supported = False
_NATIVE_OP_SUPPORT["linear"] = supported
return supported
def _lower_native(
graph_module: GraphModule,
example_inputs: list[Any],
*,
use_fusion: bool = True,
extra_output_nodes: list[Node] | None = None,
) -> _NativeLowering | None:
"""Lower the canonical graph into the native Stax IR when possible."""
try:
import tensorplay
native_module = getattr(tensorplay._C, "_stax", None)
except (AttributeError, ImportError):
native_module = None
if native_module is None:
return None
if not hasattr(native_module.Graph, "execute"):
return None
try:
import tensorplay
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return None
if len(example_inputs) != len(graph_module.graph.placeholders):
return None
if any(not isinstance(value, tensor_type) for value in example_inputs):
# Stax's native ABI is Tensor-only. The generated GraphModule remains
# the correct compiled path for scalar/keyword placeholders.
return None
graph = native_module.Graph()
values: dict[Node, Any] = {}
attribute_targets: list[str] = []
constant_values: list[Any] = []
folded_convs = _fold_eval_conv_batch_norm(graph_module, example_inputs)
# convolution inputs, weights, and outputs. Its generated wrapper uses
# ``empty_strided`` tensors with the channels-last strides and the current
# rather than assuming NCHW. Materialize the same weight layout once for
# folded constants; runtime activations are converted by a native graph
# node immediately before the first convolution that consumes them.
use_channels_last = bool(
example_inputs
and example_inputs[0].device.is_cuda()
and not tensorplay.is_grad_enabled()
)
if use_channels_last and folded_convs:
try:
folded_with_layout: dict[Node, tuple[Node, Any, Any]] = {}
for conv, (batch_norm, weight, bias) in folded_convs.items():
# NCHW logical shape, NHWC physical storage, then reinterpret
# as NCHW. This layout is represented by generated weight
# strides, e.g. [K*C*R*S, 1, S*C, C].
physical_weight = weight.permute((0, 2, 3, 1)).clone()
channels_last_weight = physical_weight.permute((0, 3, 1, 2))
folded_with_layout[conv] = (
batch_norm,
channels_last_weight,
bias,
)
folded_convs = folded_with_layout
except (AttributeError, RuntimeError, TypeError, ValueError):
# Keep the already validated folding path if this backend cannot
# materialize the specialized layout for a particular dtype.
use_channels_last = False
folded_native_inputs: dict[Node, tuple[Any, Any]] = {}
folded_batch_norm_to_conv = {
batch_norm: conv for conv, (batch_norm, _, _) in folded_convs.items()
}
# belong in the native ABI. Inference Conv+BN folding leaves the original
# parameter nodes in the captured graph, but the native graph consumes only the folded
# weight/bias constants. Passing those dead tensors through Python and
# C++ on every invocation is pure call-boundary overhead.
live_attribute_nodes: set[Node] = set()
visited_nodes: set[Node] = set()
def visit_native_dependency(value: Any) -> None:
if isinstance(value, Node):
if value in visited_nodes:
return
visited_nodes.add(value)
if value.op == "get_attr":
live_attribute_nodes.add(value)
return
if value in folded_convs or value in folded_batch_norm_to_conv:
visit_native_dependency(value.args[0])
return
for argument in value.args:
visit_native_dependency(argument)
for argument in value.kwargs.values():
visit_native_dependency(argument)
return
if isinstance(value, (tuple, list)):
for item in value:
visit_native_dependency(item)
elif isinstance(value, dict):
for item in value.values():
visit_native_dependency(item)
elif isinstance(value, slice):
visit_native_dependency(value.start)
visit_native_dependency(value.stop)
visit_native_dependency(value.step)
for output in graph_module.graph.outputs:
visit_native_dependency(output.args)
for extra_node in extra_output_nodes or []:
visit_native_dependency(extra_node)
fused_relu_convs: dict[Node, Node] = {}
fused_add_relus: dict[Node, Node] = {}
fused_relu_nodes: set[Node] = set()
layout_values: dict[Node, bool] = {}
channels_last_values: dict[Node, Any] = {}
peel_conv_bias = bool(example_inputs and example_inputs[0].device.is_cuda())
# This is the same producer/sole-consumer legality check used by
# training graph because add_relu has a generated autograd formula; the
# Conv->ReLU and Conv+BN folding paths remain inference-only.
if use_fusion:
for relu in graph_module.graph.nodes:
if (
relu.op != "call_function"
or _target_name(relu.target) != "relu"
or len(relu.args) != 1
or relu.kwargs not in ({}, {"inplace": False}, {"inplace": True})
):
continue
source = relu.args[0]
if (
isinstance(source, Node)
and source.op == "call_function"
and _target_name(source.target) == "add"
and source.users == {relu}
):
if len(source.args) == 2:
lhs, rhs = source.args
alpha = 1
elif len(source.args) == 3:
lhs, rhs, alpha = source.args
else:
continue
if (
alpha == 1
and isinstance(lhs, Node)
and isinstance(rhs, Node)
):
fused_add_relus[source] = relu
continue
if tensorplay.is_grad_enabled():
continue
conv = folded_batch_norm_to_conv.get(source)
if conv is None:
conv = source
if not (
isinstance(conv, Node)
and conv.op == "call_function"
and _target_name(conv.target) == "conv2d"
):
continue
if source.users != {relu} or conv.users != ({source} if source is not conv else {relu}):
continue
fused_relu_convs[conv] = relu
for node in graph_module.graph.placeholders:
values[node] = graph.add_input()
layout_values[node] = False
def channels_last_value(node: Any) -> Any | None:
"""Return the native value in the preferred 4-D layout."""
if not isinstance(node, Node) or node not in values:
return None
if not use_channels_last:
return values[node]
if layout_values.get(node, False):
return values[node]
cached = channels_last_values.get(node)
if cached is not None:
return cached
reorder = graph.create_node("channels_last", f"{node.name}_channels_last")
reorder.add_input(values[node])
converted = reorder.add_output()
channels_last_values[node] = converted
return converted
# Register all live module attributes before synthetic folded weights so
# Graph::execute's input order is independent of where get_attr nodes are
# placed in the Python graph.
for node in graph_module.graph.nodes:
if node.op != "get_attr" or node not in live_attribute_nodes:
continue
attribute = graph_module._get_attr(node.target)
if not isinstance(attribute, tensor_type):
return None
values[node] = graph.add_input()
attribute_targets.append(node.target)
for node in graph_module.graph.nodes:
folded = folded_convs.get(node)
if folded is None:
continue
_, folded_weight, folded_bias = folded
native_weight = graph.add_input()
native_bias = graph.add_input()
folded_native_inputs[node] = (native_weight, native_bias)
constant_values.extend((folded_weight, folded_bias))
for node in graph_module.graph.nodes:
if node.op in {"placeholder", "output", "get_attr"}:
continue
if node.op not in {"call_function", "call_method"}:
return None
if (
node.op == "call_function"
and node.target is operator.getitem
and len(node.args) == 2
and isinstance(node.args[0], Node)
and isinstance(node.args[1], int)
and isinstance(values.get(node.args[0]), tuple)
):
source_values = values[node.args[0]]
index = node.args[1]
if index < 0:
index += len(source_values)
if index < 0 or index >= len(source_values):
return None
values[node] = source_values[index]
layout_values[node] = False
continue
op_name = _target_name(node.target)
# User-defined operators enter the native dispatcher bridge. Their
# output arity is carried by the graph metadata because the native
# value table reserves one value per result.
if isinstance(node.target, _CustomOpDef):
if node.kwargs or not node.args:
return None
native_node = graph.create_node("custom_op", node.name)
native_node.set_str_attr("op_name", node.target.name)
for argument in node.args:
resolved = values.get(argument) if isinstance(argument, Node) else None
if resolved is None or isinstance(resolved, tuple):
return None
native_node.add_input(resolved)
custom = node.meta.get("custom")
output_count = 1
if isinstance(custom, dict):
output_count = int(custom.get("nested_output_count", 1))
if output_count < 1:
return None
output_values = tuple(native_node.add_output() for _ in range(output_count))
values[node] = output_values[0] if output_count == 1 else output_values
continue
# ``ReLU(inplace=True)`` is an aliasing detail of the captured graph.
# The native kernel accepts that schema explicitly; other keyword
# combinations are rejected by this lowering.
if node.kwargs:
if (
op_name != "relu"
or node.kwargs not in ({"inplace": False}, {"inplace": True})
):
return None
if op_name not in _NATIVE_OPS:
return None
if op_name == "relu" and node in fused_relu_nodes:
source = node.args[0]
if source not in values:
return None
# The producer was lowered with a fused Conv+ReLU primitive; the
# ReLU value is an alias of that already-activated output.
values[node] = values[source]
layout_values[node] = layout_values.get(source, False)
continue
if op_name == "relu" and node in fused_add_relus.values():
source = node.args[0]
if not isinstance(source, Node) or source not in values:
return None
# The residual add is lowered as add_relu below, so the ReLU
# node observes the already-activated output without another
# native launch.
values[node] = values[source]
layout_values[node] = layout_values.get(source, False)
continue
def node_value(value: Any) -> Any | None:
if not isinstance(value, Node) or value not in values:
return None
return values[value]
def add_tensor_input(native_node: Any, value: Any) -> bool:
resolved = node_value(value)
if resolved is None or isinstance(resolved, tuple):
return False
native_node.add_input(resolved)
return True
if op_name == "conv2d":
if len(node.args) != 7:
return None
input_node, weight_node, bias_node, stride, padding, dilation, groups = node.args
fused_relu = fused_relu_convs.get(node)
use_conv_relu = fused_relu is not None and not peel_conv_bias
native_node = graph.create_node(
"conv2d_relu" if use_conv_relu else "conv2d",
node.name,
)
conv_input = channels_last_value(input_node)
if conv_input is None:
return None
native_node.add_input(conv_input)
folded_inputs = folded_native_inputs.get(node)
bias_input = None
bias_tensor = None
if folded_inputs is not None:
native_node.add_input(folded_inputs[0])
bias_input = folded_inputs[1]
folded_spec = folded_convs.get(node)
bias_tensor = folded_spec[2] if folded_spec is not None else None
else:
conv_weight = channels_last_value(weight_node)
if conv_weight is None:
return None
native_node.add_input(conv_weight)
if bias_node is None:
native_node.set_int_attr("has_bias", 0)
else:
bias_input = node_value(bias_node)
if bias_input is None:
return None
try:
bias_tensor = graph_module._get_attr(bias_node.target)
except (AttributeError, TypeError):
return None
if bias_input is None:
native_node.set_int_attr("has_bias", 0)
elif not peel_conv_bias or use_conv_relu:
native_node.add_input(bias_input)
native_node.set_int_attr("has_bias", 1)
else:
# the cuDNN call because cuDNN is slower with it. Keep the
# bias as a broadcast pointwise input after the convolution.
native_node.set_int_attr("has_bias", 0)
if not all(
_set_int_list_attr(native_node, key, value)
for key, value in (
("stride", stride),
("padding", padding),
("dilation", dilation),
)
):
return None
if isinstance(groups, bool) or not isinstance(groups, int):
return None
native_node.set_int_attr("groups", int(groups))
conv_value = native_node.add_output()
values[node] = conv_value
layout_values[node] = use_channels_last
if peel_conv_bias and bias_input is not None and not use_conv_relu:
if bias_tensor is None or not hasattr(bias_tensor, "shape"):
return None
bias_shape = tuple(int(item) for item in bias_tensor.shape)
if len(bias_shape) != 1:
return None
bias_view = graph.create_node("reshape", f"{node.name}_bias_view")
bias_view.add_input(bias_input)
bias_view.set_ints_attr("shape", [1, bias_shape[0], 1, 1])
bias_value = bias_view.add_output()
add_node = graph.create_node(
"add_relu" if fused_relu is not None else "add",
f"{node.name}_bias_add",
)
add_node.add_input(conv_value)
add_node.add_input(bias_value)
values[node] = add_node.add_output()
layout_values[node] = use_channels_last and fused_relu is not None
if fused_relu is not None:
fused_relu_nodes.add(fused_relu)
elif use_conv_relu and fused_relu is not None:
fused_relu_nodes.add(fused_relu)
continue
if op_name == "batch_norm":
if len(node.args) != 8:
return None
folded_batch_norm = next(
(
batch_norm
for batch_norm, _, _ in folded_convs.values()
if batch_norm is node
),
None,
)
if folded_batch_norm is not None:
values[node] = values[node.args[0]]
layout_values[node] = layout_values.get(node.args[0], False)
continue
input_node, running_mean, running_var, weight, bias, training, momentum, eps = node.args
native_node = graph.create_node("batch_norm", node.name)
if not add_tensor_input(native_node, input_node):
return None
optional_inputs = (
("has_running_mean", running_mean),
("has_running_var", running_var),
("has_weight", weight),
("has_bias", bias),
)
for attr_name, optional_node in optional_inputs:
if optional_node is None:
native_node.set_int_attr(attr_name, 0)
continue
if not add_tensor_input(native_node, optional_node):
return None
native_node.set_int_attr(attr_name, 1)
if not isinstance(training, bool) or not isinstance(momentum, numbers.Real) or not isinstance(eps, numbers.Real):
return None
native_node.set_int_attr("training", int(training))
native_node.set_float_attr("momentum", float(momentum))
native_node.set_float_attr("eps", float(eps))
values[node] = native_node.add_output()
layout_values[node] = False
continue
if op_name == "max_pool2d":
if len(node.args) != 7:
return None
input_node, kernel_size, stride, padding, dilation, ceil_mode, return_indices = node.args
if return_indices is not False or not isinstance(ceil_mode, bool):
return None
native_node = graph.create_node("max_pool2d", node.name)
if not add_tensor_input(native_node, input_node):
return None
for key, value in (
("kernel_size", kernel_size),
("stride", stride),
("padding", padding),
("dilation", dilation),
):
if not _set_int_list_attr(native_node, key, value):
return None
native_node.set_int_attr("ceil_mode", int(ceil_mode))
values[node] = native_node.add_output()
# The cuDNN tensor descriptor and output follow the input layout;
# a later convolution can therefore consume the max-pool result
# without an NCHW round-trip.
layout_values[node] = use_channels_last and layout_values.get(
input_node, False
)
continue
if op_name == "adaptive_avg_pool2d":
if len(node.args) != 2 or not _set_int_list_attr(
native_node := graph.create_node("adaptive_avg_pool2d", node.name),
"output_size",
node.args[1],
):
return None
if not add_tensor_input(native_node, node.args[0]):
return None
values[node] = native_node.add_output()
layout_values[node] = False
continue
if op_name == "flatten":
if len(node.args) != 2 or node.args[1] != 1 or node.kwargs:
return None
native_node = graph.create_node("flatten", node.name)
if not add_tensor_input(native_node, node.args[0]):
return None
native_node.set_int_attr("start_dim", 1)
native_node.set_int_attr("end_dim", -1)
values[node] = native_node.add_output()
layout_values[node] = False
continue
if op_name == "add" and node in fused_add_relus:
if len(node.args) == 2:
input_node, other_node = node.args
alpha = 1
elif len(node.args) == 3:
input_node, other_node, alpha = node.args
else:
return None
if (
alpha != 1
or not isinstance(input_node, Node)
or not isinstance(other_node, Node)
or input_node not in values
or other_node not in values
):
return None
fused = graph.create_node("add_relu", node.name)
fused.add_input(values[input_node])
fused.add_input(values[other_node])
values[node] = fused.add_output()
layout_values[node] = use_channels_last and (
layout_values.get(input_node, False)
or layout_values.get(other_node, False)
)
continue
# ``alpha`` argument, so their graph node has
# ``(input, other, alpha)`` even when alpha is the default 1, and
# keyword-only spellings record alpha in the node kwargs. Lower
# that contract to the native pointwise IR instead of falling back
# to a Python method call. Non-unit alpha becomes a scalar multiply
# and can be consumed by Stax's mul-add fusion pass.
if op_name in {"add", "sub"} and (
len(node.args) == 3
or (len(node.args) == 2 and "alpha" in (node.kwargs or {}))
):
if len(node.args) == 3 and not node.kwargs:
input_node, other_node, alpha = node.args
else:
input_node, other_node = node.args
alpha = node.kwargs.get("alpha", 1)
if not isinstance(input_node, Node) or not _is_scalar(alpha):
return None
if input_node not in values:
return None
if isinstance(other_node, Node) and other_node not in values:
return None
if alpha == 1:
binary = graph.create_node(op_name, node.name)
binary.add_input(values[input_node])
if isinstance(other_node, Node):
binary.add_input(values[other_node])
elif _is_scalar(other_node):
_set_scalar_attr(binary, other_node, 1)
else:
return None
values[node] = binary.add_output()
layout_values[node] = False
continue
if isinstance(other_node, Node):
scale = graph.create_node("mul", f"{node.name}_alpha")
scale.add_input(values[other_node])
_set_scalar_attr(scale, alpha, 1)
scaled_other = scale.add_output()
binary = graph.create_node(op_name, node.name)
binary.add_input(values[input_node])
binary.add_input(scaled_other)
values[node] = binary.add_output()
layout_values[node] = False
continue
if _is_scalar(other_node):
binary = graph.create_node(op_name, node.name)
binary.add_input(values[input_node])
_set_scalar_attr(binary, other_node * alpha, 1)
values[node] = binary.add_output()
layout_values[node] = False
continue
return None
if op_name == "linear":
if len(node.args) not in {2, 3} or any(
not isinstance(arg, Node) and arg is not None for arg in node.args
):
return None
input_node, weight_node = node.args[:2]
bias_node = node.args[2] if len(node.args) == 3 else None
if not isinstance(input_node, Node) or not isinstance(weight_node, Node):
return None
if bias_node is not None and not isinstance(bias_node, Node):
return None
if any(
value_node not in values
for value_node in (input_node, weight_node, bias_node)
if value_node is not None
):
return None
# One node, not a transpose plus a product plus an addition: the
# bias belongs in the product's epilogue, and adding it back
# separately costs a whole pass over the output.
if _native_runs_linear():
fused = graph.create_node("linear", node.name)
fused.add_input(values[input_node])
fused.add_input(values[weight_node])
if bias_node is not None:
fused.add_input(values[bias_node])
values[node] = fused.add_output()
layout_values[node] = False
continue
transpose = graph.create_node("t", f"{node.name}_weight_t")
transpose.add_input(values[weight_node])
transposed_weight = transpose.add_output()
matmul = graph.create_node("matmul", f"{node.name}_matmul")
matmul.add_input(values[input_node])
matmul.add_input(transposed_weight)
result = matmul.add_output()
if bias_node is not None:
add = graph.create_node("add", f"{node.name}_bias")
add.add_input(result)
add.add_input(values[bias_node])
result = add.add_output()
values[node] = result
layout_values[node] = False
continue
input_nodes: list[Node] = []
scalar_args: list[tuple[int, Any]] = []
for position, arg in enumerate(node.args):
if isinstance(arg, Node):
input_nodes.append(arg)
elif _is_scalar(arg):
scalar_args.append((position, arg))
else:
return None
if len(scalar_args) > 1:
return None
if op_name in {
"neg",
"pos",
"abs",
"sin",
"cos",
"exp",
"log",
"sigmoid",
"sqrt",
"square",
"tanh",
"relu",
}:
if len(node.args) != 1 or len(input_nodes) != 1:
return None
elif len(input_nodes) not in {1, 2} or len(node.args) not in {1, 2}:
return None
if any(input_node not in values for input_node in input_nodes):
return None
if op_name == "mm":
if len(node.args) != 2 or len(input_nodes) != 2:
return None
native_node = graph.create_node("mm", node.name)
native_node.add_input(values[input_nodes[0]])
native_node.add_input(values[input_nodes[1]])
values[node] = native_node.add_output()
layout_values[node] = False
continue
native_node = graph.create_node(op_name, node.name)
for input_node in input_nodes:
native_node.add_input(values[input_node])
if scalar_args:
_set_scalar_attr(native_node, scalar_args[0][1], scalar_args[0][0])
if op_name == "relu":
# Preserve the functional schema. The executor may call relu_
# only when the captured call explicitly requested mutation.
native_node.set_int_attr(
"inplace", int(node.kwargs.get("inplace", False))
)
values[node] = native_node.add_output()
layout_values[node] = False
try:
output_arg = graph_module.graph.output_node.args[0]
output_values = _native_value_leaves(output_arg, values)
output_spec = _native_output_spec(output_arg, values)
except (IndexError, KeyError, RuntimeError, TypeError):
return None
if not output_values or any(value is None for value in output_values):
return None
for output_value in output_values:
graph.register_output(output_value)
public_output_count = len(output_values)
registered_extra_outputs = 0
for extra_node in extra_output_nodes or []:
if extra_node not in values:
return None
extra_values = values[extra_node]
if isinstance(extra_values, tuple):
if not extra_values or any(value is None for value in extra_values):
return None
for extra_value in extra_values:
graph.register_output(extra_value)
registered_extra_outputs += len(extra_values)
else:
if extra_values is None:
return None
graph.register_output(extra_values)
registered_extra_outputs += 1
if use_fusion and not registered_extra_outputs:
graph.fuse()
return _NativeLowering(
graph_module,
graph,
attribute_targets,
constant_values,
output_count=public_output_count + registered_extra_outputs,
native_values=values,
output_spec=output_spec,
public_output_count=public_output_count,
)
class _AotShape(tuple):
"""Tensor metadata that supports both ``shape`` and ``shape()`` schemas."""
def __new__(cls, value: Any):
return super().__new__(cls, (int(item) for item in value))
def __call__(self) -> tuple[int, ...]:
return tuple(self)
class _AotNativeSymbol:
"""A symbolic Tensor value used while materializing an AOT backward graph."""
__slots__ = ("builder", "value", "shape")
def __init__(self, builder: "_AotNativeGraphBuilder", value: Any, shape: Any):
self.builder = builder
self.value = value
self.shape = _AotShape(shape)
def _binary(self, op_name: str, other: Any) -> "_AotNativeSymbol":
return self.builder.binary(op_name, self, other)
def __add__(self, other: Any) -> "_AotNativeSymbol":
return self._binary("add", other)
def __radd__(self, other: Any) -> "_AotNativeSymbol":
return self.builder.binary("add", other, self)
def __sub__(self, other: Any) -> "_AotNativeSymbol":
return self._binary("sub", other)
def __rsub__(self, other: Any) -> "_AotNativeSymbol":
return self.builder.binary("sub", other, self)
def __mul__(self, other: Any) -> "_AotNativeSymbol":
return self._binary("mul", other)
def __rmul__(self, other: Any) -> "_AotNativeSymbol":
return self.builder.binary("mul", other, self)
def __truediv__(self, other: Any) -> "_AotNativeSymbol":
return self._binary("div", other)
def __rtruediv__(self, other: Any) -> "_AotNativeSymbol":
return self.builder.binary("div", other, self)
def __neg__(self) -> "_AotNativeSymbol":
return self.builder.unary("neg", self)
def __pos__(self) -> "_AotNativeSymbol":
return self.builder.unary("pos", self)
def t(self) -> "_AotNativeSymbol":
return self.builder.unary("t", self, shape=self.shape[::-1])
def mm(self, other: Any) -> "_AotNativeSymbol":
return self.builder.binary("mm", self, other)
def matmul(self, other: Any) -> "_AotNativeSymbol":
return self.builder.binary("matmul", self, other)
def reshape(self, shape: Any) -> "_AotNativeSymbol":
return self.builder.reshape(self, shape)
def view(self, shape: Any) -> "_AotNativeSymbol":
return self.builder.reshape(self, shape)
def expand(self, shape: Any) -> "_AotNativeSymbol":
raise NotImplementedError("AOT native expand lowering is not implemented")
def sum(self, dim: Any = None, keepdim: bool = False) -> "_AotNativeSymbol":
return self.builder.sum(self, dim, keepdim)
def numel(self) -> int:
result = 1
for item in self.shape:
result *= item
return result
class _AotNativeTuple:
__slots__ = ("values",)
def __init__(self, values: tuple[_AotNativeSymbol, ...]):
self.values = values
class _AotNativeGraphBuilder:
"""Small native-IR builder used by the source-derived reverse pass."""
def __init__(self, native_module: Any):
self.native_module = native_module
self.graph = native_module.Graph()
@staticmethod
def _shape(value: Any) -> tuple[int, ...]:
return tuple(int(item) for item in getattr(value, "shape", ()))
@staticmethod
def _symbol(value: Any) -> _AotNativeSymbol | None:
return value if isinstance(value, _AotNativeSymbol) else None
def input(self, example_value: Any) -> _AotNativeSymbol:
return _AotNativeSymbol(self, self.graph.add_input(), self._shape(example_value))
def _add_inputs(self, native_node: Any, args: tuple[Any, ...]) -> list[_AotNativeSymbol]:
symbols: list[_AotNativeSymbol] = []
for value in args:
if not isinstance(value, _AotNativeSymbol):
raise TypeError("AOT native op received a non-Tensor argument")
native_node.add_input(value.value)
symbols.append(value)
return symbols
@staticmethod
def _broadcast_shape(lhs: tuple[int, ...], rhs: tuple[int, ...]) -> tuple[int, ...]:
result: list[int] = []
for left, right in zip(reversed(lhs), reversed(rhs)):
if left != right and left != 1 and right != 1:
raise ValueError(f"incompatible AOT shapes: {lhs} and {rhs}")
result.append(max(left, right))
longer = lhs if len(lhs) >= len(rhs) else rhs
result.extend(reversed(longer[: abs(len(lhs) - len(rhs))]))
return tuple(reversed(result))
def binary(self, op_name: str, lhs: Any, rhs: Any) -> _AotNativeSymbol:
lhs_symbol = self._symbol(lhs)
rhs_symbol = self._symbol(rhs)
if lhs_symbol is None and rhs_symbol is None:
if op_name == "add":
return lhs + rhs
if op_name == "sub":
return lhs - rhs
if op_name == "mul":
return lhs * rhs
if op_name == "div":
return lhs / rhs
raise NotImplementedError(f"AOT scalar operation is unsupported: {op_name}")
native_node = self.graph.create_node(op_name, f"aot_{op_name}_{len(self.graph.nodes)}")
shape = lhs_symbol.shape if lhs_symbol is not None else rhs_symbol.shape
if lhs_symbol is not None and rhs_symbol is not None:
native_node.add_input(lhs_symbol.value)
native_node.add_input(rhs_symbol.value)
shape = self._broadcast_shape(lhs_symbol.shape, rhs_symbol.shape)
else:
symbol = lhs_symbol if lhs_symbol is not None else rhs_symbol
scalar = rhs if lhs_symbol is not None else lhs
native_node.add_input(symbol.value)
_set_scalar_attr(native_node, scalar, 1 if lhs_symbol is not None else 0)
return _AotNativeSymbol(self, native_node.add_output(), shape)
def unary(
self,
op_name: str,
value: _AotNativeSymbol,
*,
shape: tuple[int, ...] | None = None,
) -> _AotNativeSymbol:
native_node = self.graph.create_node(op_name, f"aot_{op_name}_{len(self.graph.nodes)}")
native_node.add_input(value.value)
return _AotNativeSymbol(self, native_node.add_output(), shape or value.shape)
def helper(
self,
op_name: str,
args: tuple[_AotNativeSymbol, ...],
*,
attrs: dict[str, Any] | None = None,
shape: tuple[int, ...] | None = None,
outputs: int = 1,
) -> _AotNativeSymbol | _AotNativeTuple:
native_node = self.graph.create_node(op_name, f"aot_{op_name}_{len(self.graph.nodes)}")
symbols = self._add_inputs(native_node, args)
del symbols
for key, value in (attrs or {}).items():
if isinstance(value, bool) or isinstance(value, int):
native_node.set_int_attr(key, int(value))
elif isinstance(value, numbers.Real):
native_node.set_float_attr(key, float(value))
elif isinstance(value, (tuple, list)) and all(
isinstance(item, int) and not isinstance(item, bool) for item in value
):
native_node.set_ints_attr(key, [int(item) for item in value])
else:
raise TypeError(f"unsupported AOT native attribute: {key}={value!r}")
if outputs == 1:
return _AotNativeSymbol(self, native_node.add_output(), shape or args[0].shape)
return _AotNativeTuple(
tuple(
_AotNativeSymbol(self, native_node.add_output(), shape or args[0].shape)
for _ in range(outputs)
)
)
def reshape(self, value: _AotNativeSymbol, shape: Any) -> _AotNativeSymbol:
normalized = tuple(int(item) for item in shape)
return self.helper("reshape", (value,), attrs={"shape": normalized}, shape=normalized) # type: ignore[return-value]
def sum(
self, value: _AotNativeSymbol, dim: Any = None, keepdim: bool = False
) -> _AotNativeSymbol:
if dim is None:
return self.helper("sum", (value,), shape=()) # type: ignore[return-value]
dims = tuple(int(item) for item in (dim if isinstance(dim, (tuple, list)) else (dim,)))
normalized_dims = tuple(item if item >= 0 else item + len(value.shape) for item in dims)
shape = list(value.shape)
if keepdim:
for item in normalized_dims:
shape[item] = 1
else:
for item in sorted(normalized_dims, reverse=True):
shape.pop(item)
return self.helper(
"sum",
(value,),
attrs={"dim": normalized_dims, "keepdim": bool(keepdim)},
shape=tuple(shape),
) # type: ignore[return-value]
def _aot_derivative_specs() -> dict[str, tuple[Any, dict[str, str]]]:
"""Read the local derivative schema used by TensorPlay code generation."""
from pathlib import Path
from tools.codegen.model import parse_derivatives_yaml, parse_schema
yaml_path = Path(__file__).resolve().parents[2] / "config" / "derivatives.yaml"
result: dict[str, tuple[Any, dict[str, str]]] = {}
for definition in parse_derivatives_yaml(str(yaml_path)):
parsed = parse_schema(definition["name"])
formulas = {
key: value for key, value in definition.items() if key != "name"
}
result[parsed.func_name] = (parsed, formulas)
return result
def _aot_formula_python(formula: str, tensor_params: set[str]) -> str:
"""Compile one derivatives.yaml formula into a Python expression.
Shares the codegen expression AST (tokenizer + parser); the emitter
renders against the runtime formula env -- builder callables like
add/mul/t/sum plus get_tuple -- instead of the C++ text the generated
autograd nodes need.
"""
from tools.codegen.gen_autograd import (
BinOp, BoolLit, Braced, Call, Method, Neg, Num, StrLit, Var,
TENSOR_METHODS, parse_expr,
)
symbols = set(tensor_params) | {"grad", "grad_output", "result"}
def is_tensor(expr: Any) -> bool:
if isinstance(expr, Var):
return expr.name in symbols
if isinstance(expr, Neg):
return looks_tensor(expr.value)
if isinstance(expr, Method):
return expr.name.rstrip("_") in TENSOR_METHODS
if isinstance(expr, Call):
leaf = expr.callee.split("::")[-1].split("<")[0]
return leaf not in ("Scalar",)
if isinstance(expr, BinOp):
return is_tensor(expr.left) or looks_tensor(expr.right)
return False
def looks_tensor(expr: Any) -> bool:
return is_tensor(expr) or isinstance(expr, BinOp)
def emit(expr: Any) -> str:
if isinstance(expr, Num):
return expr.text
if isinstance(expr, BoolLit):
return "True" if expr.text == "true" else "False"
if isinstance(expr, StrLit):
return expr.text
if isinstance(expr, Var):
return expr.name
if isinstance(expr, Neg):
inner = emit(expr.value)
return f"neg({inner})" if looks_tensor(expr.value) else f"-{inner}"
if isinstance(expr, Braced):
# Python target: a braced list renders as a tuple (builder.sum
# dims, reshape shapes), matching _aot_default_value.
items = [emit(item) for item in expr.items]
if len(items) == 1:
return f"({items[0]},)"
return f"({', '.join(items)})"
if isinstance(expr, Call):
args = ", ".join(emit(a) for a in expr.args)
get = re.fullmatch(r"std::get<(\d+)>", expr.callee)
if get:
return f"get_tuple({get.group(1)}, {args})"
callee = expr.callee.split("::")[-1]
return f"{callee}({args})"
if isinstance(expr, Method):
recv = emit(expr.receiver)
args = ", ".join(emit(a) for a in expr.args)
name = expr.name
base = name[:-1] if name.endswith("_") and name[:-1] in TENSOR_METHODS else name
if base in TENSOR_METHODS:
return f"{TENSOR_METHODS[base]}({recv}, {args})" if args \
else f"{TENSOR_METHODS[base]}({recv})"
return f"{recv}.{name}({args})" if args else f"{recv}.{name}()"
if isinstance(expr, BinOp):
left = emit(expr.left)
right = emit(expr.right)
left_tensor = is_tensor(expr.left)
right_tensor = looks_tensor(expr.right)
if expr.op in "+-" and left_tensor:
return f"{'add' if expr.op == '+' else 'sub'}({left}, {right})"
if expr.op == "*" and left_tensor:
return f"mul({left}, {right})"
if expr.op == "/" and left_tensor:
return f"div({left}, {right})"
if expr.op == "*" and right_tensor:
return f"mul({right}, {left})"
if expr.op == "-" and right_tensor:
return f"neg(sub({right}, {left}))"
return f"({left} {expr.op} {right})"
raise NotImplementedError(f"AOT formula node is unsupported: {expr!r}")
return emit(parse_expr(formula))
def _build_aot_formula_env(
builder: _AotNativeGraphBuilder,
*,
batch_norm_cache: dict[tuple[int, ...], _AotNativeTuple],
) -> dict[str, Any]:
def binary(name: str):
return lambda lhs, rhs: builder.binary(name, lhs, rhs)
def unary(name: str):
return lambda value: builder.unary(name, value)
def get_tuple(index: int, value: _AotNativeTuple):
return value.values[int(index)]
def batch_norm_backward(*args: Any):
key = tuple(id(item) if isinstance(item, _AotNativeSymbol) else hash(repr(item)) for item in args)
cached = batch_norm_cache.get(key)
if cached is not None:
return cached
grad, input_value, weight, running_mean, running_var, training, eps = args
tensor_args = (grad, input_value)
attrs = {
"has_weight": weight is not None,
"has_running_mean": running_mean is not None,
"has_running_var": running_var is not None,
"training": bool(training),
"eps": float(eps),
}
optional = tuple(item for item in (weight, running_mean, running_var) if item is not None)
value = builder.helper(
"batch_norm_backward",
tensor_args + optional,
attrs=attrs,
shape=input_value.shape,
outputs=3,
)
assert isinstance(value, _AotNativeTuple)
batch_norm_cache[key] = value
return value
def conv_grad(name: str):
def invoke(grad, input_value, weight, stride, padding, dilation, groups):
return builder.helper(
name,
(grad, input_value, weight),
attrs={
"stride": tuple(stride),
"padding": tuple(padding),
"dilation": tuple(dilation),
"groups": int(groups),
},
shape=(
input_value.shape
if name.endswith("input")
else weight.shape
if name.endswith("weight")
else (grad.shape[1],)
),
)
return invoke
def max_pool_backward(grad, input_value, kernel_size, stride, padding, dilation, ceil_mode):
return builder.helper(
"max_pool2d_backward",
(grad, input_value),
attrs={
"kernel_size": tuple(kernel_size),
"stride": tuple(stride),
"padding": tuple(padding),
"dilation": tuple(dilation),
"ceil_mode": bool(ceil_mode),
},
shape=input_value.shape,
)
def adaptive_avg_pool_backward(grad, input_value):
return builder.helper(
"adaptive_avg_pool2d_backward", (grad, input_value), shape=input_value.shape
)
def threshold_backward(grad, output, threshold):
return builder.helper(
"threshold_backward",
(grad, output),
attrs={"threshold": threshold},
shape=grad.shape,
)
return {
"add": binary("add"),
"sub": binary("sub"),
"mul": binary("mul"),
"div": binary("div"),
"matmul": binary("matmul"),
"mm": binary("mm"),
"neg": unary("neg"),
"pos": unary("pos"),
"t": unary("t"),
"reshape": builder.reshape,
"sum": builder.sum,
"get_tuple": get_tuple,
"batch_norm_backward": batch_norm_backward,
"conv2d_grad_input": conv_grad("conv2d_grad_input"),
"conv2d_grad_weight": conv_grad("conv2d_grad_weight"),
"conv2d_grad_bias": conv_grad("conv2d_grad_bias"),
"max_pool2d_backward": max_pool_backward,
"adaptive_avg_pool2d_backward": adaptive_avg_pool_backward,
"threshold_backward": threshold_backward,
}
def _aot_schema_for(
specs: dict[str, tuple[Any, dict[str, str]]], op_name: str
) -> tuple[Any, dict[str, str]] | None:
candidates = [op_name]
if op_name in {"add", "sub", "mul", "div"}:
candidates.insert(0, f"{op_name}.Tensor")
for candidate in candidates:
if candidate in specs:
return specs[candidate]
return None
def _aot_default_value(value: Any) -> Any:
if value is None:
return None
if value == "true":
return True
if value == "false":
return False
if value == "{}":
return ()
if isinstance(value, str) and value.startswith("{") and value.endswith("}"):
return tuple(int(item.strip()) for item in value[1:-1].split(",") if item.strip())
try:
return int(value)
except (TypeError, ValueError):
try:
return float(value)
except (TypeError, ValueError):
return value
def _aot_add_adjoint(
builder: _AotNativeGraphBuilder,
adjoints: dict[Node, _AotNativeSymbol],
target: Any,
contribution: Any,
) -> bool:
if not isinstance(target, Node) or not isinstance(contribution, _AotNativeSymbol):
return contribution is None
previous = adjoints.get(target)
adjoints[target] = contribution if previous is None else builder.binary(
"add", previous, contribution
)
return True
def _build_aot_backward(
graph_module: GraphModule,
native_module: Any,
forward_lowering: _NativeLowering,
saved_nodes: list[Node],
runtime_values: dict[Node, Any],
runtime_inputs: list[Any],
public_node: Node,
) -> tuple[Any, list[int]] | None:
try:
specs = _aot_derivative_specs()
except (ImportError, ModuleNotFoundError, OSError):
# Derivative tooling/config unavailable: AOT lowering is optional;
# the caller falls back to the non-AOT native path.
return None
builder = _AotNativeGraphBuilder(native_module)
external_nodes = list(graph_module.graph.placeholders) + [
node for node in graph_module.graph.nodes if node.op == "get_attr"
]
external_symbols: list[_AotNativeSymbol] = []
forward_symbols: dict[Node, _AotNativeSymbol] = {}
for index, node in enumerate(external_nodes):
symbol = builder.input(runtime_inputs[index])
external_symbols.append(symbol)
forward_symbols[node] = symbol
saved_symbols: list[_AotNativeSymbol] = []
for node in saved_nodes:
actual = runtime_values.get(node)
if actual is None:
return None
symbol = builder.input(actual)
saved_symbols.append(symbol)
forward_symbols[node] = symbol
tangent = builder.input(runtime_values[public_node])
adjoints: dict[Node, _AotNativeSymbol] = {public_node: tangent}
batch_norm_cache: dict[tuple[int, ...], _AotNativeTuple] = {}
formula_env = _build_aot_formula_env(builder, batch_norm_cache=batch_norm_cache)
for node in reversed(graph_module.graph.nodes):
if node.op in {"placeholder", "get_attr", "output"}:
continue
grad = adjoints.get(node)
if grad is None:
continue
op_name = _target_name(node.target)
if op_name == "linear":
if len(node.args) not in {2, 3} or not all(
isinstance(item, Node) for item in node.args[:2]
):
return None
input_node, weight_node = node.args[:2]
bias_node = node.args[2] if len(node.args) == 3 else None
if bias_node is not None and not isinstance(bias_node, Node):
return None
input_value = forward_symbols[input_node]
weight_value = forward_symbols[weight_node]
weight_t = builder.unary("t", weight_value, shape=weight_value.shape[::-1])
input_grad = builder.helper(
"matmul_backward_self", (grad, input_value, weight_t), shape=input_value.shape
)
weight_t_grad = builder.helper(
"matmul_backward_other", (grad, input_value, weight_t), shape=weight_t.shape
)
if not _aot_add_adjoint(builder, adjoints, input_node, input_grad):
return None
if not _aot_add_adjoint(
builder, adjoints, weight_node, builder.unary("t", weight_t_grad, shape=weight_value.shape)
):
return None
if bias_node is not None:
dims = tuple(range(max(0, len(grad.shape) - 1)))
bias_grad = builder.sum(grad, dims, False) if dims else grad
if not _aot_add_adjoint(builder, adjoints, bias_node, bias_grad):
return None
continue
if op_name == "flatten":
if not node.args or not isinstance(node.args[0], Node):
return None
source_value = runtime_values.get(node.args[0])
if source_value is None:
return None
# the input storage. Keep that metadata-only dependency out of
# the saved-tensor list so the rebuilt graph can omit the pooled
# activation while still producing the exact reshape backward.
contribution = builder.reshape(grad, tuple(source_value.shape))
if not _aot_add_adjoint(builder, adjoints, node.args[0], contribution):
return None
continue
schema = _aot_schema_for(specs, op_name)
if schema is None:
return None
parsed, formulas = schema
if node.op == "call_method":
arg_values = (node.args[0],) if node.args else ()
elif op_name == "batch_norm":
names = (
"input", "running_mean", "running_var", "weight", "bias",
"training", "momentum", "eps",
)
arg_values = tuple(zip(names, node.args))
else:
arg_values = tuple(zip((arg.name for arg in parsed.args), node.args))
if op_name == "batch_norm":
context = {name: value for name, value in arg_values}
else:
context = dict(arg_values)
for arg in parsed.args:
if arg.name not in context and arg.default is not None:
context[arg.name] = _aot_default_value(arg.default)
context["grad"] = grad
# others only need metadata or their inputs (e.g. adaptive average
# pooling). A pruned saved-tensor set must not require a symbol for
# a result that the selected formula never reads.
context["result"] = forward_symbols.get(node)
tensor_params = {
name for name, value in context.items() if isinstance(value, Node)
}
env = dict(formula_env)
env.update(
{
name: forward_symbols.get(value) if isinstance(value, Node) else value
for name, value in context.items()
}
)
try:
for arg_name, formula in formulas.items():
target = context.get(arg_name)
if not isinstance(target, Node):
continue
translated = _aot_formula_python(formula, tensor_params)
contribution = eval(translated, {"__builtins__": {}}, env)
if not _aot_add_adjoint(builder, adjoints, target, contribution):
return None
except (KeyError, NameError, NotImplementedError, TypeError, ValueError, RuntimeError):
return None
grad_positions: list[int] = []
for index, node in enumerate(external_nodes):
actual = runtime_inputs[index]
if not getattr(actual, "requires_grad", False):
continue
contribution = adjoints.get(node)
if contribution is None:
continue
builder.graph.register_output(contribution.value)
grad_positions.append(index)
if not grad_positions:
return None
needed_saved_nodes = [
node
for node, symbol in zip(saved_nodes, saved_symbols)
if getattr(symbol.value, "use_count", 0) != 0
]
return builder.graph, grad_positions, needed_saved_nodes
class _AotNativeLowering:
def __init__(
self,
graph_module: GraphModule,
forward_graph: Any,
backward_graph: Any,
attribute_targets: list[str],
grad_positions: list[int],
) -> None:
self.graph_module = graph_module
self.forward_graph = forward_graph
self.backward_graph = backward_graph
self.placeholders = graph_module.graph.placeholders
self.attribute_targets = attribute_targets
self.grad_positions = list(grad_positions)
self.input_count = len(self.placeholders) + len(self.attribute_targets)
self._tensorplay_codegen = "stax-aot-native"
self._tensorplay_backward_codegen = "stax-aot-native"
lowering = self
from ..autograd import Function
class _AotAutogradFunction(Function):
@staticmethod
def forward(ctx: Any, *inputs: Any) -> Any:
outputs = lowering.forward_graph.execute(list(inputs))
ctx.save_for_backward(*inputs, *outputs[1:])
return outputs[0]
@staticmethod
def backward(ctx: Any, *grad_outputs: Any) -> tuple[Any, ...]:
grad_output = grad_outputs[0] if grad_outputs else None
if grad_output is None:
return (None,) * lowering.input_count
saved = list(ctx.saved_tensors)
outputs = lowering.backward_graph.execute([*saved, grad_output])
by_position = dict(zip(lowering.grad_positions, outputs))
return tuple(by_position.get(index) for index in range(lowering.input_count))
self._autograd_function = _AotAutogradFunction
def _bind_inputs(self, *args: Any, **kwargs: Any) -> list[Any]:
bound = self.graph_module.signature.bind_partial(*args, **kwargs)
bound.apply_defaults()
inputs = [
bound.arguments[node.target if isinstance(node.target, str) else node.name]
for node in self.placeholders
]
inputs.extend(self.graph_module._get_attr(target) for target in self.attribute_targets)
return inputs
def __call__(self, *args: Any, **kwargs: Any) -> Any:
inputs = self._bind_inputs(*args, **kwargs)
import tensorplay
if not tensorplay.is_grad_enabled() or not any(
getattr(value, "requires_grad", False) for value in inputs
):
return self.forward_graph.execute(inputs)[0]
return self._autograd_function.apply(*inputs)
def _lower_aot_native(
graph_module: GraphModule,
example_inputs: list[Any],
) -> _AotNativeLowering | None:
"""Build separate native forward/backward graphs at the AOT boundary."""
try:
import tensorplay
native_module = getattr(tensorplay._C, "_stax", None)
tensor_type = tensorplay.Tensor
except (AttributeError, ImportError):
return None
if native_module is None or not hasattr(native_module.Graph, "execute"):
return None
if len(example_inputs) != len(graph_module.graph.placeholders):
return None
if any(not isinstance(value, tensor_type) for value in example_inputs):
return None
if not tensorplay.is_grad_enabled():
return None
external_nodes = list(graph_module.graph.placeholders) + [
node for node in graph_module.graph.nodes if node.op == "get_attr"
]
attribute_targets = [node.target for node in external_nodes if node.op == "get_attr"]
runtime_inputs = list(example_inputs)
runtime_inputs.extend(graph_module._get_attr(target) for target in attribute_targets)
if not any(getattr(value, "requires_grad", False) for value in runtime_inputs):
return None
saved_nodes = [
node
for node in graph_module.graph.nodes
if node.op in {"call_function", "call_method"}
]
output_values = [
value for output in graph_module.graph.outputs for value in _nodes(output.args)
]
if len(output_values) != 1:
return None
public_node = output_values[0]
forward_lowering = _lower_native(
graph_module,
example_inputs,
use_fusion=False,
extra_output_nodes=saved_nodes,
)
if forward_lowering is None:
return None
if len(runtime_inputs) != len(forward_lowering.graph.inputs):
return None
# Training BatchNorm updates running buffers during forward. A compiler
# trace must not perform that update a second time; restore non-gradient
# capture path separates tracing state from the user execution state.
snapshots: list[tuple[Any, Any]] = []
seen_attributes: set[int] = set()
try:
for target in attribute_targets:
value = graph_module._get_attr(target)
if (
isinstance(value, tensor_type)
and not getattr(value, "requires_grad", False)
and id(value) not in seen_attributes
):
snapshots.append((value, value.detach().clone()))
seen_attributes.add(id(value))
with tensorplay.no_grad():
forward_outputs = forward_lowering.graph.execute(runtime_inputs)
except (AttributeError, RuntimeError, TypeError, ValueError):
return None
finally:
if snapshots:
with tensorplay.no_grad():
for value, snapshot in snapshots:
value.copy_(snapshot)
if len(forward_outputs) != 1 + len(saved_nodes):
return None
runtime_values: dict[Node, Any] = {public_node: forward_outputs[0]}
for index, node in enumerate(saved_nodes, start=1):
runtime_values[node] = forward_outputs[index]
built = _build_aot_backward(
graph_module,
native_module,
forward_lowering,
saved_nodes,
runtime_values,
runtime_inputs,
public_node,
)
if built is None:
return None
_, grad_positions, needed_saved_nodes = built
# The first graph is a shape/materialization graph. Rebuild the forward
# graph with only the values that the source-derived backward graph reads,
# intermediate until backward.
forward_lowering = _lower_native(
graph_module,
example_inputs,
use_fusion=False,
extra_output_nodes=needed_saved_nodes,
)
if forward_lowering is None:
return None
rebuilt = _build_aot_backward(
graph_module,
native_module,
forward_lowering,
needed_saved_nodes,
runtime_values,
runtime_inputs,
public_node,
)
if rebuilt is None:
return None
backward_graph, grad_positions, rebuilt_saved_nodes = rebuilt
if rebuilt_saved_nodes != needed_saved_nodes:
return None
return _AotNativeLowering(
graph_module,
forward_lowering.graph,
backward_graph,
attribute_targets,
grad_positions,
)
[docs]
def stax(
graph_module: GraphModule,
example_inputs: list[Any],
*,
mode: str | None = None,
options: dict[str, Any] | None = None,
name: str | None = None,
dynamic: bool | None = None,
strict_native: bool = False,
**kwargs: Any,
):
"""Compile one canonical graph and return an executable callable.
``example_inputs`` and backend options are part of the same contract as
metadata in the frontend and uses the native graph when its lowering
contract is satisfied. ``strict_native`` makes a failed lowering a hard
compiler error, so a benchmark can never report the Python GraphModule
executor as compiled performance.
"""
del name, kwargs
if mode not in {None, "default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"}:
raise RuntimeError(f"unknown Stax optimization mode: {mode!r}")
if options is not None:
if not isinstance(options, dict):
raise TypeError(f"options must be a dict, got {type(options)!r}")
unknown = set(options).difference(
{"stax.native", "stax.fusion", "stax.triton"}
)
if unknown:
raise RuntimeError(
f"Unexpected Stax optimization option(s): {sorted(unknown)!r}"
)
if any(not isinstance(value, bool) for value in options.values()):
raise RuntimeError("Stax optimization options must be bool values")
use_native = options is None or options.get("stax.native", True)
use_fusion = options is None or options.get("stax.fusion", True)
use_triton = options is None or options.get("stax.triton", True)
if use_native and use_fusion:
fused_cpu_graph = _lower_cpu_fused_pointwise(
graph_module,
example_inputs,
strict_native=strict_native,
dynamic=bool(dynamic is True),
)
if fused_cpu_graph is not None:
graph_module._stax_native_graph = fused_cpu_graph.graph
return fused_cpu_graph
# A region whose tail is a reduction folds the whole expression into
# the reduction loop: one pass over the input, no intermediate.
fused_cpu_reduction = _lower_cpu_fused_reduction(
graph_module,
example_inputs,
strict_native=strict_native,
dynamic=bool(dynamic is True),
)
if fused_cpu_reduction is not None:
graph_module._stax_codegen = "stax-fused-cpu-reduce"
return fused_cpu_reduction
# A region whose reductions sit in the middle stages per row: each
# reduction folds the row to one value that the following work reads
# as a broadcast, so the region still reads its inputs once.
row_fused_cpu = _lower_cpu_row_fusion(
graph_module,
example_inputs,
strict_native=strict_native,
dynamic=bool(dynamic is True),
)
if row_fused_cpu is not None:
graph_module._stax_codegen = "stax-fused-cpu-rowfuse"
return row_fused_cpu
if use_native and use_triton:
# Keep Triton optional and lazy. Importing tensorplay on a CPU-only
# machine must not import Triton or its compiler toolchain.
try:
first = example_inputs[0]
is_cuda = first.device.is_cuda()
except (AttributeError, IndexError):
is_cuda = False
if is_cuda and use_fusion:
# A region whose reductions sit in the middle keeps its row
# resident across every stage; splitting it at each reduction
# would stream the input once per stage instead.
row_fused_cuda = _lower_cuda_row_fusion(
graph_module,
example_inputs,
strict_native=strict_native,
dynamic=bool(dynamic is True),
)
if row_fused_cuda is not None:
graph_module._stax_codegen = "stax-fused-cuda-rowfuse"
return row_fused_cuda
if is_cuda:
from .codegen.triton import (
compile_graph_module as compile_triton_graph,
)
triton_graph = compile_triton_graph(
graph_module,
example_inputs,
mode=mode,
strict_native=strict_native,
)
if triton_graph is not None:
graph_module._stax_codegen = "triton"
return triton_graph
if use_native and getattr(graph_module.root, "training", False):
aot_graph = _lower_aot_native(graph_module, example_inputs)
if aot_graph is not None:
graph_module._stax_native_graph = aot_graph.forward_graph
return aot_graph
if strict_native and any(
getattr(graph_module._get_attr(node.target), "requires_grad", False)
for node in graph_module.graph.nodes
if node.op == "get_attr"
):
raise RuntimeError(
"AOT backward graph for the captured training region"
)
native_graph = (
_lower_native(graph_module, example_inputs, use_fusion=use_fusion)
if use_native
else None
)
if native_graph is not None:
graph_module._stax_native_graph = native_graph.graph
return native_graph
if use_native and use_fusion:
# Nothing claimed the region whole. Its fusible runs are still worth
# compiling: each becomes one kernel, and the operators between them
# run as captured instead of the region losing every compiled route.
segmented = _lower_cpu_segmented(
graph_module,
example_inputs,
strict_native=strict_native,
dynamic=bool(dynamic is True),
)
if segmented is not None:
graph_module._stax_codegen = "stax-fused-cpu-segments"
return segmented
if strict_native:
raise RuntimeError(
"strict_native Stax lowering failed: captured graph has no native executable"
)
# No native executable exists for this graph (scalar placeholders,
# factory-only regions, unsupported surface). Fall back to the
# generated Python executor so the region still runs with captured
# semantics instead of failing to compile.
return graph_module.recompile()Help improve this page
Found an error, an unclear step, or a missing example?
Was this page helpful?