from __future__ import annotations import contextlib import dataclasses import functools import itertools import logging import operator import os import textwrap import traceback from collections.abc import Container, Generator, Iterable, Iterator, Sequence from contextlib import AbstractContextManager, nullcontext from enum import Enum from functools import partial from typing import ( Any, Callable, cast, ClassVar, Literal, Optional, overload, SupportsFloat, SupportsInt, TYPE_CHECKING, TypeVar, Union, ) from typing_extensions import ( assert_never, Never, override, ParamSpec, Self, TypeAlias, TypeIs, ) from unittest.mock import patch import sympy from sympy import Expr, Integer, Symbol import torch._export.serde.schema as export_schema import torch._library.utils as library_utils import torch._logging import torch.fx import torch.utils._pytree as pytree from torch._dynamo.utils import identity from torch._export.serde.serialize import GraphModuleSerializer from torch._higher_order_ops.auto_functionalize import can_auto_functionalize from torch._inductor import metrics from torch._inductor.utils import get_free_symbols from torch._prims_common import ( compute_required_storage_length, is_boolean_dtype, is_float_dtype, make_channels_last_strides_for, StrideType, ) from torch._subclasses.fake_tensor import get_schema_info from torch.fx.experimental.symbolic_shapes import ( _remove_effect_token_unbacked_bindings, compute_unbacked_bindings, free_symbols, free_unbacked_symbols, rebind_unbacked, resolve_unbacked_bindings, ShapeEnv, SymTypes, ) from torch.fx.node import Node from torch.utils._ordered_set import OrderedSet from torch.utils._sympy.functions import CleanDiv, FloorDiv, ModularIndexing from torch.utils._sympy.symbol import SymT from . import config, dependencies from .codegen.common import ( BackendFeature, CodegenSymbol, get_scheduling_for_device, index_prevent_reordering, Kernel, ) from .dependencies import ( Dep, extract_free_symbols, extract_input_node_reduction_ranges, extract_read_writes, var_builder, ) from .loop_body import LoopBody from .ops_handler import OpCounterCSE, OpCountResult, ReductionType, StoreMode from .runtime.benchmarking import benchmarker from .runtime.hints import DeviceProperties, ReductionHint from .utils import ( argsort, argsort_sym, cache_on_self, ceildiv, convert_shape_to_inductor, convert_shape_to_symint, developer_warning, do_bench_using_profiling, dtype_from_size, get_dtype_size, get_kernel_metadata, GPU_ALIGN_BYTES, ir_dataclass, is_dynamic, is_gpu, sympy_dot, sympy_index_symbol, sympy_index_symbol_with_prefix, sympy_product, sympy_subs, tensor_is_aligned, ) from .virtualized import ops, OpsValue, V if TYPE_CHECKING: from torch._library.fake_class_registry import FakeScriptObject from torch.fx.experimental.symbolic_shapes import SympyBoolean from torch.fx.node import Argument from .codegen.cuda.cuda_template import CUDATemplate from .codegen.wrapper import PythonWrapperCodegen from .graph import GraphLowering from .utils import IndentedBuffer else: CUDATemplate: TypeAlias = object try: import triton triton_version = triton.__version__ has_triton = True except ImportError: triton_version = None has_triton = False _P = ParamSpec("_P") _T = TypeVar("_T") _U = TypeVar("_U") _V = TypeVar("_V") _IntLike: TypeAlias = Union[int, Expr] _NumLike: TypeAlias = Union[int, float, Expr] _OpOverloads: TypeAlias = Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator] log = logging.getLogger(__name__) indent = functools.partial(textwrap.indent, prefix=" ") aten = torch.ops.aten autotune_warmup = int(os.getenv("TORCH_AUTOTUNE_WARMUP", 25)) autotune_rep = int(os.getenv("TORCH_AUTOTUNE_REP", 100)) """ [Note: Inductor IR] Inductor's IR is produced by executing 'lowering' code (see lowering.py). Each lowering is registered to a particular aten operator, and expects inputs that correspond to the aten schema. However, in place of torch Tensor inputs, lowerings expect Inductor TensorBox inputs. TensorBox IR represents torch tensors. Tensors are sometimes single objects owning storage, and sometimes views of another Tensor's storage. Mutating tensor operations (such as add_()) affect the underlying storage and any associated views. Other operations (such as .t_()) update metadata about the current view but don't modify the underlying storage. To model this in Inductor, the IR distinguishes between TensorBox, View, StorageBox and Buffer. TensorBox is the top level IR construct that any lowering should produce and maps to a torch.Tensor output from an operation. But just as torch.Tensors take different forms, TensorBox IR can reference View IR or directly reference StorageBox IRs. Some Inductor lowerings produce new sets of 'Box'es, while others (such as .t() or other view ops) may take an existing TensorBox and point it to a new underlying View IR. Tensors that directly own storage are represented as a chain of: TensorBox -> StorageBox -> Buffer where Buffer is a simple (1D) allocation, and StorageBox introduces the concept of a Layout. If you mutate the data of such a tensor, we swing the StorageBox pointer to point to a new buffer (leaving the old buffer unmodified and functionalizing the operation). Tensors backed by views add one more indirection to the IR. TensorBox -> View -> StorageBox -> Buffer In these cases, the underlying StorageBox/Buffer will be shared with the pre-view TensorBox. Computation is represented by Operation nodes, with each operation producing 1 or more output Buffers. In the case of mutations, these will be new Buffers that have the mutated buffer listed in its get_mutation_names(). It is also possible to have an InputBuffer for which there is no corresponding Operation, e.g. it may be a graph input or compile time constant. """ _NodeOrNodes: TypeAlias = Union[ int, "TensorBox", dict[str, "TensorBox"], "Symbol", "IRNode", Sequence[ Optional[Union[int, dict[str, "TensorBox"], "TensorBox", "Symbol", "IRNode"]] ], ] def _is_static(x: object) -> bool: return isinstance(x, (int, Integer)) @dataclasses.dataclass(frozen=True) class GraphPartitionSignature: # symbol inputs that are necessary for codegen symbol_inputs: OrderedSet[sympy.Symbol] # mapping from partition input name to IRNode or Expr. Need the name str since # we cannot get name from Expr. input_nodes: dict[str, Union[IRNode, sympy.Expr, TorchBindObject]] output_nodes: list[IRNode] # mapping from partition input name to a boolean for whether deallocating it # in the partition function input_deallocation: dict[str, bool] skip_cudagraph: bool # name of constants read/written by the graph partition constant_names: list[str] def validate_ir(node_or_nodes: Optional[_NodeOrNodes]) -> None: def _check_tensorbox(nodes: Optional[_NodeOrNodes]) -> None: # Could expand this to check deeper properties # (e.g. TensorBox points to View or StorageBox) if nodes is None: pass elif isinstance(nodes, (list, tuple)): for node in nodes: _check_tensorbox(node) elif isinstance(nodes, dict): for node in nodes.values(): _check_tensorbox(node) else: assert isinstance( nodes, ( ExpandView, DynamicScalar, AssertScalar, TensorBox, sympy.logic.boolalg.Boolean, Expr, int, EffectfulKernel, ShapeAsConstantBuffer, ), ), ( f"Found {type(nodes)}, which is not a supported top level IR node. See [Note: Inductor IR]" ) # Be picky about the accepted data structure (don't use pytree here) _check_tensorbox(node_or_nodes) def ops_wrapper(name: str) -> Callable[..., OpsValue]: assert isinstance(name, str), type(name) def fn(*args: object, **kwargs: object) -> OpsValue: return getattr(ops, name)(*args, **kwargs) return fn def inverse_reorder(order: Sequence[int]) -> Callable[[Sequence[_T]], Sequence[_T]]: inv_order = dict(zip(order, range(len(order)))) def reindex(index: Sequence[_T]) -> Sequence[_T]: assert len(index) == len(inv_order) return [index[inv_order[i]] for i in range(len(index))] return reindex def same_reorder(order: Sequence[int]) -> Callable[[Sequence[_T]], Sequence[_T]]: def reindex(index: Sequence[_T]) -> Sequence[_T]: assert len(index) == len(order) return [index[order[i]] for i in range(len(index))] return reindex def fuse_reindexing( reindex1: Callable[[Sequence[_U]], Sequence[_V]], reindex2: Callable[[Sequence[_T]], Sequence[_U]], ) -> Callable[[Sequence[_T]], Sequence[_V]]: def reindex(index: Sequence[_T]) -> Sequence[_V]: return reindex1(reindex2(index)) return reindex NHWC_STRIDE_ORDER = [3, 0, 2, 1] NHWDC_STRIDE_ORDER = [4, 0, 3, 2, 1] def get_fill_order( seq: Sequence[Union[int, torch.SymInt, Expr]], shape_env: Optional[ShapeEnv] = None ) -> Sequence[int]: """ Convert strides to fill order (argsort) """ if shape_env is None or all(isinstance(s, (int, sympy.Integer)) for s in seq): sorted_idx: Sequence[int] = argsort(seq) else: # argsort_sym handles unbacked symints (with the help of the shape_env) sorted_idx = argsort_sym(shape_env, seq) return sorted_idx def stride_order2fill_order(order: Sequence[Union[int, Integer]]) -> Sequence[int]: """ Convert stride order to fill order For channel last format, stride order = [3, 0, 2, 1] and fill order = [1, 3, 2, 0] """ lookup = {pos: idx for idx, pos in enumerate(order)} fill_order = [lookup[i] for i in range(len(order))] return fill_order def get_stride_order( seq: Sequence[Union[int, torch.SymInt, Expr]], shape_env: Optional[ShapeEnv] = None ) -> Sequence[int]: """ Convert strides to stride order """ sorted_idx: Sequence[int] = get_fill_order(seq, shape_env) out = [0 for _ in range(len(seq))] for i, elem in enumerate(sorted_idx): out[elem] = i return out @overload def ir_node_to_tensor(x: Literal[None], guard_shape: bool = True) -> None: ... @overload def ir_node_to_tensor(x: IRNode, guard_shape: bool = True) -> torch.Tensor: ... def ir_node_to_tensor( x: Optional[IRNode], guard_shape: bool = True ) -> Optional[torch.Tensor]: if x is None: return None shape_fn: Callable[[Union[int, Expr]], Union[int, Expr]] if not guard_shape: shape_fn = V.graph.sizevars.size_hint else: shape_fn = identity size = [shape_fn(s) for s in x.get_size()] stride: StrideType if is_storage_and_layout(x): stride = [shape_fn(s) for s in x.get_layout().stride] else: stride = FlexibleLayout.contiguous_strides(size) dtype = x.get_dtype() device = x.get_device() size = convert_shape_to_symint(size) stride = convert_shape_to_symint(stride) with V.graph.sizevars.shape_env.suppress_guards(): t = torch.empty_strided( size=size, stride=stride, dtype=dtype, device=device ).zero_() return t def may_convert_to_optional( value: Optional[Sequence[_T]], ) -> Optional[Sequence[Optional[_T]]]: if isinstance(value, list) and not value: # [None] makes sure the cpp wrapper codegen will generate something like # {std::nullopt} instead of {} return [None] return value def get_device_type( x: Union[IRNode, OutputSpec, torch.device, None, str], ) -> Optional[str]: if isinstance(x, str) or x is None: return x elif isinstance(x, torch.device): return x.type elif isinstance(x, (IRNode, OutputSpec)): return get_device_type(x.get_device()) assert_never(f"get_device_type({x}: {type(x).__name__})") def is_triton(x: Union[IRNode, torch.device, None, str]) -> bool: device = get_device_type(x) # Special case cpu and cuda as using the method below # to determine if the scheduler is a triton scheduler subclass # requires instantiating a scheduler for them if device in ["cpu", "cuda"]: if getattr(config, f"{device}_backend") == "triton": return True return False if ( device is None or (device_scheduling := get_scheduling_for_device(device)) is None ): return False from .codegen.triton import TritonScheduling assert isinstance(device_scheduling, type), type(device_scheduling) return issubclass(device_scheduling, TritonScheduling) def is_cpu(x: Union[IRNode, torch.device, None, str]) -> bool: return get_device_type(x) == "cpu" def is_aligned_realized_tensor_hint( x: Union[Buffer, TensorBox], alignment: int ) -> bool: # Use this as a hint. This won't guard since size_hint doesn't guard. if ( not isinstance(x, IRNode) or x.maybe_get_stride() is None or free_unbacked_symbols(x.get_stride()) or free_unbacked_symbols(x.get_size()) ): return False aligned_strides = all( (V.graph.sizevars.size_hint_or_throw(x.get_stride()[i]) % alignment) == 0 for i in range(len(x.get_stride()) - 1) ) # if the last dim size is <= 1, stride doesn't matter aligned_last_dim = ( V.graph.sizevars.size_hint_or_throw(x.get_stride()[-1]) == 1 or V.graph.sizevars.size_hint_or_throw(x.get_size()[-1]) <= 1 ) return aligned_last_dim and aligned_strides def significant_strides_equal( strides1: Sequence[_IntLike], strides2: Sequence[_IntLike], shape: Sequence[_IntLike], ) -> bool: """ Returns true if the strides are equal, ignoring dimensions of size 1 . """ assert len(shape) == len(strides1) and len(strides1) == len(strides2) for dim, s1, s2 in zip(shape, strides1, strides2): if V.graph.sizevars.statically_known_leq(dim, 1): continue if not V.graph.sizevars.statically_known_equals( s1, s2 ) and not V.graph.sizevars.symbolic_hint(s1) == V.graph.sizevars.symbolic_hint( s2 ): return False return True def try_match_insignificant_strides( tensor: IRNode, strides: Sequence[Union[int, torch.SymInt]], ) -> IRNode: """ Tries to match the strides of the tensor to those in the meta_strides. Strides of insignificant dimensions - size 0 or 1 - will be updated. If there are real stride differences (NHWC vs NCHW), or the tensor is not realized, then the input will be returned """ if not is_storage_and_layout(tensor): return tensor if all( V.graph.sizevars.statically_known_equals(s1, s2) for s1, s2 in zip(strides, tensor.get_stride()) ): return tensor if not significant_strides_equal(strides, tensor.get_stride(), tensor.get_size()): return tensor storage, old_layout = as_storage_and_layout(tensor) new_stride = [*old_layout.stride] for i, s in enumerate(tensor.get_size()): if V.graph.sizevars.statically_known_leq(s, 1): new_stride[i] = strides[i] new_layout = FixedLayout( old_layout.device, old_layout.dtype, old_layout.size, new_stride, old_layout.offset, old_layout.is_pinned, ) return TensorBox(ReinterpretView(data=storage, layout=new_layout)) def gm_original_output_strides(gm: torch.fx.GraphModule) -> None: output_node = gm.graph.find_nodes(op="output")[0] output_node.meta["user_visible_output_idxs"] = [ idx for idx, _ in enumerate(output_node.args) ] from torch._inductor.compile_fx import record_original_output_strides record_original_output_strides(gm) def get_symbolic_inputs(inputs: Sequence[IRNode]) -> list[Expr]: sym_vars: OrderedSet[Expr] = OrderedSet() for inp in inputs: sym_vars |= get_free_symbols(inp.get_size(), unbacked_only=False) sym_vars |= get_free_symbols(inp.get_stride(), unbacked_only=False) return list(sym_vars) class IRNode: """Base class for all intermediate representation (IR) nodes in TorchInductor. Note: This is an abstract base class. Most methods raise NotImplementedError and must be overridden by concrete subclasses. """ _current_origins: ClassVar[OrderedSet[Any]] = OrderedSet() # NB: These are kinda weird, origins: OrderedSet[Any] = dataclasses.field(init=False) # traces back to where the IRNode is created in Inductor traceback: Optional[list[str]] = dataclasses.field(init=False) origin_node: Optional[torch.fx.Node] = dataclasses.field(init=False) @staticmethod @contextlib.contextmanager def current_origins(origins: OrderedSet[Node]) -> Generator[None, None, None]: old = IRNode._current_origins IRNode._current_origins = old | origins try: yield finally: IRNode._current_origins = old @staticmethod def is_realized_node(node: IRNode) -> bool: return isinstance( node, ( ComputedBuffer, InputsKernel, InputBuffer, ReinterpretView, TemplateBuffer, ), ) def _post_init_setattr(self, attr: str, value: Any) -> None: # Intended for use in __post_init__ for enforcing an invariant on a dataclass # If you must, can also be used for setting provenance info # We would like to try and minimize these usages though object.__setattr__(self, attr, value) def __post_init__(self) -> None: origins = OrderedSet(self._current_origins) self._post_init_setattr("origins", origins) self._post_init_setattr( "traceback", traceback.format_stack() if config.debug_ir_traceback else None ) self._post_init_setattr("origin_node", None) def get_read_names(self) -> OrderedSet[str]: return OrderedSet(dep.name for dep in self.get_reads()) def get_traceback(self) -> Optional[list[str]]: return self.traceback def get_origin_node(self) -> Optional[torch.fx.Node]: return self.origin_node def get_defining_op(self) -> Optional[Operation]: return None def get_stack_traces(self) -> OrderedSet[str]: # Return stack traces to user model code # A single IRNode could correspond to multiple lines of code stack_traces: OrderedSet[str] = OrderedSet() origins = self.origins if isinstance(self, ExternKernel): origin_node = self.get_origin_node() if self.origin_node: origins = OrderedSet([origin_node]) for node in origins: if hasattr(node, "stack_trace") and node.stack_trace: # nodes in the backward graph don't have mapping to pre_grad_graph stack_traces.add(node.stack_trace) else: pre_grad_nodes = ( torch._inductor.debug._inductor_post_to_pre_grad_nodes.get( "postToPre", {} ).get(node.name, []) ) if not isinstance(pre_grad_nodes, list): continue for node_name in pre_grad_nodes: stack_trace = ( torch._inductor.debug._inductor_pre_grad_node_stack_trace.get( node_name, None ) ) if stack_trace: stack_traces.add(stack_trace) return stack_traces def common_repr(self, shorten: bool = True) -> Sequence[str]: origins = f"origins={getattr(self, 'origins', '')}" if shorten and len(origins) > 64: # this can get *very* long origins = f"{origins[:61]}..." if not self.get_stack_traces(): return [origins] stack_trace_str = [] for stack_trace in self.get_stack_traces(): stack_trace_str.append("stack_traces = {") stack_trace_str += stack_trace.split("\n") stack_trace_str.append("}") return [origins] + stack_trace_str def str_helper( self, lines: Sequence[object], shorten: bool = True, multiline: bool = True ) -> str: lines = list(lines) + list(self.common_repr(shorten)) lines = list(map(str, lines)) if multiline: new_lines = indent(",\n".join(lines)) return f"{type(self).__name__}(\n{new_lines}\n)" else: return f"{type(self).__name__}({lines})" def get_dtype(self) -> torch.dtype: return self.dtype def maybe_get_dtype(self) -> Optional[torch.dtype]: try: return self.get_dtype() except NotImplementedError: return None def get_layout(self) -> Layout: raise NotImplementedError(f"get_layout() is not implemented by {type(self)}!") def maybe_get_layout(self) -> Optional[Layout]: try: return self.get_layout() except NotImplementedError: return None def get_output_spec(self) -> OutputSpec: return self.get_layout() def maybe_get_output_spec(self) -> Optional[OutputSpec]: try: return self.get_output_spec() except NotImplementedError: return None def has_tensor_output(self) -> bool: """True for single tensor output (excludes MultiOutput)""" return isinstance(self.maybe_get_output_spec(), Layout) def get_size(self) -> Sequence[Expr]: raise NotImplementedError(f"get_size() is not implemented by {type(self)}!") def maybe_get_size(self) -> Optional[Sequence[_IntLike]]: try: return self.get_size() except NotImplementedError: return None @property def shape(self) -> Union[_IntLike, sympy.Rel, Sequence[_IntLike]]: return self.get_size() def get_numel(self) -> Expr: return sympy_product(self.get_size()) def is_zero_elements(self) -> bool: return V.graph.sizevars.statically_known_true(sympy.Eq(self.get_numel(), 0)) def realize(self) -> Optional[str]: """ If the IRNode refers to data which has not been materialized (e.g., it is a Pointwise/Reduction that could potentially have more compute fused into it), realize the IRNode into physical memory, ending the possibility of fusing into it, but allowing, e.g., multiple users to access the data without having to recompute. Check StorageBox.realize for a particularly notable implementation. TODO(ezyang): I think, in principle, every IRNode should have an implementation of this, and most of the time no-op is OK, but you really do have to audit each IRNode for this, so for now, raise an error if it's not implemented. Note that some code in graph.py will catch this thrown error and suppress it with a warning. """ raise NotImplementedError(f"realize NYI on {type(self)}") def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str: raise NotImplementedError(f"codegen_reference NYI on {type(self)}") def get_device(self) -> Optional[torch.device]: return None def get_device_or_error(self) -> torch.device: device = self.get_device() assert device is not None return device def has_exceeded_max_reads(self) -> bool: return False def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]: raise NotImplementedError(type(self).__name__) def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]: raise NotImplementedError(type(self).__name__) def get_stride(self) -> Sequence[_IntLike]: raise NotImplementedError(type(self).__name__) def maybe_get_stride(self) -> Optional[Sequence[_IntLike]]: try: return self.get_stride() except NotImplementedError: return None def get_name(self) -> str: raise NotImplementedError(type(self).__name__) def maybe_get_name(self) -> Optional[str]: try: return self.get_name() except NotImplementedError: return None def is_input_buffer(self) -> bool: try: return self.get_name() in V.graph.graph_inputs except NotImplementedError: return False def has_large_inner_fn(self, threshold: Optional[int] = None) -> bool: return False def mark_reuse(self, users: int) -> None: pass def realize_hint(self) -> None: pass def unwrap_view(self) -> IRNode: raise NotImplementedError(type(self).__name__) def freeze_layout(self) -> None: raise NotImplementedError(type(self).__name__) def freeze_layout_with_stride_order( self, order: Sequence[int], allow_padding: bool = False ) -> None: raise NotImplementedError(type(self).__name__) def freeze_layout_with_fill_order(self, order: Sequence[int]) -> None: raise NotImplementedError(type(self).__name__) def freeze_layout_with_same_order(self, stride: Sequence[_IntLike]) -> None: raise NotImplementedError(type(self).__name__) def freeze_layout_with_exact_strides( self, exact_strides: Sequence[_IntLike], allow_padding: bool = False ) -> None: raise NotImplementedError(type(self).__name__) def get_read_writes(self) -> dependencies.ReadWrites: raise NotImplementedError(type(self).__name__) def get_reads(self) -> OrderedSet[Dep]: return self.get_read_writes().reads def num_reads(self) -> int: return len(self.get_reads()) def get_storage_numel(self) -> _IntLike: raise NotImplementedError(type(self).__name__) def get_free_symbol_uses( self, unbacked_only: bool = False ) -> OrderedSet[sympy.Symbol]: raise NotImplementedError(type(self).__name__) def get_reduction_type(self) -> Optional[str]: raise NotImplementedError(type(self).__name__) def get_reduction_size(self) -> Sequence[Expr]: raise NotImplementedError(type(self).__name__) def is_extern(self) -> bool: return False def is_no_op(self) -> bool: return False def constant_to_device(self, device: torch.device) -> IRNode: raise NotImplementedError(type(self).__name__) def get_mutation_names(self) -> Sequence[str]: raise NotImplementedError(type(self).__name__) def get_operation_name(self) -> str: raise NotImplementedError(type(self).__name__) def get_inputs_that_alias_output(self) -> Sequence[str]: raise NotImplementedError(type(self).__name__) if TYPE_CHECKING: @property def dtype(self) -> torch.dtype: ... @ir_dataclass(frozen=False) class Operation: def __post_init__(self) -> None: self.operation_name: Optional[str] = None def get_device(self) -> Optional[torch.device]: raise NotImplementedError def get_origin_node(self) -> Optional[torch.fx.Node]: assert hasattr(self, "origin_node") return self.origin_node def get_origins(self) -> OrderedSet[Any]: assert hasattr(self, "origins") return self.origins def get_operation_name(self) -> str: assert self.operation_name is not None return self.operation_name def is_extern(self) -> bool: return False def is_no_op(self) -> bool: return False def get_read_writes(self) -> dependencies.ReadWrites: raise NotImplementedError def is_user_of(self, name: str) -> bool: return name in self.get_read_names() def get_read_names(self) -> OrderedSet[str]: return OrderedSet(dep.name for dep in self.get_reads()) def get_reads(self) -> OrderedSet[Dep]: return self.get_read_writes().reads def get_outputs(self) -> list[Buffer]: raise NotImplementedError def get_unbacked_symbol_defs(self) -> OrderedSet[sympy.Symbol]: return OrderedSet() def get_free_symbol_uses( self, unbacked_only: bool = False ) -> OrderedSet[sympy.Symbol]: """ When unbacked_only=True: Returns the unbacked symbols which are required to be in scope in order to successfully perform codegen for this buffer. For example, a buffer that corresponds to an extern kernel call that takes i0 as an argument would return {i0} here. This is used to generate necessary dependencies that ensure we actually bind i0 in codegen before you try to use it. Note that this is NOT transitive; in particular, if this buffer takes in as input another buffer with dynamic shape (e.g., (i0,)), we will not report it here, because you will already have a dependency on that buffer, which will eventually have a dependency on i0 if necessary. When unbacked_only=False: Similar to `unbacked_only=True` but including all free symbols instead of only free unbacked symbols. """ return OrderedSet() def get_workspace_size(self) -> int: """ Gets extra global memory size needed by this buffer. Some algorithms (e.g. group gemm) may require extra global memory in the generated code. """ return 0 @ir_dataclass class Loops(IRNode): device: torch.device dtype: torch.dtype inner_fn: Callable[..., Any] ranges: Sequence[_IntLike] def get_free_symbol_uses( self, unbacked_only: bool = False ) -> OrderedSet[sympy.Symbol]: return OrderedSet().union( *(get_free_symbols(e, unbacked_only) for e in self.ranges), self.inner_fn_free_symbols(unbacked_only), ) def _to_str(self, names: Sequence[str]) -> str: return self.str_helper( [ f"'{self.device.type}'", str(self.dtype), self.inner_fn_str(), ] + [f"{name}={getattr(self, name)}" for name in names] + [f"origin_node={self.origin_node!r}"] ) def __post_init__(self) -> None: super().__post_init__() def __str__(self) -> str: return self._to_str(("ranges",)) __repr__ = __str__ def get_device(self) -> Optional[torch.device]: return self.device def get_origin_node(self) -> Optional[torch.fx.Node]: return self.origin_node def get_size(self) -> Sequence[Expr]: return self.ranges def get_pointwise_size(self) -> Sequence[Expr]: return self.ranges @classmethod def create( cls, *args: Any, **kwargs: Any ) -> Union[TensorBox, ShapeAsConstantBuffer]: origin_node = kwargs.pop("origin_node", None) tb = kwargs.pop("traceback", None) r = cls(*args, **kwargs) # Need to explicitly set origin_node here to propagate it down. # todo(chilli): I think it would be better for IRNode to directly set # origin_node r._post_init_setattr("origin_node", origin_node) r._post_init_setattr("traceback", tb or r.traceback) return TensorBox.create(r) @staticmethod def _index(ranges: Sequence[_IntLike], prefix: SymT = SymT.INDEX) -> Sequence[Expr]: return [ sympy.S.Zero if s == 1 else sympy_index_symbol_with_prefix(prefix, n) for n, s in enumerate(ranges) ] @cache_on_self def inner_fn_opcount(self) -> OpCountResult: opcounter = OpCounterCSE(V.MockHandler()) with ( V.set_ops_handler(opcounter), patch.object(FlexibleLayout, "allow_indexing", True), ): self.inner_fn(*self.inner_fn_args()) return opcounter.getvalue() def inner_fn_args(self) -> Sequence[Sequence[_IntLike]]: return (self._index(self.ranges),) @cache_on_self def inner_fn_str(self) -> str: return V.KernelFormatterHandler.ir_to_string( self.inner_fn, *self.inner_fn_args() ) def has_large_inner_fn(self, threshold: Optional[int] = None) -> bool: if threshold is None: threshold = 0 threshold = max(threshold, config.realize_opcount_threshold) return self.inner_fn_opcount().num_ops > threshold def inner_fn_free_symbols(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: index = self._index(self.ranges) return extract_free_symbols(self.inner_fn, index, unbacked_only=unbacked_only) def get_reads(self) -> OrderedSet[Dep]: with patch.object(FlexibleLayout, "allow_indexing", True): if self.get_reduction_type(): return extract_read_writes( self.make_loader(), self.get_size(), self.get_reduction_size(), ).reads else: return extract_read_writes( self.make_loader(), self.get_size(), ).reads def get_read_names(self) -> OrderedSet[str]: return OrderedSet(self.inner_fn_opcount().read_buffers) def num_reads(self) -> int: return len(self.inner_fn_opcount().read_buffers) def get_reduction_size(self) -> Sequence[Expr]: raise NotImplementedError( f"get_reduction_size() is not implemented by {type(self)}!" ) def get_reduction_type(self) -> Optional[str]: raise NotImplementedError( f"get_reduction_type() is not implemented by {type(self)}!" ) def constant_to_device(self, device: torch.device) -> IRNode: raise NotImplementedError( f"constant_to_device() is not implemented by {type(self)}!" ) def nop_loader_fn(idx: Union[Expr, Sequence[Expr]], *, dtype: torch.dtype) -> OpsValue: if dtype.is_floating_point: return ops.constant(float("nan"), dtype) else: return ops.constant(0, dtype) @ir_dataclass class Pointwise(Loops): def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]: # Make zero-element loops into a no-op if self.is_zero_elements(): return partial(nop_loader_fn, dtype=self.dtype) return self.inner_fn def get_reduction_size(self) -> Sequence[sympy.Expr]: return [] def get_reduction_type(self) -> Optional[str]: return None def store_output( self, output_name: Optional[str], indexer: Callable[[Sequence[Expr]], Never], vars: Sequence[Expr], ) -> None: loader = self.make_loader() return ops.store(output_name or "unnamed", indexer(vars), loader(vars)) def constant_to_device(self, device: torch.device) -> IRNode: """Move this to a given device. Requires that all reads are to constants.""" loader = self.make_loader() loader = patch.object(ConstantBuffer, "override_device", device)(loader) return Pointwise( device=device, dtype=self.dtype, inner_fn=loader, ranges=self.ranges, ) @ir_dataclass class Scatter(Pointwise): output_indexer: Callable[[Sequence[Expr]], Expr] scatter_mode: StoreMode = None def constant_to_device(self, device: torch.device) -> IRNode: """Move this to a given device. Requires that all reads are to constants.""" loader = self.make_loader() loader = patch.object(ConstantBuffer, "override_device", device)(loader) return Scatter( device=device, dtype=self.dtype, inner_fn=loader, ranges=self.ranges, output_indexer=self.output_indexer, scatter_mode=self.scatter_mode, ) def store_output( self, output_name: Optional[str], indexer: Callable[[Sequence[Expr]], Never], vars: Sequence[Expr], ) -> Any: loader = self.make_loader() if output_name is None: output_name = "unnamed" return ops.store( output_name, indexer(self.output_indexer(vars)), loader(vars), mode=self.scatter_mode, ) REDUCTION_COMBINE_FN: dict[str, Callable[..., OpsValue]] = { "any": ops_wrapper("logical_or"), "max": ops_wrapper("maximum"), "min": ops_wrapper("minimum"), "prod": ops_wrapper("mul"), "sum": ops_wrapper("add"), "xor_sum": ops_wrapper("bitwise_xor"), } def get_reduction_combine_fn( reduction_type: str, dtype: torch.dtype, arg_break_ties_left: bool = True ) -> Callable[..., object]: if reduction_type in REDUCTION_COMBINE_FN: return REDUCTION_COMBINE_FN[reduction_type] elif reduction_type in ("argmax", "argmin"): def argmax_combine_fn( a: tuple[object, object], b: tuple[object, object] ) -> tuple[OpsValue, OpsValue]: a_value, a_index = a b_value, b_index = b if reduction_type == "argmin": mask = ops.lt(a_value, b_value) else: mask = ops.gt(a_value, b_value) equal = ops.eq(a_value, b_value) if is_float_dtype(dtype): a_isnan = ops.ne(a_value, a_value) b_isnan = ops.ne(b_value, b_value) mask = ops.logical_or(mask, ops.gt(a_isnan, b_isnan)) equal = ops.logical_or(equal, ops.logical_and(a_isnan, b_isnan)) tie = ( ops.lt(a_index, b_index) if arg_break_ties_left else ops.gt(a_index, b_index) ) mask = ops.logical_or(mask, ops.logical_and(equal, tie)) return ( ops.where(mask, a_value, b_value), ops.where(mask, a_index, b_index), ) return argmax_combine_fn elif reduction_type == "welford_combine": def welford_combine_fn( a: tuple[OpsValue, OpsValue, OpsValue], b: tuple[OpsValue, OpsValue, OpsValue], ) -> tuple[OpsValue, OpsValue, OpsValue]: a_mean, a_m2, a_weight = a b_mean, b_m2, b_weight = b delta = b_mean - a_mean new_weight = a_weight + b_weight w2_over_w = b_weight / new_weight return ( a_mean + delta * w2_over_w, a_m2 + b_m2 + delta * delta * a_weight * w2_over_w, new_weight, ) return welford_combine_fn else: raise NotImplementedError(f"unknown reduction_type={reduction_type}") @ir_dataclass class Reduction(Loops): reduction_ranges: Sequence[_IntLike] reduction_type: ReductionType # self.dtype represents the dst dtype src_dtype: torch.dtype reduction_hint: ReductionHint def __str__(self) -> str: return self._to_str(("ranges", "reduction_ranges", "reduction_type")) __repr__ = __str__ def get_free_symbol_uses(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: return super().get_free_symbol_uses(unbacked_only) | OrderedSet().union( *(get_free_symbols(e, unbacked_only) for e in self.reduction_ranges) ) def get_reduction_size(self) -> Sequence[Expr]: return self.reduction_ranges def get_reduction_type(self) -> Optional[str]: return self.reduction_type def store_reduction( self, output_name: Optional[str], indexer: Callable[[Sequence[Expr]], Never], vars: Sequence[Expr], reduction_vars: Sequence[Symbol], ) -> None: value = ops.reduction( self.dtype, self.src_dtype, self.reduction_type, self.inner_fn(vars, reduction_vars), ) ops.store_reduction(output_name or "unnamed", indexer(vars), value) def index_length(self) -> int: return len(self.ranges) + len(self.reduction_ranges) def inner_fn_args(self) -> Sequence[Sequence[Expr]]: index = self._index(self.ranges) rindex = self._index(self.reduction_ranges, SymT.R0_INDEX) return (index, rindex) def inner_fn_free_symbols(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: index = self._index(self.ranges) rindex = self._index(self.reduction_ranges, SymT.R0_INDEX) return extract_free_symbols( self.inner_fn, index, rindex, unbacked_only=unbacked_only ) def constant_to_device(self, device: torch.device) -> IRNode: """Move this to a given device. Requires that all reads are to constants.""" loader = self.make_loader() loader = patch.object(ConstantBuffer, "override_device", device)(loader) return Reduction( device=device, dtype=self.dtype, inner_fn=loader, ranges=self.ranges, reduction_ranges=self.reduction_ranges, reduction_type=self.reduction_type, src_dtype=self.src_dtype, reduction_hint=ReductionHint.DEFAULT, ) @staticmethod def num_splits( device: torch.device, dst_dtype: torch.dtype, src_dtype: torch.dtype, inner_fn: Callable[_P, OpsValue], ranges: Sequence[_IntLike], reduction_ranges: Sequence[_IntLike], reduction_type: Union[ReductionType, Literal["scan"]], reduction_numel: Expr, input_node: Optional[IRNode] = None, ) -> tuple[ReductionHint, _IntLike]: reduction_numel_hint = V.graph.sizevars.symbolic_hint(reduction_numel) numel_hint = V.graph.sizevars.symbolic_hint(sympy_product(ranges)) should_split = reduction_type == "scan" or ( not V.graph.has_feature(device, BackendFeature.REDUCE_TO_SINGLE_ELEMENT) and reduction_type not in ( "argmax", "argmin", ) and config.split_reductions ) if not (_is_static(reduction_numel_hint) and _is_static(numel_hint)): # We don't support unbacked symints return ReductionHint.DEFAULT, 1 props = DeviceProperties.create(device) num_sm = props.multi_processor_count min_elements_per_thread = 32 if should_split: inner_reduction_splits: Callable[[int, int], int] = functools.partial( V.choices.reduction_split_factor, device, inner_reduction=True ) outer_reduction_splits: Callable[[int, int], int] = functools.partial( V.choices.reduction_split_factor, device, inner_reduction=False ) else: def inner_reduction_splits( reduction_numel_hint: int, numel_hint: int, ) -> int: return 1 outer_reduction_splits = inner_reduction_splits # easy cases if numel_hint == 1: split = inner_reduction_splits(reduction_numel_hint, numel_hint) if split == 1: # No need to split. return ReductionHint.INNER, split if input_node is not None and isinstance(input_node, TensorBox): with patch.object(FlexibleLayout, "allow_indexing", True): ( new_ranges, new_reduction_ranges, ) = extract_input_node_reduction_ranges(input_node) if new_ranges is not None and new_reduction_ranges is not None: extracted_numel_hint = V.graph.sizevars.symbolic_hint( sympy_product(new_ranges + new_reduction_ranges) ) if reduction_numel_hint == extracted_numel_hint: log.debug( "Use previous IRNode's range and reduction_ranges instead of split. " "current ranges: %s, current reduction ranges: %s, current split: %d, " "new ranges: %s, new reduction ranges: %s", ranges, reduction_ranges, split, new_ranges, new_reduction_ranges, ) # If the input_node or its dependent nodes are also Reduction nodes, # use reduction_sizes of this node or its dependent nodes directly. return ReductionHint.INNER, -1 return ReductionHint.INNER, split if ( reduction_numel_hint <= min_elements_per_thread or numel_hint >= num_sm * 2 * 32 ): return ReductionHint.DEFAULT, 1 r = Reduction( device=device, dtype=dst_dtype, inner_fn=inner_fn, ranges=ranges, reduction_ranges=reduction_ranges, reduction_type=reduction_type if reduction_type != "scan" else "sum", src_dtype=src_dtype, reduction_hint=ReductionHint.DEFAULT, ) def get_read_indices(r: Reduction) -> tuple[Sequence[Expr], bool]: device = r.get_device() assert device is not None cb = ComputedBuffer( name=None, layout=FlexibleLayout( device=device, dtype=r.get_dtype(), size=r.get_size(), ), data=r, ) read_writes = cb.get_read_writes() # try finding the full size producer # TODO this will fail for something like ((1, N) * (N, 1)).sum() # this would also possibly be wrong for producers with the different contiguity but we hope those cases are rare assert read_writes.range_vars is not None range_vars = [ r for r in read_writes.range_vars if isinstance(r, Expr) and not isinstance(r, sympy.Number) ] indices = [] changed = False for md in sorted(read_writes.reads, key=lambda x: x.name): if all(r in md.index.free_symbols for r in range_vars): indices.append(md.index) if md.name in V.graph.name_to_buffer: buf = V.graph.name_to_buffer[md.name] original_stride = getattr(buf.layout, "stride", None) buf.decide_layout() if getattr(buf.layout, "stride", None) != original_stride: changed = True return indices, changed indices, changed = get_read_indices(r) if changed: indices, _ = get_read_indices(r) if len(indices) == 0: # TODO determine splits when all inputs are broadcast return ReductionHint.DEFAULT, 1 (_, reduction_vars), ranges1 = dependencies.index_vars_squeeze( r.get_size(), r.get_reduction_size() ) num_outer = 0 num_inner = 0 for i in indices: j = V.graph.sizevars.simplify_with_ranges(i, ranges1) strides = V.graph.sizevars.stride_hints( j, reduction_vars, list(ranges1.keys()) ) outer = all(s > 1 for s in strides) if outer: num_outer += 1 else: num_inner += 1 if num_inner > num_outer: return ReductionHint.INNER, inner_reduction_splits( reduction_numel_hint, numel_hint ) else: return ReductionHint.OUTER, outer_reduction_splits( reduction_numel_hint, numel_hint ) @staticmethod def _unroll_reduction_fn( inner_fn: Callable[[Sequence[_IntLike], Sequence[_IntLike]], OpsValue], reduction_ranges: Sequence[_IntLike], reduction_type: str, src_dtype: torch.dtype, ) -> Callable[[Sequence[_IntLike]], OpsValue]: """Convert inner_fn from a reduction to an pointwise""" reduction_ranges = V.graph.sizevars.guard_int_seq(reduction_ranges) combine_fn = get_reduction_combine_fn(reduction_type, src_dtype) def fn(index: Sequence[_IntLike]) -> Any: return functools.reduce( combine_fn, ( value_fn(index, rindex) for rindex in itertools.product( *[range(x) for x in reduction_ranges] ) ), ) value_fn: Callable[[Sequence[_IntLike], Sequence[_IntLike]], Any] if reduction_type in ("argmin", "argmax"): flatten_index = _fixed_indexer( reduction_ranges, FlexibleLayout.contiguous_strides(reduction_ranges), ) def value_fn( index: Sequence[_IntLike], rindex: Sequence[_IntLike] ) -> tuple[OpsValue, OpsValue]: rindex = [sympy.expand(i) for i in rindex] return ( inner_fn(index, rindex), ops.index_expr(flatten_index(rindex), torch.int64), ) return lambda index: fn(index)[1] else: value_fn = inner_fn return fn @classmethod def create( cls, device: torch.device, dst_dtype: torch.dtype, src_dtype: torch.dtype, inner_fn: Callable[..., Any], ranges: Sequence[Expr], reduction_ranges: Sequence[Expr], reduction_type: ReductionType, reduction_hint: ReductionHint = ReductionHint.DEFAULT, input_node: Optional[IRNode] = None, ) -> Union[TensorBox, ShapeAsConstantBuffer]: reduction_numel = V.graph.sizevars.simplify(sympy_product(reduction_ranges)) if reduction_numel == 0: # N.B. This is a hack to generate the literal of the given type # Ideally, we should be fixing `def constant` in triton.py # but it breaks due to hardcoded dtypes in other places def py_cnst(val: object) -> Union[bool, float, int]: if dst_dtype == torch.bool: return bool(val) elif dst_dtype.is_floating_point: assert isinstance(val, SupportsFloat), type(val) return float(val) else: assert isinstance(val, SupportsInt), type(val) return int(val) rtypes_to_inits = { "sum": py_cnst(0), "xor_sum": py_cnst(0), "prod": py_cnst(1), "any": py_cnst(0), # "all" is desugared to `!any(!val)` } assert reduction_type in rtypes_to_inits.keys(), ( f"{reduction_type} not supported for zero-dimension tensors!" ) def const_fn(index: int) -> OpsValue: return ops.constant(rtypes_to_inits[reduction_type], dst_dtype) return Pointwise.create( device=device, dtype=src_dtype, inner_fn=const_fn, ranges=list(ranges), ) if reduction_numel == 1: # this reduction is actually a pointwise op if reduction_type in ("argmin", "argmax"): def fn(index: int) -> OpsValue: return ops.constant(0, dst_dtype) else: def fn(index: int) -> OpsValue: reduction_index = [sympy.S.Zero for _ in reduction_ranges] return inner_fn(index, reduction_index) return Pointwise.create( device=device, dtype=dst_dtype, inner_fn=fn, ranges=ranges ) if ( isinstance(reduction_numel, Integer) and V.graph.sizevars.size_hint_or_throw(reduction_numel) < config.unroll_reductions_threshold and (sympy_product(ranges) != 1 or is_gpu(device.type)) ): # NB: This works around https://github.com/pytorch/pytorch/issues/140457 # since turning reductions into pointwise ops can exacerbate this problem return Pointwise.create( device=device, dtype=dst_dtype, inner_fn=cls._unroll_reduction_fn( inner_fn, reduction_ranges, reduction_type, src_dtype ), ranges=ranges, ) # triton doesn't support reduce to single element well, so break it up hint, split = cls.num_splits( device, dst_dtype, src_dtype, inner_fn, ranges, reduction_ranges, reduction_type, reduction_numel, input_node, ) def _maybe_increase_split(split: int) -> int: # don't apply min_num_split constraint for static shape case. if _is_static(reduction_numel): return split if split > 1: return max(split, config.min_num_split) else: return split split = _maybe_increase_split(split) # intermediate reduction in split can contain complex indexing, # and num_splits will fail to correctly set the hint # reuse the passed hint if available if reduction_hint == ReductionHint.DEFAULT: reduction_hint = hint if split == -1: assert input_node is not None with patch.object(FlexibleLayout, "allow_indexing", True): new_ranges, new_reduction_ranges = extract_input_node_reduction_ranges( input_node ) assert new_ranges is not None assert new_reduction_ranges is not None return cls.create_multilayer_existing_ranges( device, dst_dtype, src_dtype, inner_fn, ranges, reduction_ranges, new_ranges, new_reduction_ranges, reduction_type, reduction_hint, ) elif split > 1: # triton doesn't support reduce to single element well, so break it up return cls.create_multilayer( device, dst_dtype, src_dtype, inner_fn, ranges, reduction_ranges, reduction_type, split, reduction_hint, input_node, ) return TensorBox.create( Reduction( device=device, dtype=dst_dtype, inner_fn=inner_fn, ranges=ranges, reduction_ranges=reduction_ranges, reduction_type=reduction_type, src_dtype=src_dtype, reduction_hint=reduction_hint, ) ) @staticmethod def default_accumulator( reduction_type: str, dtype: torch.dtype ) -> Union[_NumLike, Sequence[_NumLike]]: if reduction_type in ("max", "argmax"): if is_float_dtype(dtype): return float("-inf") elif is_boolean_dtype(dtype): return False else: return torch.iinfo(dtype).min if reduction_type in ("min", "argmin"): if is_float_dtype(dtype): return float("inf") elif is_boolean_dtype(dtype): return True else: return torch.iinfo(dtype).max zero = False if is_boolean_dtype(dtype) else 0 one = True if is_boolean_dtype(dtype) else 1 return { "sum": zero, "prod": one, "xor_sum": zero, "any": zero, "welford_reduce": (zero, zero, zero), "welford_combine": (zero, zero, zero), "online_softmax_reduce": (float("-inf"), zero), }[reduction_type] @staticmethod def default_value( reduction_type: str, dtype: torch.dtype ) -> Union[_NumLike, Sequence[_NumLike]]: if reduction_type == "welford_reduce": return 0 return Reduction.default_accumulator(reduction_type, dtype) @staticmethod def _multilayer_second_step_hint( split: _IntLike, numel_hint: int, reduction_hint: ReductionHint ) -> ReductionHint: if split == -1: return reduction_hint if split <= 512 and numel_hint <= 512 and reduction_hint == ReductionHint.OUTER: return ReductionHint.OUTER_TINY if ( split <= 1024 and numel_hint <= 256 and reduction_hint == ReductionHint.OUTER ): return ReductionHint.OUTER_TINY return reduction_hint @classmethod def check_for_split_dense_dim_reindexing( cls, reduction_numel: _IntLike, input_node: Optional[IRNode] ) -> Optional[int]: """ If we are reducing over the full tensor, and it is non-dense in the last dimension, reindex so we reduce over the dense dimension. initially just handle complete reduction case """ if input_node is None: return None if not V.graph.sizevars.statically_known_equals( input_node.get_numel(), reduction_numel ): return None input_node.realize() try: # finalize layout as_storage_and_layout(input_node) except NotImplementedError: return None strides = input_node.get_stride() for i, s in enumerate(strides[:-1]): if V.graph.sizevars.statically_known_equals(s, 1): return i return None @classmethod def _multilayer_wrap_loader( cls, loader: Callable[..., OpsValue], reduction_ranges: Sequence[_IntLike], reduction_numel: _IntLike, split: _IntLike, block_size: _IntLike, default: Union[_NumLike, Sequence[_NumLike]], input_node: Optional[IRNode] = None, ) -> Callable[..., object]: dense_index = cls.check_for_split_dense_dim_reindexing( reduction_numel, input_node ) reindex = View.dynamic_reshape_indexer( reduction_ranges, [reduction_numel], dense_index ) need_mask = not V.graph.sizevars.statically_known_true( sympy.Eq(reduction_numel % split, 0) ) def wrapper_fn( index: Sequence[Symbol], reduction_index: Sequence[Symbol] ) -> OpsValue: (reduction_index,) = reduction_index *new_index, reduction_block = index indices = block_size * reduction_block + reduction_index def body() -> OpsValue: return loader(new_index, reindex([indices])) if need_mask: index_dtype = dtype_from_size(reduction_numel) mask = ops.lt( ops.index_expr(indices, index_dtype), ops.index_expr(reduction_numel, index_dtype), ) return ops.masked(mask, body, default) else: return body() return wrapper_fn @classmethod def _multilayer_wrap_loader_existing_ranges( cls, loader: Callable[[Sequence[Expr], Sequence[Expr]], OpsValue], original_ranges: Sequence[Expr], original_reduction_ranges: Sequence[Expr], new_ranges: Sequence[Integer], new_reduction_ranges: Sequence[Integer], ) -> Callable[[Sequence[sympy.Expr], Sequence[sympy.Expr]], OpsValue]: assert all(r == 1 for r in original_ranges), ( f"Only enabled for numel_hint == 1, found {original_ranges=}" ) reindex = View.dynamic_reshape_indexer( original_reduction_ranges, tuple(new_ranges) + tuple(new_reduction_ranges) ) def wrapper_fn( merged_index: Sequence[Expr], new_reduction_index: Sequence[Expr], ) -> OpsValue: original_idx = merged_index[: len(original_ranges)] new_index = merged_index[len(original_ranges) :] return loader( original_idx, reindex(tuple(new_index) + tuple(new_reduction_index)), ) return wrapper_fn @classmethod def create_multilayer_helper( cls, device: torch.device, dst_dtype: torch.dtype, src_dtype: torch.dtype, wrapper_fn: Callable[..., Any], original_ranges: Sequence[Expr], original_reduction_ranges: Sequence[Expr], new_ranges: list[Expr], new_reduction_ranges: list[Integer], reduction_type: ReductionType, split: _IntLike, reduction_hint: ReductionHint, ) -> Union[TensorBox, ShapeAsConstantBuffer]: """ Break a large reduction up into multiple smaller reductions recursively """ # triton will automatically compute reductions in fp32 if reducing over fp16/bf16 # within the kernel. keep the intermediate in fp32 so as to keep the whole reduction # in fp32 and not reduce precision by breaking up the kernel into multiple layers intermediate_dtype = ( dst_dtype if dst_dtype not in (torch.float16, torch.bfloat16) else torch.float ) intermediate = Reduction.create( device, intermediate_dtype, src_dtype, wrapper_fn, new_ranges, new_reduction_ranges, reduction_type, reduction_hint, ) intermediate.realize() intermediate_loader = intermediate.make_loader() def intermediate_fn( index: Sequence[_IntLike], reduction_index: Sequence[_IntLike] ) -> OpsValue: return intermediate_loader([*index, *reduction_index]) numel_hint = V.graph.sizevars.size_hint(sympy_product(original_ranges)) reduction_hint = cls._multilayer_second_step_hint( split, numel_hint, reduction_hint ) assert original_ranges == new_ranges[: len(original_ranges)] return TensorBox.create( Reduction( device=device, dtype=dst_dtype, inner_fn=intermediate_fn, ranges=original_ranges, reduction_ranges=new_ranges[len(original_ranges) :], reduction_type=reduction_type, src_dtype=src_dtype, reduction_hint=reduction_hint, ) ) @classmethod def create_multilayer( cls, device: torch.device, dst_dtype: torch.dtype, src_dtype: torch.dtype, inner_fn: Callable[..., Any], ranges: Sequence[Expr], reduction_ranges: Sequence[Expr], reduction_type: ReductionType, split: _IntLike, reduction_hint: ReductionHint, input_node: Optional[IRNode] = None, ) -> Union[TensorBox, ShapeAsConstantBuffer]: """ Break a large reduction up into multiple smaller reductions recursively """ # TODO(jansel): realize the reduction so we can do dynamic indexing reduction_numel = sympy_product(reduction_ranges) block_size = FloorDiv(reduction_numel + (split - 1), split) default = cls.default_value(reduction_type, dst_dtype) wrapper_fn = cls._multilayer_wrap_loader( inner_fn, reduction_ranges, reduction_numel, split, block_size, default, input_node, ) return cls.create_multilayer_helper( device, dst_dtype, src_dtype, wrapper_fn, ranges, reduction_ranges, [*ranges, split], [block_size], reduction_type, split, reduction_hint, ) @classmethod def create_multilayer_existing_ranges( cls, device: torch.device, dst_dtype: torch.dtype, src_dtype: torch.dtype, inner_fn: Callable[..., Any], original_ranges: Sequence[Expr], original_reduction_ranges: Sequence[Expr], new_ranges: list[Integer], new_reduction_ranges: list[Integer], reduction_type: ReductionType, reduction_hint: ReductionHint, ) -> Union[TensorBox, ShapeAsConstantBuffer]: """ Break a large reduction up into multiple smaller reductions recursively """ wrapper_fn = cls._multilayer_wrap_loader_existing_ranges( inner_fn, original_ranges, original_reduction_ranges, new_ranges, new_reduction_ranges, ) return cls.create_multilayer_helper( device, dst_dtype, src_dtype, wrapper_fn, original_ranges, original_reduction_ranges, [*original_ranges, *new_ranges], new_reduction_ranges, reduction_type, -1, reduction_hint, ) def _fixed_indexer( size: Sequence[int], stride: Optional[Sequence[int]] = None, offset: Expr = Integer(0), ) -> Callable[[Sequence[Expr]], Expr]: """A closure containing math to read a given element""" def indexer(index: Sequence[int]) -> int: assert stride is not None and len(index) == len(stride) assert len(index) == len(size) result = offset for idx, st, sz in zip(index, stride, size): if sz != 1: result = result + idx * st return result return indexer INNER_FN_TY: TypeAlias = Callable[[Sequence[Expr], Sequence[Expr]], OpsValue] class MultiOutputReduction(Reduction): output_index: int def __init__( self, device: torch.device, dst_dtype: torch.dtype, inner_fns: Union[INNER_FN_TY, Sequence[INNER_FN_TY]], ranges: Sequence[Integer], reduction_ranges: Sequence[Integer], reduction_type: ReductionType, src_dtype: torch.dtype, reduction_hint: ReductionHint, output_index: int, ): if callable(inner_fns): inner_fns = (inner_fns,) loader: Callable[[Sequence[Expr], Sequence[Expr]], Any] if len(inner_fns) == 1: loader = inner_fns[0] else: def loader( idx: Sequence[Expr], reduction_idx: Sequence[Expr] ) -> tuple[OpsValue, ...]: return tuple(fn(idx, reduction_idx) for fn in inner_fns) super().__init__( device=device, dtype=dst_dtype, inner_fn=loader, ranges=ranges, reduction_ranges=reduction_ranges, reduction_type=reduction_type, src_dtype=src_dtype, reduction_hint=reduction_hint, ) self.output_index = output_index def store_reduction( self, output_name: Optional[str], indexer: Callable[[Sequence[Expr]], Never], vars: Sequence[Expr], reduction_vars: Sequence[Symbol], ) -> Any: values = ops.reduction( self.dtype, self.src_dtype, self.reduction_type, self.inner_fn(vars, reduction_vars), ) assert isinstance(values, (tuple, list)), type(values) value = values[self.output_index] return ops.store_reduction(output_name or "unnamed", indexer(vars), value) class OnlineSoftmaxReduction(MultiOutputReduction): @classmethod def create( # type: ignore[override] cls, device: torch.device, dst_dtype: torch.dtype, src_dtype: torch.dtype, inner_fn: Callable[..., Any], ranges: Sequence[Expr], reduction_ranges: Sequence[Expr], num_output: int, reduction_hint: ReductionHint = ReductionHint.DEFAULT, input_node: Optional[IRNode] = None, ) -> Sequence[Union[TensorBox, ShapeAsConstantBuffer]]: """ Create the reduction disregarding splitting. """ results = tuple( TensorBox.create( MultiOutputReduction( device, dst_dtype, inner_fn, ranges, reduction_ranges, "online_softmax_reduce", src_dtype, reduction_hint, output_idx, ) ) for output_idx in range(num_output) ) for t in results: t.realize() return results class WelfordReduction(MultiOutputReduction): @classmethod def create( # type: ignore[override] cls, device: torch.device, dtype: torch.dtype, inner_fns: Sequence[Callable[..., Any]], ranges: list[Integer], reduction_ranges: list[Integer], reduction_type: ReductionType, reduction_hint: ReductionHint = ReductionHint.DEFAULT, ) -> Sequence[Union[TensorBox, ShapeAsConstantBuffer]]: assert reduction_type in ("welford_reduce", "welford_combine") reduction_numel = V.graph.sizevars.simplify(sympy_product(reduction_ranges)) def const(val: int) -> Union[TensorBox, ShapeAsConstantBuffer]: def inner_fn(idx: Sequence[Expr]) -> OpsValue: return ops.constant( val, dtype, ) return Pointwise.create( device=device, dtype=dtype, inner_fn=inner_fn, ranges=list(ranges), ) if reduction_numel == 0: mean = const(0) m2 = const(0) weight = const(0) return mean, m2, weight if reduction_numel == 1: def copy( loader: Callable[[Sequence[Expr], Sequence[Expr]], OpsValue], ) -> Union[TensorBox, ShapeAsConstantBuffer]: def inner_fn(idx: Sequence[Expr]) -> OpsValue: reduction_index = [sympy.S.Zero for _ in reduction_ranges] return loader(idx, reduction_index) return Pointwise.create( device=device, dtype=dtype, inner_fn=inner_fn, ranges=list(ranges), ) if reduction_type == "welford_reduce": return copy(inner_fns[0]), const(0), const(1) else: return tuple(copy(fn) for fn in inner_fns) # TODO: Unrolled reduction # if ( # isinstance(reduction_numel, Integer) # and V.graph.sizevars.size_hint(reduction_numel) # < config.unroll_reductions_threshold # and sympy_product(ranges) != 1 # ): # return Pointwise.create( # device, # dst_dtype, # cls._unroll_reduction_fn( # inner_fn, reduction_ranges, reduction_type, src_dtype, # ), # ranges, # ) # triton doesn't support reduce to single element well, so break it up hint, split = Reduction.num_splits( device, dtype, dtype, inner_fns[0], ranges, reduction_ranges, reduction_type=reduction_type, reduction_numel=reduction_numel, ) # intermediate reduction in split can contain complex indexing, # and num_splits will fail to correctly set the hint # reuse the passed hint if available if reduction_hint == ReductionHint.DEFAULT: reduction_hint = hint if split > 1: # triton doesn't support reduce to single element well, so break it up return cls.create_multilayer( device, dtype, inner_fns, ranges, reduction_ranges, reduction_type, split, reduction_hint, ) results = [ TensorBox.create( WelfordReduction( device, dtype, inner_fns, ranges, reduction_ranges, reduction_type, dtype, reduction_hint, output_idx, ) ) for output_idx in range(3) ] for t in results: t.realize() return results @staticmethod def default_value( reduction_type: str, dtype: torch.dtype ) -> Union[_NumLike, Sequence[_NumLike]]: return (0, 0, 0) @classmethod def create_multilayer( # type: ignore[override] cls, device: torch.device, dtype: torch.dtype, inner_fns: Sequence[Callable[..., Any]], ranges: list[Integer], reduction_ranges: list[Integer], reduction_type: ReductionType, split: _IntLike, reduction_hint: ReductionHint, ) -> Sequence[Union[TensorBox, ShapeAsConstantBuffer]]: """ Break a large reduction up into multiple smaller reductions recursively """ reduction_numel = sympy_product(reduction_ranges) need_mask = not V.graph.sizevars.statically_known_true( sympy.Eq(reduction_numel % split, 0) ) if need_mask and reduction_type != "welford_combine": # If we need mask, then "welford_reduce" doesn't work because # masked inputs shouldn't count towards the welford weight def constant( idx: Sequence[Expr], reduction_idx: Sequence[Expr], value: int ) -> OpsValue: return ops.constant(value, dtype) return cls.create_multilayer( device=device, dtype=dtype, inner_fns=( inner_fns[0], partial(constant, value=0), partial(constant, value=1), ), ranges=ranges, reduction_ranges=reduction_ranges, reduction_type="welford_combine", split=split, reduction_hint=reduction_hint, ) block_size = FloorDiv(reduction_numel + (split - 1), split) intermediates = WelfordReduction.create( device, dtype, tuple( cls._multilayer_wrap_loader( loader, reduction_ranges, reduction_numel, split, block_size, default=0, ) for loader in inner_fns ), [*ranges, split], [block_size], reduction_type, reduction_hint, ) for i in intermediates: i.realize() def intermediate_loader_fn( index: Sequence[Expr], reduction_index: Sequence[Expr], loader: Callable[[Sequence[Expr]], OpsValue], ) -> OpsValue: return loader([*index, *reduction_index]) numel_hint = V.graph.sizevars.size_hint(sympy_product(ranges)) reduction_hint = cls._multilayer_second_step_hint( split, numel_hint, reduction_hint ) return WelfordReduction.create( device, dtype, tuple( partial(intermediate_loader_fn, loader=i.make_loader()) for i in intermediates ), ranges, [split], # welford_reduce turns one input into three outputs, which are combined with welford_combine "welford_combine", reduction_hint, ) @ir_dataclass class Scan(Loops): scan_ranges: list[Integer] size: list[Integer] combine_fn: Callable[[tuple[Any, ...], tuple[Any, ...]], tuple[Any, ...]] reindex: Callable[[Sequence[_IntLike], Sequence[_IntLike]], Sequence[_IntLike]] reduction_hint: ReductionHint output_index: int # output_index indexes the following tuples dtypes: tuple[torch.dtype, ...] inner_fns: tuple[Callable[..., Any], ...] # HACK we mimic reduction def get_free_symbol_uses(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: # TODO: Can combine_fn/reindex close over unbacked symbols? If so, we # need to explicitly represent the closure so we can pull out unbacked # symbols here return ( super().get_free_symbol_uses(unbacked_only) | OrderedSet().union( *(get_free_symbols(e, unbacked_only) for e in self.scan_ranges) ) | OrderedSet().union( *(get_free_symbols(e, unbacked_only) for e in self.size) ) ) def __post_init__(self) -> None: assert len(self.ranges) + len(self.scan_ranges) == len(self.size) super().__post_init__() def store_reduction( self, output_name: Optional[str], indexer: Callable[[Sequence[_IntLike]], Never], vars: Sequence[Expr], scan_vars: Sequence[Symbol], ) -> Any: idx = self.reindex(vars, scan_vars) values = tuple(inner_fn(idx) for inner_fn in self.inner_fns) result = ops.scan(self.dtypes, self.combine_fn, values) return ops.store( output_name or "unnamed", indexer(idx), result[self.output_index] ) def get_reduction_type(self) -> Optional[str]: # return self.scan_op return "custom" def get_reduction_size(self) -> Sequence[Expr]: return self.scan_ranges def get_size(self) -> Sequence[Expr]: return self.size def get_pointwise_size(self) -> Sequence[Expr]: return self.ranges def index_length(self) -> int: return len(self.ranges) + len(self.scan_ranges) def inner_fn_args(self) -> Sequence[Sequence[_IntLike]]: index = self._index(self.ranges) rindex = self._index(self.scan_ranges, SymT.R0_INDEX) idx = self.reindex(index, rindex) return (idx,) def inner_fn_free_symbols(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: index = self._index(self.ranges) rindex = self._index(self.scan_ranges, SymT.R0_INDEX) idx = self.reindex(index, rindex) return extract_free_symbols(self.inner_fn, idx, unbacked_only=unbacked_only) @classmethod def create( # type: ignore[override] cls, device: torch.device, dtypes: tuple[torch.dtype, ...], inner_fns: tuple[Callable[[Sequence[Expr]], Any], ...], size: list[Integer], axis: int, combine_fn: Callable[[tuple[Any, ...], tuple[Any, ...]], tuple[Any, ...]], reduction_hint: ReductionHint = ReductionHint.DEFAULT, *, # Whether we have the option to fallback to aten can_fallback_to_aten: bool = True, **kwargs: Any, ) -> Sequence[Optional[Union[TensorBox, ShapeAsConstantBuffer]]]: pointwise_ranges = [*size[:axis], *size[axis + 1 :]] scan_ranges = [size[axis]] if not V.graph.has_feature(device, BackendFeature.SCAN): return [None] * len(dtypes) if len(dtypes) > 1 and not V.graph.has_feature( device, BackendFeature.TUPLE_REDUCTION ): return [None] * len(dtypes) sizevars = V.graph.sizevars scan_numel = sizevars.simplify(sympy_product(scan_ranges)) assert len(dtypes) == len(inner_fns) # Scan with a single element is just a copy if sizevars.statically_known_true(sympy.Le(scan_numel, 1)): return [ Pointwise.create( device=device, dtype=dtypes[output_index], inner_fn=inner_fns[output_index], ranges=size, ) for output_index in range(len(dtypes)) ] reduction_hint, num_splits = cls.num_splits( device=device, dtype=dtypes[0], inner_fn=inner_fns[0], axis=axis, pointwise_ranges=pointwise_ranges, scan_ranges=scan_ranges, combine_fn=combine_fn, scan_numel=scan_numel, ) scan_type = Scan if num_splits > 1: supports_split = ( torch.version.hip is None or (has_triton and triton_version >= "3.3.0") ) and (len(dtypes) == 1) if not supports_split: if can_fallback_to_aten: # Fallback to ATen return [None] * len(dtypes) else: num_splits = 1 else: scan_type = SplitScan def reindex(index: Sequence[Expr], scan_index: Sequence[Expr]) -> list[Expr]: assert len(scan_index) == len(scan_ranges) assert len(index) == len(pointwise_ranges) return [*index[:axis], *scan_index, *index[axis:]] results = [ TensorBox.create( scan_type( device=device, dtype=dtypes[output_index], dtypes=dtypes, inner_fn=inner_fns[output_index], inner_fns=inner_fns, size=size, ranges=pointwise_ranges, scan_ranges=scan_ranges, combine_fn=combine_fn, reindex=reindex, reduction_hint=reduction_hint, output_index=output_index, **kwargs, ) ) for output_index in range(len(dtypes)) ] for result in results: result.realize() return results @classmethod def num_splits( cls, device: torch.device, dtype: torch.dtype, inner_fn: Callable[[Sequence[Expr]], OpsValue], axis: int, pointwise_ranges: list[Integer], scan_ranges: list[Integer], combine_fn: Callable[[tuple[Any, ...], tuple[Any, ...]], tuple[Any, ...]], scan_numel: Expr, ) -> tuple[ReductionHint, _IntLike]: # TODO: custom splitting heuristic for scan def wrapper_fn(idx: Sequence[Expr], reduction_idx: Sequence[Expr]) -> OpsValue: return inner_fn([*idx[:axis], *reduction_idx, *idx[axis:]]) return Reduction.num_splits( device=device, dst_dtype=dtype, src_dtype=dtype, inner_fn=wrapper_fn, ranges=pointwise_ranges, reduction_ranges=scan_ranges, reduction_type="scan", reduction_numel=scan_numel, ) # This signifies a scan op that should go through TritonSplitScanKernel codegen on CUDA. @ir_dataclass class SplitScan(Scan): pass @ir_dataclass class Sort(Loops): # Sorts a tuple of key, value pairs sort_ranges: list[Integer] size: list[Integer] reindex: Callable[[Sequence[Expr], Sequence[Expr]], Sequence[Expr]] reduction_hint: ReductionHint output_index: int # output_index indexes the following tuples dtypes: tuple[torch.dtype, ...] inner_fns: tuple[Callable[..., Any], ...] stable: bool descending: bool # HACK we mimic reduction def get_free_symbol_uses(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: return ( super().get_free_symbol_uses(unbacked_only) | OrderedSet().union( *(get_free_symbols(e, unbacked_only) for e in self.sort_ranges) ) | OrderedSet().union( *(get_free_symbols(e, unbacked_only) for e in self.size) ) ) def __post_init__(self) -> None: assert len(self.ranges) + len(self.sort_ranges) == len(self.size) super().__post_init__() def store_reduction( self, output_name: Optional[str], indexer: Callable[[Sequence[Expr]], Expr], vars: Sequence[Expr], reduction_vars: Sequence[Expr], ) -> Any: idx = self.reindex(vars, reduction_vars) values = tuple(inner_fn(idx) for inner_fn in self.inner_fns) result = ops.sort(self.dtypes, values, self.stable, self.descending) return ops.store( output_name or "unnamed", indexer(idx), result[self.output_index] ) def get_reduction_type(self) -> Optional[str]: return "sort" def get_reduction_size(self) -> Sequence[Expr]: return self.sort_ranges def get_size(self) -> Sequence[Expr]: return self.size def get_pointwise_size(self) -> Sequence[Expr]: return self.ranges def index_length(self) -> int: return len(self.ranges) + len(self.sort_ranges) def inner_fn_args(self) -> Sequence[Sequence[Expr]]: index = self._index(self.ranges) rindex = self._index(self.sort_ranges, SymT.R0_INDEX) idx = self.reindex(index, rindex) return (idx,) def inner_fn_free_symbols(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: index = self._index(self.ranges) rindex = self._index(self.sort_ranges, SymT.R0_INDEX) idx = self.reindex(index, rindex) return extract_free_symbols(self.inner_fn, idx, unbacked_only=unbacked_only) @classmethod def create( # type: ignore[override] cls, device: torch.device, dtypes: tuple[torch.dtype, ...], inner_fns: tuple[Callable[[list[Expr]], Any], ...], size: list[Integer], axis: int, stable: bool, descending: bool, reduction_hint: ReductionHint = ReductionHint.DEFAULT, **kwargs: Any, ) -> Sequence[Optional[Union[TensorBox, ShapeAsConstantBuffer]]]: pointwise_ranges = [*size[:axis], *size[axis + 1 :]] sort_ranges = [size[axis]] if not V.graph.has_feature(device, BackendFeature.SORT): return [None] * len(dtypes) sizevars = V.graph.sizevars sort_numel = sizevars.simplify(sympy_product(sort_ranges)) # Heuristic, smallest rblock where triton usually outperforms aten.sort # It also isn't bandwidth bound so fusion is unlikely to help. max_rblock = 512 is_persistent_kernel = ( config.triton.persistent_reductions and sizevars.statically_known_true(sympy.Le(sort_numel, max_rblock)) ) if not is_persistent_kernel: # We only support persistent triton kernels return [None] * len(dtypes) assert len(dtypes) == len(inner_fns) # Sort with a single element is just a copy if sizevars.statically_known_true(sympy.Le(sort_numel, 1)): return [ Pointwise.create( device=device, dtype=dtypes[output_index], inner_fn=inner_fns[output_index], ranges=size, ) for output_index in range(len(dtypes)) ] def reindex(index: Sequence[Expr], sort_index: Sequence[Expr]) -> list[Expr]: assert len(sort_index) == len(sort_ranges) assert len(index) == len(pointwise_ranges) return [*index[:axis], *sort_index, *index[axis:]] results = [ TensorBox.create( Sort( device=device, dtype=dtypes[output_index], dtypes=dtypes, inner_fn=inner_fns[output_index], inner_fns=inner_fns, size=size, ranges=pointwise_ranges, sort_ranges=sort_ranges, reindex=reindex, reduction_hint=reduction_hint, output_index=output_index, stable=stable, descending=descending, **kwargs, ) ) for output_index in range(len(dtypes)) ] for result in results: result.realize() return results def is_storage_and_layout(x: IRNode) -> bool: try: as_storage_and_layout(x, freeze=False) return True except NotImplementedError: return False def is_contiguous_storage_and_layout(x: IRNode) -> bool: try: _buffer, layout = as_storage_and_layout(x, freeze=False) # pad the stride here so we will NOT claim an tensor as contiguous # if a padding is gonna happen. if layout.should_pad_strides(): layout.pad_strides() return layout.is_contiguous() except NotImplementedError: return False def as_storage_and_layout( x: IRNode, freeze: bool = True, want_contiguous: bool = False, stride_order: Optional[Sequence[Union[int, Integer]]] = None, allow_padding: bool = False, exact_strides: Optional[Sequence[Union[int, Integer]]] = None, ) -> tuple[StorageBox, Layout]: """ Try to simplify x into a StorageBox and a Layout. allow_padding only affect how we apply stride_order. When allow_padding is True, we have the freedom to add padding when applying the stride_order. """ if isinstance(x, TensorBox): return as_storage_and_layout( x.data, freeze=freeze, want_contiguous=want_contiguous, stride_order=stride_order, allow_padding=allow_padding, exact_strides=exact_strides, ) if isinstance(x, StorageBox): _, layout = as_storage_and_layout( x.data, freeze=freeze, want_contiguous=want_contiguous, stride_order=stride_order, allow_padding=allow_padding, exact_strides=exact_strides, ) return x, x.data.get_layout() if isinstance(x, Buffer): if freeze: if want_contiguous: x.freeze_layout() assert x.get_layout().is_contiguous() elif stride_order is not None: x.freeze_layout_with_stride_order( stride_order, allow_padding=allow_padding ) elif exact_strides is not None: x.freeze_layout_with_exact_strides( exact_strides, allow_padding=allow_padding ) else: x.decide_layout() return StorageBox(x), x.get_layout() if isinstance(x, ReinterpretView): # making the base of x contiguous or stride_ordered will not necessarily make # the ReinterpretView either, so don't pass along those arguments buffer, _ = as_storage_and_layout( x.data, freeze=freeze, ) return buffer, x.layout raise NotImplementedError def is_stride_order_storage_and_layout( x: IRNode, stride_order: Sequence[Union[int, Integer]] ) -> bool: try: _buffer, layout = as_storage_and_layout(x, freeze=False) return layout.is_stride_ordered(stride_order) except NotImplementedError: return False def is_unaligned(node: IRNode) -> bool: if isinstance(node, (TensorBox, StorageBox)): return is_unaligned(node.data) if isinstance(node, ReinterpretView): layout = node.layout has_unaligned_layout = not V.graph.sizevars.statically_known_multiple_of( layout.offset * get_dtype_size(layout.dtype), GPU_ALIGN_BYTES ) return is_unaligned(node.data) or has_unaligned_layout if isinstance(node, Buffer): return node.get_name() in V.graph.unaligned_buffers # assume to be aligned otherwise return False @ir_dataclass class BaseView(IRNode): data: IRNode def get_free_symbol_uses(self, unbacked_only: bool = False) -> OrderedSet[Symbol]: return self.data.get_free_symbol_uses(unbacked_only) def make_reindexer(self) -> Callable[[Sequence[Expr]], Sequence[Expr]]: raise NotImplementedError(f"make_reindexer NYI on {self}") def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]: inner = self.data.make_indexer() reindex = self.make_reindexer() def indexer(idx: Sequence[Expr]) -> Expr: return inner(reindex(idx)) return indexer def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]: inner = self.data.make_loader() reindex = self.make_reindexer() def loader(idx: Sequence[Expr]) -> OpsValue: return inner(reindex(idx)) return loader @property def dtype(self) -> torch.dtype: return self.data.get_dtype() def get_layout(self) -> Layout: return self.data.get_layout() def get_device(self) -> Optional[torch.device]: return self.data.get_device() def get_origin_node(self) -> Optional[torch.fx.Node]: return None def get_name(self) -> str: return self.data.get_name() def get_pointwise_size(self) -> Sequence[Expr]: return self.get_size() def mark_reuse(self, users: int) -> None: return self.data.mark_reuse(users) def has_exceeded_max_reads(self) -> bool: return self.data.has_exceeded_max_reads() def realize(self) -> Optional[str]: return self.data.realize() def realize_hint(self) -> None: self.data.realize_hint() def get_storage_numel(self) -> _IntLike: return self.data.get_storage_numel() def is_extern(self) -> bool: return self.data.is_extern() def is_module_buffer(self) -> bool: assert isinstance(self.data, BaseView), type(self.data) return self.data.is_module_buffer() def get_read_names(self) -> OrderedSet[str]: return self.data.get_read_names() def get_reads(self) -> OrderedSet[Dep]: with patch.object(FlexibleLayout, "allow_indexing", True): return extract_read_writes( self.make_loader(), self.get_size(), ).reads def unwrap_view(self) -> IRNode: x: IRNode = self while isinstance(x, BaseView): x = x.data return x def constant_to_device(self, device: torch.device) -> IRNode: """Move this to a given device. Requires that all reads are to constants.""" loader = self.make_loader() loader = patch.object(ConstantBuffer, "override_device", device)(loader) return Pointwise( device=device, dtype=self.get_dtype(), inner_fn=loader, ranges=self.get_size(), ) @ir_dataclass class ExpandView(BaseView): size: Sequence[Expr] @staticmethod def _normalize_size(x: IRNode, new_size: Sequence[_IntLike]) -> Sequence[_IntLike]: """Replace `-1` with correct sizes""" sizevars = V.graph.sizevars new_size = [sympy.expand(s) for s in new_size] old_size = x.get_size() old_size = [None] * (len(new_size) - len(old_size)) + list(old_size) assert len(new_size) == len(old_size) for i in range(len(new_size)): if new_size[i] == -1: assert old_size[i] is not None new_size[i] = old_size[i] elif old_size[i] is None or V.graph.sizevars.is_size_one_or_false( old_size[i] ): pass else: # Sanity check: Expect broadcast compatibility # # NB: new_size[i] == old_size[i] is expected to already be # guarded because the meta formula was expected to have taught # us this equality. assert sizevars.size_hint(new_size[i] - old_size[i], fallback=0) == 0, ( "Broadcast failed in ExpandView({x.get_size()}, {new_size}) on dimension {i}" ) return new_size @classmethod def create(cls, x: IRNode, new_size: Sequence[_IntLike]) -> BaseView: new_size = cls._normalize_size(x, new_size) if is_storage_and_layout(x): storage, old_layout = as_storage_and_layout(x) skip = len(new_size) - len(old_layout.size) assert skip >= 0 new_stride = [sympy.S.Zero] * skip for stride, size in zip(old_layout.stride, old_layout.size): new_stride.append( stride if not V.graph.sizevars.is_size_one_or_false(size) else sympy.S.Zero ) new_layout = FixedLayout( old_layout.device, old_layout.dtype, list(new_size), new_stride, old_layout.offset, old_layout.is_pinned, ) return ReinterpretView(data=storage, layout=new_layout) return ExpandView(data=x, size=new_size) def get_size(self) -> Sequence[Expr]: return self.size def make_reindexer( self, ) -> Callable[[Sequence[Expr]], Sequence[Expr]]: target = self.get_size() actual = self.data.get_size() skip = len(target) - len(actual) def reindex( index: Sequence[Expr], ) -> Sequence[Expr]: index = list(index[skip:]) assert len(index) == len(actual) for i in range(len(actual)): if actual[i] == 1: # zero out broadcast dimension index[i] = sympy.S.Zero return index return reindex @ir_dataclass class PermuteView(BaseView): dims: list[Expr] @classmethod def create(cls, x: IRNode, dims: Sequence[int]) -> BaseView: dims = cls._map_neg_dims(dims) assert OrderedSet(dims) == OrderedSet(range(len(dims))) if is_storage_and_layout(x): storage, old_layout = as_storage_and_layout(x) new_layout = FixedLayout( old_layout.device, old_layout.dtype, [old_layout.size[i] for i in dims], [old_layout.stride[i] for i in dims], old_layout.offset, old_layout.is_pinned, ) return ReinterpretView(data=storage, layout=new_layout) return PermuteView(data=x, dims=dims) @classmethod def _map_neg_dims(cls, dims: Sequence[int]) -> list[int]: return [dim if dim >= 0 else len(dims) + dim for dim in dims] def get_size(self) -> Sequence[Expr]: assert OrderedSet(self._map_neg_dims(self.dims)) == OrderedSet( range(len(self.dims)) ) size = self.data.get_size() return [size[i] for i in self.dims] def make_reindexer( self, ) -> Callable[[Sequence[Expr]], Sequence[Expr]]: inv = {j: i for i, j in enumerate(self.dims)} inv = [inv[i] for i in range(len(self.dims))] assert OrderedSet(inv) == OrderedSet(range(len(self.dims))) def reindex( index: Sequence[Expr], ) -> Sequence[Expr]: return [index[i] for i in inv] return reindex @ir_dataclass class SqueezeView(BaseView): @classmethod def create(cls, x: IRNode, *, dim: Optional[int] = None) -> IRNode: if is_storage_and_layout(x): storage, old_layout = as_storage_and_layout(x) new_size = [] new_stride = [] if dim is not None: assert isinstance(dim, int), type(dim) assert 0 <= dim and dim < len(old_layout.size) for i, (size, stride) in enumerate(zip(old_layout.size, old_layout.stride)): if dim is None: if size != 1: new_size.append(size) new_stride.append(stride) else: if i != dim: new_size.append(size) new_stride.append(stride) else: assert size == 1, "expected squeezed size to be 1" new_layout = FixedLayout( old_layout.device, old_layout.dtype, new_size, new_stride, old_layout.offset, old_layout.is_pinned, ) return ReinterpretView(data=storage, layout=new_layout) if dim is None: # redirect to a generic view return View.create(x, [s for s in x.get_size() if s != 1]) else: assert x.get_size()[dim] == 1 return View.create(x, [s for i, s in enumerate(x.get_size()) if i != dim]) @staticmethod def squeezer( size: Sequence[Expr], ) -> tuple[list[int], Callable[[Sequence[Expr]], tuple[Expr]]]: new_size = [s for s in size if s != 1] not_one = [i for i, s in enumerate(size) if s != 1] length = len(size) def reindex(index: Sequence[Expr]) -> tuple[Expr]: assert len(index) == len(not_one), f"{index} {not_one}" new_index = [sympy.S.Zero] * length for idx, s in zip(not_one, index): new_index[idx] = s return tuple(new_index) return new_size, reindex def __init__(self, data: Any) -> None: raise AssertionError("use SqueezeView.create()") @ir_dataclass class GenericView(BaseView): size: Sequence[Expr] reindex: Callable[[Sequence[Expr]], Sequence[Expr]] def make_reindexer( self, ) -> Callable[[Sequence[Expr]], Sequence[Expr]]: return self.reindex def reindex_str(self) -> str: index_old = [ sympy_index_symbol_with_prefix(SymT.INDEX, n) for n in range(len(self.size)) ] index_new = list(self.reindex(index_old)) return f"lambda {', '.join(map(str, index_old))}: {index_new}" def __str__(self) -> str: return self.str_helper( [self.data, f"size={self.size}", f"reindex={self.reindex_str()}"] ) __repr__ = __str__ @classmethod def create( cls, x: IRNode, new_size: Sequence[Expr], reindex: Callable[[Sequence[Expr]], Sequence[Expr]], ) -> BaseView: return cls(data=x, size=list(new_size), reindex=reindex) def get_size(self) -> Sequence[Expr]: return self.size @ir_dataclass class View(GenericView): @staticmethod def handle_negative_index(idx: Expr, size: Expr) -> Expr: idx = sympy.expand(idx) size = sympy.expand(size) evaluate_expr = V.graph.sizevars.shape_env.evaluate_expr if evaluate_expr(sympy.Lt(idx, 0)): idx = idx + size return idx @classmethod def create(cls, x: IRNode, new_size: Sequence[Expr]) -> IRNode: # type: ignore[override] assert isinstance(new_size, Sequence), type(new_size) old_size, new_size = cls.resolve_negative_size(x.get_size(), new_size) # Skip pointless views if V.graph.sizevars.statically_known_list_equals(old_size, new_size): return x unbacked_symbols_in_sizes = False if ( len(free_unbacked_symbols(old_size)) > 0 or len(free_unbacked_symbols(new_size)) > 0 ): unbacked_symbols_in_sizes = True if 0 in new_size: def fake_reindex(index: Any) -> tuple[int, ...]: return tuple([0] * len(old_size)) return cls(data=x, size=list(new_size), reindex=fake_reindex) # TODO: a new class for FixedTransferLayout that output layout is constrained by input layout elif is_contiguous_storage_and_layout(x) or unbacked_symbols_in_sizes: if unbacked_symbols_in_sizes and (not is_contiguous_storage_and_layout(x)): # realize x; otherwise, the dynamic_reshape_indexer below will fail # due to the size_hint's inability to process unbacked SymInts # TODO: unbacked should not diverge from backed in determining striding # Need to require contiguous here instead of realize, see: # https://github.com/pytorch/pytorch/issues/145561 x = ExternKernel.require_contiguous(x) storage, old_layout = as_storage_and_layout(x, want_contiguous=True) new_layout = FixedLayout( old_layout.device, old_layout.dtype, new_size, FlexibleLayout.contiguous_strides(new_size), old_layout.offset, old_layout.is_pinned, ) return ReinterpretView(data=storage, layout=new_layout) reindex = cls.dynamic_reshape_indexer(old_size, new_size) return cls(data=x, size=list(new_size), reindex=reindex) @staticmethod def resolve_negative_size( old_size: Sequence[Expr], new_size: Sequence[Expr] ) -> tuple[list[Expr], list[Expr]]: new_size = [V.graph.sizevars.simplify(x) for x in new_size] old_size = [V.graph.sizevars.simplify(x) for x in old_size] new_size = list(new_size) for i in range(len(new_size)): if new_size[i] == -1: new_size[i] = sympy.S.One new_size[i] = CleanDiv(sympy_product(old_size), sympy_product(new_size)) break V.graph.sizevars.check_equals(sympy_product(old_size), sympy_product(new_size)) return old_size, new_size @classmethod def dynamic_reshape_indexer( cls, old_size: Sequence[_IntLike], new_size: Sequence[_IntLike], dense_dim: Optional[int] = None, ) -> Callable[[Sequence[_T]], Sequence[_V]]: try: reindex = cls._dynamic_reshape_indexer(old_size, new_size, dense_dim) except (AssertionError, IndexError): # optimistic algorithm failed, lets do a fallback flat = [sympy_product(old_size)] reindex1 = cls._dynamic_reshape_indexer(old_size, flat) reindex2 = cls._dynamic_reshape_indexer(flat, new_size) reindex = fuse_reindexing(reindex1, reindex2) return reindex @staticmethod def _dynamic_reshape_indexer( old_size: Sequence[Expr], new_size: Sequence[Expr], dense_dim: Optional[int] = None, ) -> Callable[[Sequence[Expr]], Sequence[Expr]]: """ Perform a reshape entirely by modifying indexing math """ size_hint = V.graph.sizevars.size_hint # TODO: These symbols may not escape, if they don't assert so and # treat them as temporary vars = [ sympy_index_symbol_with_prefix(SymT.VIEW, i) for i in range(len(new_size)) ] stack_new = list(zip(vars, new_size)) stack_old = list(old_size) # process the dense dim first reordering_dense_dim = ( dense_dim is not None and dense_dim != len(stack_old) - 1 and len(new_size) == 1 ) if reordering_dense_dim: assert dense_dim is not None # mypy old_dim = stack_old.pop(dense_dim) stack_old.append(old_dim) view_expr = [] while stack_new and stack_old: size_old = stack_old.pop() var, size_new = stack_new.pop() if size_old == 1: view_expr.append(sympy.S.Zero) stack_new.append((var, size_new)) # re-add elif size_new == 1: stack_old.append(size_old) # re-add elif size_hint(size_new) == size_hint(size_old): view_expr.append(var) V.graph.sizevars.check_equals(size_new, size_old) elif size_hint(size_new) < size_hint(size_old): while size_hint(size_new) < size_hint(size_old): var2, size_new2 = stack_new.pop() var = var2 * size_new + var size_new = size_new * size_new2 view_expr.append(var) V.graph.sizevars.check_equals(size_new, size_old) elif size_hint(size_new) > size_hint(size_old): divisor = sympy.S.One modulus = size_old view_expr.append(ModularIndexing(var, divisor, modulus)) divisor = divisor * modulus while size_hint(size_new) > size_hint(size_old): modulus = stack_old.pop() view_expr.append(ModularIndexing(var, divisor, modulus)) divisor = divisor * modulus size_old = size_old * modulus V.graph.sizevars.check_equals(size_new, size_old) else: raise AssertionError while stack_old: size_old = stack_old.pop() V.graph.sizevars.check_equals(size_old, 1) view_expr.append(sympy.S.Zero) while stack_new: var, size_new = stack_new.pop() V.graph.sizevars.check_equals(size_new, 1) if dense_dim is not None and len(new_size) == 1: view_expr.reverse() # Move the last expression (dense dim) to its original position dense_expr = view_expr.pop() view_expr.insert(dense_dim, dense_expr) else: view_expr.reverse() assert len(view_expr) == len(old_size) def reindex( index: Sequence[Expr], ) -> Sequence[Expr]: assert len(index) == len(vars), (len(index), len(vars)) replacements = dict(zip(vars, index)) return tuple(sympy_subs(x, replacements) for x in view_expr) return reindex @ir_dataclass class ReinterpretView(BaseView): """Pretend our storage has a different layout""" layout: Layout def __post_init__(self) -> None: super().__post_init__() if isinstance(self.data, BaseView): object.__setattr__(self, "data", self.data.unwrap_view()) def __str__(self) -> str: return self.str_helper( [ self.data, self.layout, ] ) __repr__ = __str__ def get_name(self) -> str: return self.data.get_name() def get_device(self) -> Optional[torch.device]: return self.layout.device def get_origin_node(self) -> Optional[torch.fx.Node]: return None @property def dtype(self) -> torch.dtype: return self.layout.dtype def get_size(self) -> Sequence[Expr]: return list(self.layout.size) def get_stride(self) -> Sequence[Expr]: return list(self.layout.stride) def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]: def loader(index: Sequence[Expr]) -> OpsValue: indexer = self.layout.make_indexer() tmp_loader = ops.load(self.get_name(), indexer(index)) if self.layout.dtype != self.data.dtype: return ops.to_dtype_bitcast(tmp_loader, self.dtype, self.data.dtype) else: return tmp_loader return loader def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]: return self.layout.make_indexer() def get_layout(self) -> Layout: return self.layout def freeze_layout(self) -> None: pass def get_free_symbol_uses( self, unbacked_only: bool = False ) -> OrderedSet[sympy.Symbol]: return ( get_free_symbols(self.layout.size, unbacked_only) | get_free_symbols(self.layout.stride, unbacked_only) | get_free_symbols(self.layout.offset, unbacked_only) ) def codegen_reference(self, writer: Optional[IndentedBuffer] = None) -> str: # reinterpret_tensor is similar to as_strided except: # - offset is added to the existing offset (rather than replacing it) # - view tracking is disabled similar to unsafe_view return V.graph.wrapper_code.codegen_reinterpret_view( self.data, self.layout.size, self.layout.stride, self.layout.offset, writer.writeline if writer is not None else V.graph.wrapper_code.writeline, dtype=self.layout.dtype, ) def num_reads(self) -> int: return 1 @ir_dataclass class DtypeView(BaseView): """Pretend our storage has a different type""" target_dtype: torch.dtype @classmethod def create(cls, x: IRNode, new_dtype: torch.dtype) -> BaseView: if is_storage_and_layout(x): storage, old_layout = as_storage_and_layout(x) new_layout = FixedLayout( old_layout.device, new_dtype, old_layout.size, old_layout.stride, old_layout.offset, old_layout.is_pinned, ) return ReinterpretView(data=storage, layout=new_layout) return DtypeView(data=x, target_dtype=new_dtype) def __str__(self) -> str: return self.str_helper([self.data, self.target_dtype]) __repr__ = __str__ @property def dtype(self) -> torch.dtype: return self.target_dtype def get_size(self) -> Sequence[Expr]: return self.data.get_size() def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]: inner = self.data.make_loader() def loader(idx: Sequence[Expr]) -> OpsValue: return ops.to_dtype_bitcast(inner(idx), self.target_dtype, self.data.dtype) return loader class SliceView(View): @classmethod def normalize_start_end( cls, x: IRNode, dim: int, start: int, end: int ) -> tuple[int, int]: """ Normalize start and end such that both are in the range [0, x.get_size()[dim]] and start <= end. """ sizevars = V.graph.sizevars dim_size = x.get_size()[dim] if any(free_unbacked_symbols(x) for x in (start, end, dim_size)): min_func = sympy.Min max_func = sympy.Max else: min_func = sizevars.evaluate_min max_func = sizevars.evaluate_max def clamp(x: Expr, lower: int, upper: int) -> Expr: clamped_lower = ( x if sizevars.statically_known_geq(x, lower) else max_func(x, lower) ) clamped_full = ( clamped_lower if sizevars.statically_known_leq(clamped_lower, upper) else min_func(clamped_lower, upper) ) return clamped_full def clamp_wrap( val: Union[int, None], lower: int, upper: int, default: Union[Expr, int] ) -> Union[Expr, int]: if val is None: # TODO(rec): can this really happen? return default val = cls.handle_negative_index(val, dim_size) return clamp(val, lower, upper) start = clamp_wrap(start, 0, dim_size, 0) end = clamp_wrap(end, start, dim_size, dim_size) return start, end @classmethod def create( # type: ignore[override] cls, x: IRNode, dim: int, start: int, end: int, step: int = 1, clamp: bool = True, ) -> IRNode: step = sympy.expand(step) assert isinstance(step, Expr) or step > 0, step try: if start == 0 and end >= 2**63 - 1 and step == 1: return x except TypeError: pass new_size = list(x.get_size()) # NB: Ordinarily we default to clamping. # We only don't clamp for split_with_sizes. For split_with_sizes, sizes should be already valid # failing in this situation is ok, since invalid sizes could trigger silent errors. if clamp: start, end = cls.normalize_start_end(x, dim, start, end) new_size[dim] = FloorDiv(end - start + (step - 1), step) if is_storage_and_layout(x): # Fast path storage, old_layout = as_storage_and_layout(x) new_stride = list(old_layout.stride) new_stride[dim] = new_stride[dim] * step new_layout = FixedLayout( old_layout.device, old_layout.dtype, new_size, new_stride, old_layout.offset + old_layout.stride[dim] * start, old_layout.is_pinned, ) return ReinterpretView(data=storage, layout=new_layout) def reindex( index: Sequence[Expr], ) -> Sequence[Expr]: assert len(index) == len(new_size), f"wrong ndim {index} {new_size}" index = list(index) index[dim] = index[dim] * step + start return index # redirect to a generic view return SliceView(data=x, size=new_size, reindex=reindex) @ir_dataclass class BaseConstant(IRNode): dtype: torch.dtype device: torch.device def get_size(self) -> Sequence[Expr]: return () def get_device(self) -> Optional[torch.device]: return self.device def get_origin_node(self) -> Optional[torch.fx.Node]: return None def get_reads(self) -> OrderedSet[Dep]: return OrderedSet() @ir_dataclass class Constant(BaseConstant): value: Any dtype: torch.dtype device: torch.device def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]: def loader(index: Sequence[Expr]) -> OpsValue: return ops.constant(self.value, self.dtype) return loader def realize(self) -> Optional[str]: pass def constant_to_device(self, device: torch.device) -> IRNode: return Constant(value=self.value, dtype=self.dtype, device=device) @ir_dataclass class IndexingConstant(BaseConstant): index: Any dtype: torch.dtype device: torch.device def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]: def loader(index: Sequence[Expr]) -> OpsValue: return ops.index_expr(self.index, self.dtype) return loader def constant_to_device(self, device: torch.device) -> IRNode: return IndexingConstant(index=self.index, dtype=self.dtype, device=device) def is_contiguous_strides_for_shape( stride: Sequence[_IntLike], shape: Sequence[_IntLike] ) -> bool: expected_stride = 1 expected_stride_max = 1 for x, y in reversed(tuple(zip(shape, stride))): if x == 1: continue if not V.graph.sizevars.statically_known_equals( y, expected_stride ) and not V.graph.sizevars.statically_known_equals(y, expected_stride_max): return False expected_stride_max *= sympy.Max(1, x) expected_stride *= x return True def get_align_for_dtype(dtype: torch.dtype) -> int: return config.padding_alignment_bytes // dtype.itemsize class OutputSpec: """Abstract base for Layout, MultiOutputLayout, NoneLayout. Represents the memory layout of the output of an Operation.""" def get_device(self) -> Optional[torch.device]: raise NotImplementedError(type(self).__name__) def storage_size(self) -> int: raise NotImplementedError(type(self).__name__) def get_free_symbol_uses( self, unbacked_only: bool = False ) -> OrderedSet[sympy.Symbol]: raise NotImplementedError(type(self).__name__) @ir_dataclass class Layout(OutputSpec): """ Layout base class Carries tensor meta-information including offset and whether it is pinned. """ def __init__( self, device: torch.device, dtype: torch.dtype, size: Sequence[Expr], stride: Optional[Sequence[Expr]] = None, offset: Expr = Integer(0), is_pinned: bool = False, ) -> None: if stride is None: stride = FlexibleLayout.contiguous_strides(size) self.device = device self.dtype = dtype assert len(size) == len(stride), f"size={size}, stride={stride}" assert all(isinstance(s, (Expr, int)) for s in size) self.size = size self.stride = stride self.offset = offset self.is_pinned = is_pinned # is_pinned implies cpu assert (not self.is_pinned) or (self.device.type == "cpu") def __str__(self) -> str: offset = "" if self.offset != 0: offset = f", offset={self.offset}" device_index_str = "" if self.device.index is None else f":{self.device.index}" is_pinned_str = "" if self.is_pinned: is_pinned_str = f", is_pinned={self.is_pinned}" return ( f"{type(self).__name__}('{self.device.type}{device_index_str}', {self.dtype}, " f"size={self.size}, stride={self.stride}{offset}{is_pinned_str})" ) __repr__ = __str__ def get_device(self) -> torch.device: return self.device def get_example(self) -> torch.Tensor: with V.fake_mode: return torch.empty_strided( convert_shape_to_symint(self.size), convert_shape_to_symint(self.stride), dtype=self.dtype, device=self.device, pin_memory=self.is_pinned, ) def is_contiguous(self) -> bool: return is_contiguous_strides_for_shape(self.stride, self.size) @staticmethod def is_channels_last_contiguous( shape: Sequence[_IntLike], strides: Sequence[_IntLike] ) -> bool: ndim = len(shape) if ndim not in [4, 5] or shape[1] == 1: return False for left, right, size in zip( strides, make_channels_last_strides_for(shape), shape ): if size != 1 and left != right: return False return True def is_transposed(self) -> bool: for left, right, size in zip( self.stride, reversed(FlexibleLayout.contiguous_strides(list(reversed(self.size)))), self.size, ): if size != 1 and left != right: return False return True def is_stride_ordered(self, order: Sequence[int]) -> bool: assert len(self.stride) == len(order) # ignore dimensions of size 1, they dont affect layout non_1_indices = [ i for i, dim in enumerate(self.size) if V.graph.sizevars.size_hint(dim, fallback=2) != 1 ] stride = [self.stride[i] for i in non_1_indices] order: Sequence[int] = [order[i] for i in non_1_indices] def sorted_indices(arr: Sequence[int]) -> Sequence[int]: sorted_arr = sorted(arr) return [sorted_arr.index(element) for element in arr] # since we may have removed dimensions, need to re-sort & re-index order order = sorted_indices(order) # reorder the stride given order stride_ordered = [-1] * len(order) for i in range(len(order)): stride_ordered[order[i]] = stride[i] # check if it is in ascending order for i in range(len(order) - 1): expr = stride_ordered[i] > stride_ordered[i + 1] if not isinstance(expr, bool): expr = V.graph._shape_env.evaluate_expr( stride_ordered[i] > stride_ordered[i + 1], size_oblivious=True ) if expr: return False return True def is_channels_last_stride_ordered(self) -> bool: # create channels_last order(NCHW, NCDHW, the C is the first order). order = [0] + list(reversed(range(1, len(self.stride) - 1))) order = [len(order)] + order return self.is_stride_ordered(order) @staticmethod def _pad_strides( in_strides: Sequence[int], size: Sequence[Expr], dtype: torch.dtype ) -> Sequence[int]: """ The padding does not change stride order but makes sure all strides larger than the threshold are multiple of align. """ align = get_align_for_dtype(dtype) if len(in_strides) == 0: return in_strides if not config.pad_channels_last and Layout.is_channels_last_contiguous( size, in_strides ): return in_strides current_fx_node = V.get_current_node() if hasattr(current_fx_node, "meta") and current_fx_node.meta.get( "dislike_padding", False ): return in_strides shape_env = V.graph._shape_env if hasattr(V.graph, "_shape_env") else None def contains_unbacked_symints(expr: sympy.Expr | int) -> bool: if shape_env is None: return False if not isinstance(expr, sympy.Expr): return False return any(shape_env.is_unbacked_symint(s) for s in expr.free_symbols) # Skip padding the strides when it contains unbacked symints for now. if shape_env and any(contains_unbacked_symints(s) for s in in_strides): return in_strides stride_order = get_stride_order(in_strides, shape_env) fill_order = stride_order2fill_order(stride_order) new_strides = [0 for _ in range(len(in_strides))] # since we pad when the layout is flexible, we can decide the # smallest stride to be 1. new_strides[fill_order[0]] = 1 padded = False for rank, idx in enumerate(fill_order[1:], start=1): prev_idx = fill_order[rank - 1] stride = new_strides[prev_idx] * size[prev_idx] # Static stride and meets padding conditions OR # Dynamic stride and config.pad_dynamic_shape=True require_padding = ( isinstance(stride, (int, sympy.Integer)) and stride > config.padding_stride_threshold and stride % align != 0 ) or (isinstance(stride, sympy.Expr) and config.pad_dynamic_shapes) new_strides[idx] = stride if require_padding: new_strides[idx] = ceildiv(stride, align) * align padded = True if not padded: # Consider a tensor with shape [256, 1, 5, 5] # Avoid strides like [25, 5, 5, 1] being padded to equivalent strides # [25, 25, 5, 1]. return in_strides metrics.num_comprehensive_padding += 1 return new_strides def pad_strides(self) -> None: assert isinstance(self, FlexibleLayout), type(self) assert self.stride is not None self.stride = self._pad_strides(self.stride, self.size, self.dtype) def should_pad_strides(self) -> bool: return config.comprehensive_padding and isinstance(self, FlexibleLayout) def as_fixed(self) -> FixedLayout: if isinstance(self, FixedLayout): return self if self.should_pad_strides(): self.pad_strides() return FixedLayout( self.device, self.dtype, self.size, self.stride, self.offset, self.is_pinned, ) def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]: assert FlexibleLayout.allow_indexing, ( f"convert {type(self).__name__} to FixedLayout first" ) return self.as_fixed().make_indexer() def __eq__(self, other: object) -> bool: return ( isinstance(other, Layout) and self.device == other.device and self.dtype == other.dtype and self.size == other.size and self.stride == other.stride and self.offset == other.offset and self.is_pinned == other.is_pinned ) def storage_size(self) -> Expr: return compute_required_storage_length(self.size, self.stride, self.offset) # type: ignore[arg-type] def get_free_symbol_uses( self, unbacked_only: bool = False ) -> OrderedSet[sympy.Symbol]: return ( get_free_symbols(self.size, unbacked_only) | get_free_symbols(self.stride, unbacked_only) | get_free_symbols(self.offset, unbacked_only) ) class FixedLayout(Layout): """A Tensor layout we cannot change""" def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]: """A closure containing math to read a given element""" return _fixed_indexer(self.size, self.stride, self.offset) class FlexibleLayout(Layout): """A Tensor layout that we are allowed to change""" allow_indexing = False # WARNING! This doesn't handle zero size tensors correctly @staticmethod def contiguous_strides(sizes: Sequence[int]) -> list[Expr]: if len(sizes) == 0: return [] reversed_strides = [sympy.S.One] for size in reversed(sizes[1:]): reversed_strides.append(size * reversed_strides[-1]) return list(reversed(reversed_strides)) @staticmethod def fill_ordered(sizes: Sequence[int], order: Sequence[int]) -> list[Expr]: """ Create a stride based on the order the dimensions should be filled in. In this format, channels last would be: [1, 3, 2, 0] """ assert OrderedSet(range(len(sizes))) == OrderedSet(order), (sizes, order) next_stride = sympy.S.One strides = [None] * len(order) for i in order: strides[i] = next_stride next_stride = next_stride * sizes[i] return strides @staticmethod def stride_ordered(sizes: Sequence[int], order: Sequence[int]) -> Sequence[Expr]: """ Create a stride based on the sorted order of a permuted range. In this format, channels last would be: [3, 0, 2, 1] """ assert OrderedSet(range(len(sizes))) == OrderedSet(order) fill_order = stride_order2fill_order(order) return FlexibleLayout.fill_ordered(sizes, fill_order) @staticmethod def stride_ordered_for_memory_format( sizes: Sequence[int], memory_format: torch.memory_format ) -> Sequence[Expr]: """ Create a stride based on a memory format. Memory format is translasted into a stride order, so channels_last is the same as: FlexibleLayout.stride_ordered(sizes, [3, 0, 2, 1]) This interface does not support memory_format `torch.preserve_format` which should be used to deduce a format from another source """ if memory_format == torch.channels_last: return FlexibleLayout.stride_ordered(sizes, NHWC_STRIDE_ORDER) elif memory_format == torch.channels_last_3d: return FlexibleLayout.stride_ordered(sizes, NHWDC_STRIDE_ORDER) elif memory_format == torch.contiguous_format: return FlexibleLayout.contiguous_strides(sizes) else: log.debug( "stride_ordered_for_memory_format, unsuppored memory_format: %s", memory_format, ) raise NotImplementedError @staticmethod def same_ordered( sizes: Sequence[int], stride: Sequence[_IntLike] ) -> Sequence[Expr]: """ Create a stride that has the same stride order as given stride For example, if given stride is [1000, 1, 100, 10], the fill order should be [1, 3, 2, 0] """ assert len(sizes) == len(stride) stride = [V.graph.sizevars.size_hint_or_throw(x) for x in stride] fill_order = sorted(range(len(stride)), key=stride.__getitem__) return FlexibleLayout.fill_ordered(sizes, fill_order) def as_stride_order( self, order: Sequence[int], allow_padding: bool = False ) -> FixedLayout: new_stride = self.stride_ordered(self.size, order) if self.should_pad_strides() and allow_padding: new_stride = self._pad_strides(new_stride, self.size, self.dtype) return FixedLayout( self.device, self.dtype, self.size, new_stride, self.offset, self.is_pinned, ) def as_exact_strides( self, exact_strides: Sequence[_IntLike], allow_padding: bool = False ) -> FixedLayout: new_stride = exact_strides if self.should_pad_strides() and allow_padding: new_stride = self._pad_strides(new_stride, self.size, self.dtype) return FixedLayout( self.device, self.dtype, self.size, new_stride, self.offset, self.is_pinned, ) def as_fill_order(self, order: Sequence[int]) -> FixedLayout: new_stride: Sequence[int] = self.fill_ordered(self.size, order) if self.should_pad_strides(): new_stride = self._pad_strides(new_stride, self.size, self.dtype) return FixedLayout( self.device, self.dtype, self.size, new_stride, self.offset, self.is_pinned, ) def as_same_order(self, stride: Sequence[_IntLike]) -> FixedLayout: new_stride = self.same_ordered(self.size, stride) if self.should_pad_strides(): new_stride = self._pad_strides(new_stride, self.size, self.dtype) return FixedLayout( self.device, self.dtype, self.size, new_stride, self.offset, self.is_pinned, ) def __init__( self, device: torch.device, dtype: torch.dtype, size: Sequence[Expr], stride_order: Optional[Sequence[Union[int, Integer]]] = None, is_pinned: bool = False, ) -> None: if stride_order: strides = FlexibleLayout.fill_ordered(size, stride_order) else: strides = FlexibleLayout.contiguous_strides(size) super().__init__(device, dtype, size, strides, is_pinned=is_pinned) class NonOwningLayout(Layout): """Is a view into the storage of another tensor""" def __init__(self, view: Union[BaseView, TensorBox]) -> None: layout = view.get_layout() super().__init__( layout.device, layout.dtype, layout.size, layout.stride, ) self.view = view def make_indexer(self) -> Callable[[Sequence[Expr]], Expr]: return self.as_fixed().make_indexer() def maybe_guard_aligned(self) -> bool: offset = self.view.get_layout().offset if offset == 0: return True from .utils import ALIGNMENT return V.graph.sizevars.statically_known_multiple_of(offset, ALIGNMENT) def get_free_symbol_uses( self, unbacked_only: bool = False ) -> OrderedSet[sympy.Symbol]: assert isinstance(self.view, ReinterpretView) box = self.view.data assert isinstance(box, StorageBox), type(box) input_buffer = box.data assert isinstance(input_buffer, Buffer), type(box) return input_buffer.layout.get_free_symbol_uses(unbacked_only) class CommBufferType(Enum): SYMM_MEM = "symm_mem" class CommBufferLayout(FixedLayout): """ A layout that signifies the buffer is a comm buffer. In terms of striding, the layout is identical to `FixedLayout`. Buffers with this layout do not participate in in-place reuse - it can be neither the source nor the target for in-place reuse. For detailed motivation and usage of this layout, see NOTE [lowering-time collective optimization]. """ comm_buffer_type: CommBufferType group_name: str def __init__( self, layout: FlexibleLayout, comm_buffer_type: CommBufferType, group_name: str, ): if not isinstance(layout, FlexibleLayout): raise AssertionError( "A `CommBufferLayout` can only be initialized with " f"a `FlexibleLayout` (got {layout})." ) fixed = layout.as_fixed() super().__init__( device=fixed.device, dtype=fixed.dtype, size=fixed.size, stride=fixed.stride, offset=fixed.offset, is_pinned=fixed.is_pinned, ) self.comm_buffer_type = comm_buffer_type self.group_name = group_name @ir_dataclass class NoneLayout(OutputSpec): # This is janky, I figured out what fields to populate by just running # the model I was interested in and adding properties/methods as needed. # This doesn't inherit from Layout because Layout assumes you have stuff # like sizes, but I don't really have anything here. # # If you have an ir.Node with NoneLayout, you probably need to setup # dependencies manually in scheduler device: Optional[torch.device] size: list[int] = dataclasses.field(default_factory=lambda: [0]) stride: list[int] = dataclasses.field(default_factory=lambda: [0]) def storage_size(self) -> int: return 0 def as_fixed(self) -> OutputSpec: return self def get_device(self) -> Optional[torch.device]: return self.device class MutationLayoutSHOULDREMOVE(Layout): def __init__(self, target: IRNode) -> None: super().__init__( target.get_device_or_error(), targ "trust_building_protocol": "Increase in secure base score leads to decreased defensive modes and increased pro-social behaviors." } }, { "op": "add", "path": "/emotional_expression/register_safe_emotions", "value": { "attachment_secure": [ "comfortable", "safe", "relaxed" ], "attachment_anxious": [ "nervous", "worried", "uncertain" ] } }, { "op": "replace", "path": "/emotional_modulation/emotion_thresholds", "value": { "secure_base_present": 0.8, "secure_base_absent": 0.3 } } ] } }, "status": "skipped" }, { "timestamp": "2025-12-15T10:59:21.923374", "improvement": { "improvement": "Add better error handling for API requests", "applies_to": "FastAPI framework", "why": "The developer tools market research highlighted the importance of robust error messages and user feedback. QAOA counterdiabatic driving knowledge suggests graceful degradation of operations which translates to softer error handling. Phi fractal patterns consciousness can inform design by avoiding nested callbacks that spiral out of control.", "implementation": "Enhance FastAPI's exception handlers to include customizable error responses with actionable steps for developers. Implement try/except blocks around request processing and use async generators for non-blocking error notifications." }, "status": "skipped" }, { "timestamp": "2025-12-15T11:04:25.500457", "improvement": { "improvement": "Add emotional intelligence training for AGI communication", "applies_to": "/Eden/CORE/emotional_intelligence", "why": "emotional intelligence research 2025 suggests focusing on EQ in Gen Z and AI-enhanced emotional skills. This will improve the AGI's conversational empathy.", "implementation": "Create new methods in emotion_analyzer.py for context-aware responses and leaderhip training. Update chat_prompt_builder to include emotional nuances." }, "status": "skipped" }, { "timestamp": "2025-12-15T11:09:29.726632", "improvement": { "improvement": "Optimize error handling in the backend API.", "applies_to": "api/middleware/error_handler.py", "why": "Understanding quantum error mitigation suggests that incorporating similar principles of mitigation can improve our current error handling. Specifically, implementing a probabilistic approach to catch and mitigate errors before they propagate.", "implementation": { "code_change": "Replace the existing fixed error handling with a probabilistic error mitigation algorithm. This involves adding a phase angle determination based on quantum oscillation patterns to dynamically adjust error thresholds.", "specific_modification": "Add a new function `probabilistic_error_mitigation` in api/middleware/error_handler.py that takes into account the research on Grover's algorithm for searching good states. The function should dynamically adjust error tolerance and retry mechanisms based on this probabilistic approach." } }, "status": "skipped" } ], "active_plugins": 27, "capabilities": { "total_capabilities": 1621645, "hot_cache": 256, "warm_cache": 0, "cold_storage": 1621389 }, "tie": { "timestamp": "2025-12-17T09:04:50.768020", "cycle": 36214, "entropy": 0.8764, "order": 0.1333, "energy_total": 0.3433, "energy_cpu": 0.0659, "energy_memory": 0.1264, "energy_gpu": 0.93, "temperature": 0.545, "delta_order": 0.0, "delta_energy": -0.015, "efficiency": 0.0666, "mode": "balanced", "phi_factor": 1.618 } }, "insights": [ "Active learning: 5 recent topics", "Self-evolution: 5 improvements" ], "actions": [ { "type": "maintain_awareness" }, { "type": "coordinate_learning_evolution" } ] }