_tensor_iterator.py

"""Python access to ATen's TensorIterator build pipeline.

:class:`TensorIterator` builds an iterator from a set of operands and flags
that mirror ``at::TensorIteratorConfig``, then exposes the post-build
shape / dtype / device / stride information for inspection.

This is a build-pipeline-only surface: there is no ``for_each`` here. Use it
to debug shape and dtype inference, validate custom-op contracts, or inspect
how ATen would lay out a kernel's iteration.

Construction is canonical, *not* a faithful replay of arbitrary
``TensorIteratorConfig`` call sequences. Operands are always registered in
the order ``outputs -> inputs -> const_inputs`` and setters are applied in
one fixed order. Notably:

* The C++ builder distinguishes ``add_input(a); add_const_input(b)`` from
  ``add_const_input(b); add_input(a)`` -- ``input(0)`` refers to different
  operands. This Python surface cannot express that distinction: every
  ``inputs[i]`` precedes every ``const_inputs[j]``.
* Some C++ setters have order-dependent side effects (e.g.
  ``promote_inputs_to_common_dtype(true)`` also flips
  ``check_all_same_dtype`` to ``false``). The Python surface materializes
  the *final* boolean state of each knob, not the call order, so it can't
  reproduce a sequence where an intermediate setter observed a
  since-overwritten value.

Every in-tree caller of ``at::TensorIteratorConfig`` fits the canonical-
recipe shape; the lossiness is theoretical, not practical.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch
from torch._C import (
    _TensorIterator as _CTensorIterator,
    _TensorIteratorSpec as _CTensorIteratorSpec,
)


if TYPE_CHECKING:
    from collections.abc import Sequence


__all__ = [
    "TensorIterator",
    "binary_op",
    "binary_float_op",
    "comparison_op",
    "unary_op",
    "unary_float_op",
    "nullary_op",
    "reduce_op",
]


class TensorIterator:
    """A built TensorIterator. Read-only view of the build-pipeline result.

    Constructor kwargs mirror ``at::TensorIteratorConfig``; defaults match
    the C++ defaults. See module docstring for canonical-recipe caveats.

    Examples
    --------
    Build a TI for a binary op with type promotion and inspect the result::

        >>> import torch
        >>> from torch._tensor_iterator import TensorIterator
        >>> a = torch.zeros(3, 4, dtype=torch.float32)
        >>> b = torch.zeros(3, 4, dtype=torch.float64)
        >>> it = TensorIterator(
        ...     outputs=[None],
        ...     const_inputs=[a, b],
        ...     promote_inputs_to_common_dtype=True,
        ...     cast_common_dtype_to_outputs=True,
        ... )
        >>> it.common_dtype
        torch.float64
        >>> it.numel
        12
        >>> it.ndim       # contiguous (3, 4) collapses to a single dim
        1

    Declare a static shape (the iterator skips broadcast/coalesce)::

        >>> a = torch.zeros(2, 6)
        >>> out = torch.empty(2, 6)
        >>> it = TensorIterator(
        ...     outputs=[out],
        ...     const_inputs=[a],
        ...     resize_outputs=False,
        ...     static_shape=(2, 6),
        ... )

    Build a reduction TI (output must be pre-allocated to the reduced shape)::

        >>> a = torch.zeros(3, 4)
        >>> out = torch.empty(3, 1)
        >>> it = TensorIterator(
        ...     outputs=[out],
        ...     const_inputs=[a],
        ...     resize_outputs=False,
        ...     is_reduction=True,
        ... )

    Inspect post-coalesce strides (in bytes; use ``element_strides`` for
    element units)::

        >>> a = torch.zeros(3, 4)
        >>> b = torch.zeros(3, 4)
        >>> it = TensorIterator(outputs=[None], const_inputs=[a, b])
        >>> tuple(it.strides(0))           # output byte strides
        (4,)
        >>> it.element_strides(0)          # ... in elements
        (1,)
    """

    def __init__(
        self,
        *,
        outputs: list[torch.Tensor | None] | None = None,
        inputs: list[torch.Tensor] | None = None,
        const_inputs: list[torch.Tensor] | None = None,
        check_all_same_dtype: bool = True,
        check_all_same_device: bool = True,
        promote_inputs_to_common_dtype: bool = False,
        promote_integer_inputs_to_float: bool = False,
        cast_common_dtype_to_outputs: bool = False,
        enforce_safe_casting_to_output: bool = False,
        enforce_linear_iteration: bool = False,
        resize_outputs: bool = True,
        check_mem_overlap: bool = True,
        allow_cpu_scalars: bool = False,
        is_reduction: bool = False,
        static_dtype: torch.dtype | None = None,
        static_device: torch.device | None = None,
        static_shape: Sequence[int] | None = None,
        squash_dims: Sequence[int] = (),
    ) -> None:
        spec = _CTensorIteratorSpec()
        spec.outputs = list(outputs) if outputs is not None else []
        spec.inputs = list(inputs) if inputs is not None else []
        spec.const_inputs = list(const_inputs) if const_inputs is not None else []

        spec.check_all_same_dtype = check_all_same_dtype
        spec.check_all_same_device = check_all_same_device
        spec.promote_inputs_to_common_dtype = promote_inputs_to_common_dtype
        spec.promote_integer_inputs_to_float = promote_integer_inputs_to_float
        spec.cast_common_dtype_to_outputs = cast_common_dtype_to_outputs
        spec.enforce_safe_casting_to_output = enforce_safe_casting_to_output
        spec.enforce_linear_iteration = enforce_linear_iteration
        spec.resize_outputs = resize_outputs
        spec.check_mem_overlap = check_mem_overlap
        spec.allow_cpu_scalars = allow_cpu_scalars
        spec.is_reduction = is_reduction

        if static_dtype is not None:
            spec.static_dtype = static_dtype
        if static_device is not None:
            spec.static_device = static_device
        if static_shape is not None:
            spec.static_shape = list(static_shape)
            spec.squash_dims = list(squash_dims)

        self._impl: _CTensorIterator = spec.build()

    @property
    def ndim(self) -> int:
        return self._impl.ndim

    @property
    def shape(self) -> memoryview:
        """Iterator shape (post coalesce/reorder), as a zero-copy
        ``memoryview`` of ``int64`` elements. The view holds a reference
        to this iterator and keeps it alive for as long as the view is
        reachable; copy via ``tuple(it.shape)`` if you need a snapshot
        you can outlive the iterator with. Hot-path readers (dispatch
        conditionals) get index access without an allocation."""
        return self._impl.shape

    @property
    def numel(self) -> int:
        return self._impl.numel

    @property
    def ntensors(self) -> int:
        return self._impl.ntensors

    @property
    def ninputs(self) -> int:
        return self._impl.ninputs

    @property
    def noutputs(self) -> int:
        return self._impl.noutputs

    @property
    def is_contiguous(self) -> bool:
        return self._impl.is_contiguous

    @property
    def is_trivial_1d(self) -> bool:
        return self._impl.is_trivial_1d

    @property
    def common_dtype(self) -> torch.dtype | None:
        """The inferred computation dtype, or ``None`` if no single dtype
        was inferred. Populated whenever TensorIterator can resolve a
        common dtype -- under promotion flags, or when every input
        already shares a dtype. ``None`` does not mean "promotion was
        not requested"; it means inference produced no answer."""
        return self._impl.common_dtype

    def tensor(self, index: int) -> torch.Tensor:
        """Return the iterator's current operand at ``index``.

        Note: under ``promote_inputs_to_common_dtype`` /
        ``cast_common_dtype_to_outputs`` (CPU paths), this may be a
        promoted/cast kernel temporary rather than the tensor that was
        registered with the config -- it's the iterator's view of the
        operand a kernel would actually iterate over."""
        return self._impl.tensor(index)

    def input(self, index: int = 0) -> torch.Tensor:
        """Return the iterator's current input. See :meth:`tensor` for the
        promoted-temporary caveat."""
        return self._impl.input(index)

    def output(self, index: int = 0) -> torch.Tensor:
        """Return the iterator's current output. See :meth:`tensor` for the
        cast-temporary caveat."""
        return self._impl.output(index)

    def dtype(self, index: int = 0) -> torch.dtype:
        return self._impl.dtype(index)

    def device(self, index: int = 0) -> torch.device:
        return self._impl.device(index)

    def strides(self, index: int) -> memoryview:
        """Per-operand strides in bytes (post reorder/coalesce), as a
        zero-copy ``memoryview`` of ``int64`` elements. The view holds a
        reference to this iterator and keeps it alive for as long as the
        view is reachable; copy via ``tuple(it.strides(i))`` if you need
        a snapshot you can outlive the iterator with."""
        return self._impl.strides(index)

    def element_strides(self, index: int) -> tuple[int, ...]:
        """Per-operand strides in elements (byte stride / element size).

        Allocates a fresh tuple on every call. Don't use on a hot path:
        cache the result, or read :meth:`strides` once and divide by
        :meth:`element_size` of the operand inline."""
        return self._impl.element_strides(index)

    def __repr__(self) -> str:
        return repr(self._impl)


# --- Factory shortcuts. These mirror the C++ named constructors at
# aten/src/ATen/TensorIterator.cpp:1069+. Pass ``out=None`` to ask the iterator
# to allocate a fresh output tensor of the inferred shape/dtype/device.


def binary_op(
    out: torch.Tensor | None, a: torch.Tensor, b: torch.Tensor
) -> TensorIterator:
    """Equivalent of ``at::TensorIterator::binary_op``."""
    return TensorIterator(
        outputs=[out],
        const_inputs=[a, b],
        allow_cpu_scalars=True,
        promote_inputs_to_common_dtype=True,
        cast_common_dtype_to_outputs=True,
        enforce_safe_casting_to_output=True,
    )


def binary_float_op(
    out: torch.Tensor | None, a: torch.Tensor, b: torch.Tensor
) -> TensorIterator:
    """Equivalent of ``at::TensorIterator::binary_float_op``."""
    return TensorIterator(
        outputs=[out],
        const_inputs=[a, b],
        allow_cpu_scalars=True,
        promote_inputs_to_common_dtype=True,
        cast_common_dtype_to_outputs=True,
        enforce_safe_casting_to_output=True,
        promote_integer_inputs_to_float=True,
    )


def comparison_op(
    out: torch.Tensor | None, a: torch.Tensor, b: torch.Tensor
) -> TensorIterator:
    """Equivalent of ``at::TensorIterator::comparison_op``.

    When ``out`` is ``None``, the output dtype is forced to bool. When ``out``
    is a defined non-bool tensor, the common dtype is cast back to its dtype
    via ``cast_common_dtype_to_outputs``. The bool-output case skips that cast
    as a performance optimization.
    """
    static_dtype = torch.bool if out is None else None
    cast_to_outputs = out is not None and out.dtype != torch.bool
    return TensorIterator(
        outputs=[out],
        const_inputs=[a, b],
        allow_cpu_scalars=True,
        promote_inputs_to_common_dtype=True,
        cast_common_dtype_to_outputs=cast_to_outputs,
        static_dtype=static_dtype,
    )


def unary_op(out: torch.Tensor | None, a: torch.Tensor) -> TensorIterator:
    """Equivalent of ``at::TensorIterator::unary_op``."""
    return TensorIterator(outputs=[out], const_inputs=[a])


def unary_float_op(out: torch.Tensor | None, a: torch.Tensor) -> TensorIterator:
    """Equivalent of ``at::TensorIterator::unary_float_op``."""
    return TensorIterator(
        outputs=[out],
        const_inputs=[a],
        promote_inputs_to_common_dtype=True,
        cast_common_dtype_to_outputs=True,
        enforce_safe_casting_to_output=True,
        promote_integer_inputs_to_float=True,
    )


def nullary_op(out: torch.Tensor) -> TensorIterator:
    """Equivalent of ``at::TensorIterator::nullary_op``.

    Unlike the binary/unary factories, ``out`` must be a defined tensor;
    the C++ named constructor takes a non-undefined output.
    """
    if out is None:
        raise TypeError(
            "nullary_op requires a defined output tensor; None is not allowed."
        )
    return TensorIterator(
        outputs=[out],
        check_all_same_dtype=False,
        resize_outputs=False,
    )


def reduce_op(
    out: torch.Tensor,
    a: torch.Tensor,
    *,
    out2: torch.Tensor | None = None,
) -> TensorIterator:
    """Equivalent of ``at::TensorIterator::reduce_op``.

    Pass ``out2`` for the two-output reduction overload (e.g. ``min`` returning
    values + indices). The output tensor(s) must be pre-allocated and shaped
    correctly: this factory does not allocate or resize. With ``out2``, both
    outputs must live on ``a``'s device and share its sizes/strides -- the
    C++ named constructor asserts this and the same checks are mirrored
    here.
    """
    # Mirror the TORCH_INTERNAL_ASSERTs in TensorIterator::reduce_op. We
    # surface them as RuntimeError up front since the binding doesn't
    # rebind the named constructor itself; without this, mismatched
    # out1/out2 would build a kernel-shape that doesn't match the
    # tensors and corrupt memory once a kernel is invoked.
    if out2 is not None:
        if not out.device == a.device == out2.device:
            raise RuntimeError(
                "reduce_op: out, out2, and the input must share a device, "
                f"got out={out.device}, out2={out2.device}, a={a.device}"
            )
        if out.dim() != out2.dim():
            raise RuntimeError(
                f"reduce_op: out and out2 must have the same dim, got "
                f"{out.dim()} and {out2.dim()}"
            )
        if out.shape != out2.shape or out.stride() != out2.stride():
            raise RuntimeError(
                "reduce_op: out and out2 must have identical sizes and strides"
            )
        return TensorIterator(
            outputs=[out, out2],
            const_inputs=[a],
            check_all_same_dtype=False,
            resize_outputs=False,
            check_mem_overlap=False,
            is_reduction=True,
        )
    return TensorIterator(
        outputs=[out],
        const_inputs=[a],
        promote_inputs_to_common_dtype=True,
        resize_outputs=False,
        check_mem_overlap=False,
        is_reduction=True,
    )