# mypy: allow-untyped-defs
import logging
import threading
from collections.abc import Callable, Sequence
from contextlib import nullcontext
from functools import lru_cache
from itertools import chain
from typing import cast

import torch
from torch._guards import detect_fake_mode
from torch._logging import LazyString
from torch._ops import OpOverload
from torch._subclasses import FakeTensorMode
from torch.distributed._functional_collectives import _are_we_tracing
from torch.distributed.device_mesh import DeviceMesh
from torch.distributed.tensor._decompositions import DecompShardingStrategy
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
from torch.distributed.tensor._op_schema import (
    OpInfo,
    OpSchema,
    OpSpec,
    OpStrategy,
    OutputSharding,
    OutputSpecType,
    RuntimeSchemaInfo,
    StrategyType,
    TupleStrategy,
)
from torch.distributed.tensor._ops.single_dim_strategy import (
    _expand_single_dim_strategy_to_mesh,
    _SingleDimStrategyInfo,
)
from torch.distributed.tensor._utils import (
    compute_local_shape_and_global_offset,
    compute_local_stride,
    try_find_mesh_from_args,
)
from torch.distributed.tensor.placement_types import _StridedShard, Shard
from torch.utils._pytree import tree_map


aten = torch.ops.aten

log = logging.getLogger(__name__)


def _propagate_use_strided_shard_flag(
    op_strategy: OpStrategy,
    op_schema: OpSchema,
) -> None:
    """Propagate use_strided_shard_as_shard_order from input specs to output specs.

    When inputs carry _StridedShard with an explicit flag, all output (and input)
    DTensorSpecs in the strategy that also contain _StridedShard must agree.
    Strategy functions may forget to propagate the flag; this function fixes
    them up centrally after the strategy is produced.
    """
    _use_strided: bool | None = None
    for spec in op_schema.args_spec:
        if any(isinstance(p, _StridedShard) for p in spec.placements):
            val = spec.use_strided_shard_as_shard_order
            if _use_strided is not None and _use_strided != val:
                _use_strided = True
            else:
                _use_strided = val

    if _use_strided is None:
        return

    def _fixup(spec: DTensorSpec) -> None:
        if not any(isinstance(p, _StridedShard) for p in spec.placements):
            return
        if spec.use_strided_shard_as_shard_order == _use_strided:
            return
        spec.use_strided_shard_as_shard_order = _use_strided
        if _use_strided:
            spec.shard_order = None  # pyrefly: ignore[bad-assignment]
        else:
            spec.shard_order = DTensorSpec.compute_default_shard_order(spec.placements)

    for op_spec in op_strategy.strategies:
        out = op_spec.output_specs
        if out is not None:
            if isinstance(out, DTensorSpec):
                _fixup(out)
            else:
                for s in out:
                    if s is not None:
                        _fixup(s)
        if op_spec.input_specs is not None:
            for s in op_spec.input_specs:
                _fixup(s)


def _length(obj) -> int:
    if obj is None:
        return 0
    if not isinstance(obj, Sequence):
        return 1
    return len(obj)


def _get_expected_num_tensor_outputs(op: OpOverload) -> int | None:
    """
    Get the expected number of tensor outputs for an operator based on its schema.

    Returns:
        The number of tensor outputs expected. Returns 0 for ops that don't return tensors
        (e.g., _linalg_check_errors). Returns 1 for single tensor return, and >1 for
        tuple returns where each element is a tensor. Returns None for List[Tensor]
        returns where the length is unknown at schema time.
    """
    return_types = op._schema.returns
    if len(return_types) == 0:
        return 0

    first_return = return_types[0]
    if isinstance(first_return.type, torch.TensorType):
        # Could be single tensor or tuple of tensors
        return len(return_types)
    elif isinstance(first_return.type, torch.ListType):
        # List[Tensor] - we don't know the length at schema time
        return None
    else:
        # Not a tensor return type
        return 0


def _validate_tensor_meta_count(
    op_schema: OpSchema,
    tensor_meta: TensorMeta | Sequence[TensorMeta | None] | None,
) -> None:
    """
    Validate that the tensor_meta matches the expected number of outputs for the op.

    Raises AssertionError if the count doesn't match, providing a helpful error message.
    """
    expected_outputs = _get_expected_num_tensor_outputs(op_schema.op)

    # Compute actual count:
    # - None means 0 outputs
    # - TensorMeta (single instance) means 1 output
    # - Sequence of TensorMeta means len(sequence) outputs
    # Note: TensorMeta is a NamedTuple (subclass of tuple), so we must check for it first
    if tensor_meta is None:
        actual_outputs = 0
    elif isinstance(tensor_meta, TensorMeta):
        actual_outputs = 1
    else:
        actual_outputs = len(tensor_meta)

    if expected_outputs is None:
        # List[Tensor] return type: length unknown at schema time, but
        # tensor_meta must be a list of TensorMeta.
        if not isinstance(tensor_meta, list):
            raise AssertionError(
                f"Tensor meta for {op_schema.op} should be a list[TensorMeta] "
                f"(op returns List[Tensor]), but got {type(tensor_meta).__name__}"
            )
        return

    if actual_outputs != expected_outputs:
        raise AssertionError(
            f"Tensor meta count mismatch for {op_schema.op}: "
            f"expected {expected_outputs} tensor output(s) based on op schema, "
            f"but _propagate_tensor_meta returned {actual_outputs}. "
            f"This usually indicates a bug in fake tensor propagation for this op."
        )


class LocalLRUCache(threading.local):
    def __init__(self, user_function: Callable) -> None:
        self.cache = lru_cache(None)(user_function)

    def __call__(self, *args, **kwargs) -> object:
        # Fast path: log.handlers check is very cheap (just checking if list is non-empty)
        # Only do the more expensive isEnabledFor check if handlers exist
        if log.handlers and log.isEnabledFor(logging.DEBUG):
            info_before = self.cache.cache_info()
            result = self.cache(*args, **kwargs)
            info_after = self.cache.cache_info()
            cache_hit = info_after.hits > info_before.hits
            op_schema = args[0] if args else None
            output_spec = getattr(result, "output_spec", None)
            log.debug(
                "sharding_prop python cache %s: %s -> %s",
                "HIT" if cache_hit else "MISS",
                op_schema,
                output_spec,
            )
            return result
        return self.cache(*args, **kwargs)

    def cache_info(self):
        return self.cache.cache_info()

    def cache_clear(self):
        return self.cache.cache_clear()


def _format_unbacked_hinting_log(
    op_schema: OpSchema,
    strategies: list[OpSpec],
    strategy_index: int,
    replacements: dict,
) -> str:
    """Format log message for unbacked hinting strategy selection (only called if debug logging enabled)."""
    args_spec = tuple(str(spec) for spec in op_schema.args_schema)
    strat = strategies[strategy_index]
    if strat.input_specs is None:
        placements_in = None
    else:
        placements_in = tuple(
            spec.format_shard_order_str(spec.placements, spec.shard_order)
            for spec in strat.input_specs
        )
    placements_out = tree_map(
        lambda spec: spec.format_shard_order_str(spec.placements, spec.shard_order),
        strat.output_specs,
        is_leaf=lambda x: isinstance(x, DTensorSpec),
    )
    return (
        f"Selected strategy {placements_in} -> {placements_out} "
        f"for {op_schema.op} with input {args_spec}, using unbacked hints: {replacements}"
    )


def _select_min_redistribute_cost(
    costs: list[torch.types.FloatLikeType],
    strategies: list[OpSpec],
    op_schema: OpSchema | None = None,
) -> int:
    """
    Given a list of costs and corresponding op strategies, selects the minimum cost strategy, returning the index.
    If unbacked symbols are involved, replaces them with known upper-bound values, falling back to hardcoded values.
    """
    from torch.fx.experimental.symbolic_shapes import (
        free_unbacked_symbols,
        is_concrete_float,
    )
    from torch.utils._sympy.interp import sympy_interp
    from torch.utils._sympy.numbers import int_oo
    from torch.utils._sympy.reference import PythonReferenceAnalysis

    int_fallback = 8192
    free_unbacked = list(set(chain(*[free_unbacked_symbols(cost) for cost in costs])))

    # Easy path: no unbacked shapes involved, choose min cost strategy.
    # Doing the hard path for backed could also make sense?
    if all(is_concrete_float(c) for c in costs) or not free_unbacked:
        return costs.index(min(costs))

    # Figure out heuristic hints for unbacked shapes.
    # If available, use shape upper bound. If not, fallback to some integer (inductor size-hinting style).
    shape_env = next(iter(x for x in costs if not is_concrete_float(x))).node.shape_env  # type: ignore[arg-type]
    replacements = {}
    for sym in free_unbacked:
        # TODO(laithsakka): unify with optimization_hint API
        if (hint := shape_env.var_to_hint_override.get(sym)) is not None:
            replacements[sym] = hint
        elif (upper := shape_env.bound_sympy(sym).upper) is not int_oo:
            replacements[sym] = upper
        else:
            replacements[sym] = int_fallback

    # Use replacements for redistribute cost hints
    proxy_costs = [
        float(cost)
        if is_concrete_float(cost)
        else sympy_interp(
            PythonReferenceAnalysis,
            replacements,
            cost.node.expr.xreplace(replacements),  # type: ignore[arg-type]
        )
        for cost in costs
    ]
    min_cost = min(proxy_costs)
    strategy_index = proxy_costs.index(min_cost)

    if op_schema:
        log.debug(
            "%s",
            LazyString(
                _format_unbacked_hinting_log,
                op_schema,
                strategies,
                strategy_index,
                replacements,
            ),
        )
    return strategy_index


def _select_min_cost_strategy(
    strategy: OpStrategy, op_schema: OpSchema | None = None
) -> OpSpec:
    from torch.fx.experimental.symbolic_shapes import guard_or_false

    if len(strategy.strategies) == 1:
        # short cut with only one possible OpSpec
        return strategy.strategies[0]

    op_spec_costs: list[torch.types.FloatLikeType] = []
    no_redistribute_strategy_index: int = -1
    negative_cost_index: int = -1
    zero_cost_index: int = -1
    for strategy_idx, op_spec in enumerate(strategy.strategies):
        if op_spec.redistribute_cost is None:
            raise AssertionError("must set redistribute cost each OpSpec!")
        redistribute_cost = sum(chain.from_iterable(op_spec.redistribute_cost))
        op_spec_costs.append(redistribute_cost)

        # If there are strategies with negative/zero/no redistribute cost,
        # we record those indices.
        # TODO: Currently this only applies to OpStrategy selection. Requires extra
        # logic to make it work for TupleStrategy, if needed.
        if op_schema is not None:
            if guard_or_false(redistribute_cost < 0):
                if (
                    negative_cost_index == -1
                    or redistribute_cost < op_spec_costs[negative_cost_index]
                ):
                    negative_cost_index = strategy_idx
            elif guard_or_false(redistribute_cost == 0):
                needs_redistribute = False
                for spec_idx, input_spec in enumerate(op_schema.args_spec):
                    desired_spec = (
                        op_spec.output_spec
                        if op_spec.input_specs is None
                        else op_spec.input_specs[spec_idx]
                    )
                    if input_spec.placements != desired_spec.placements:
                        needs_redistribute = True
                        break

                if not needs_redistribute:
                    no_redistribute_strategy_index = strategy_idx
                elif zero_cost_index == -1:
                    zero_cost_index = strategy_idx

    # prioritize negative/zero/no redistribute cost strategies
    if negative_cost_index != -1:
        # If there's negative cost, we select the one with the minimal cost,
        # even if this means we need to redistribute, e.g. via local chunking.
        # E.g. this can happen for ops in self.op_to_shape_and_stride_idx
        # when the inputs / outputs are sharded.
        selected_strategy_index = negative_cost_index
    elif no_redistribute_strategy_index != -1:
        selected_strategy_index = no_redistribute_strategy_index
    elif zero_cost_index != -1:
        selected_strategy_index = zero_cost_index
    else:
        # default to choosing minimal redistribute cost
        selected_strategy_index = _select_min_redistribute_cost(
            op_spec_costs, strategy.strategies, op_schema
        )

    return strategy.strategies[selected_strategy_index]


class ShardingPropagator:
    # Lock to protect FakeTensorMode context during tensor meta propagation.
    # By default this is a no-op (nullcontext) for performance. Multi-threaded
    # tests should set this to threading.Lock() to prevent race conditions
    # when multiple threads enter different FakeTensorMode contexts.
    _fake_mode_lock = nullcontext()

    def __init__(self) -> None:
        self.op_to_rules: dict[OpOverload, Callable[[OpSchema], OutputSharding]] = {}
        self.op_strategy_funcs: dict[
            OpOverload,
            Callable[[OpSchema], StrategyType],
        ] = {}
        self.op_single_dim_strategy_funcs: dict[
            OpOverload,
            _SingleDimStrategyInfo,
        ] = {}
        # op map to save static argnum to decide to reuse sharding prop cache or
        # re-run sharding prop
        self.op_to_schema_info: dict[OpOverload, RuntimeSchemaInfo] = {}
        self.op_to_schema_info_for_single_dim_strategy: dict[
            OpOverload, RuntimeSchemaInfo
        ] = {}
        self.propagate_op_sharding = LocalLRUCache(
            self.propagate_op_sharding_non_cached
        )
        self.decomp_strategy = DecompShardingStrategy(self)
        # op map to save indices of shape (and stride) args which may need to be
        # modified in sharding prop
        self.op_to_shape_and_stride_idx: dict[OpOverload, int | tuple[int, int]] = {
            # new factory ops
            aten.new_empty.default: 1,
            aten.new_full.default: 1,
            aten.new_ones.default: 1,
            aten.new_zeros.default: 1,
            aten.new_empty_strided.default: (1, 2),
            # view ops
            aten.expand.default: 1,
            aten.expand_copy.default: 1,
            aten.reshape.default: 1,
            aten.view.default: 1,
            aten.view_copy.default: 1,
            aten._unsafe_view.default: 1,
            aten.select_backward.default: 1,
            aten.slice_backward.default: 1,
            # upsample backward ops: input_size (arg 2) is the full [N,C,...] shape
            aten.upsample_nearest1d_backward.default: 2,
            aten.upsample_nearest2d_backward.default: 2,
            aten.upsample_nearest3d_backward.default: 2,
            aten._upsample_nearest_exact1d_backward.default: 2,
            aten._upsample_nearest_exact2d_backward.default: 2,
            aten._upsample_nearest_exact3d_backward.default: 2,
            aten._upsample_bilinear2d_aa_backward.default: 2,
            aten._upsample_bicubic2d_aa_backward.default: 2,
            aten._upsample_lanczos2d_aa_backward.default: 2,
            aten.upsample_bicubic2d_backward.default: 2,
            aten.upsample_bilinear2d_backward.default: 2,
            aten.upsample_linear1d_backward.default: 2,
            aten.upsample_trilinear3d_backward.default: 2,
        }
        # ops with individual scalar shape args that need local adjustment
        # maps op -> callable(input_specs, schema) -> adjusted schema
        # populated by op modules (e.g. _math_ops.py) at registration time
        self.op_to_scalar_shape_adjuster: dict[
            OpOverload,
            Callable[[list[DTensorSpec], OpSchema], OpSchema],
        ] = {}
        # squeeze ops that need dim arg rewritten to only globally-singleton dims
        self.squeeze_op_to_dims_variant: dict[OpOverload, OpOverload] = {
            aten.squeeze.default: aten.squeeze.dims,
            aten.squeeze.dim: aten.squeeze.dims,
            aten.squeeze.dims: aten.squeeze.dims,
            aten.squeeze_.default: aten.squeeze_.dims,
            aten.squeeze_.dim: aten.squeeze_.dims,
            aten.squeeze_.dims: aten.squeeze_.dims,
        }

    def register_sharding_prop_rule(
        self,
        op_overload: OpOverload,
        rule_func: Callable[[OpSchema], OutputSharding],
        schema_info: RuntimeSchemaInfo | None = None,
    ):
        """
        Register a sharding propagation rule for an operator.
        """
        self.op_to_rules[op_overload] = rule_func
        if schema_info is not None:
            self.op_to_schema_info[op_overload] = schema_info

    def register_single_dim_op_strategy(
        self,
        op_overload: OpOverload,
        strategy_info: _SingleDimStrategyInfo,
        schema_info: RuntimeSchemaInfo | None = None,
    ):
        """
        Register a strategy over a single mesh-dim, relying on infra to automatically expand to the full mesh.
        """
        self.op_single_dim_strategy_funcs[op_overload] = strategy_info
        if schema_info is not None:
            self.op_to_schema_info_for_single_dim_strategy[op_overload] = schema_info

    def register_op_strategy(
        self,
        op_overload: OpOverload,
        strategy_func: Callable[[OpSchema], StrategyType],
        schema_info: RuntimeSchemaInfo | None = None,
    ):
        """
        Register a :class:`OpStrategy` generator for an operator.

        During the sharding propagation, DTensor wants to enumerate all
        acceptable sharding specs (:class:`OpSpec`) for an operator,
        and by "acceptable" we mean that the operator can be executed on
        the ``_local_tensor`` of DTensor args/kwargs (with ``OpSpec.input_specs``)
        and the output(s) constitute valid DTensor(s) (with ``OpSpec.output_specs``).

        ``strategy_func`` is the function that enumerates such acceptable specs
        for the operator ``op_overload``. One general approach to write ``strategy_func``
        is, if the operator has simple arguments structure (e.g. mm, bmm), first enumerating
        all sharding specs for the operands, and then filtering out the ones that
        are not valid. For example, for ``mm``, the operands are two 2D tensors, and
        if both ``input`` and ``mat2`` have sharding placements ``[Shard(0)]``, then this
        is not an acceptable ``input_specs``.

        Once we have a way to enumerate all acceptable sharding specs, we can use each
        of them to construct a :class:`OpSpec`. The ``OpSpec.input_specs`` directly comes
        from the sharding spec, and the ``OpSpec.output_specs`` is therefore determined
        (e.g. ``[Shard(1)]`` @ ``[Shard(0)]`` yields ``[Partial()]``). In addition,
        :class:`OpSpec` also contains ``redistribute_cost`` which records the redistribution
        cost from each :class:`OpSpec` in the source :class:`OpStrategy.strategies` to
        the target sharding spec, for each operand.

        The ``strategy_func`` should return a :class:`OpStrategy` which contains a list of
        all the :class:`OpSpec`s generated in the above.

        The optional ``schema_info`` tells which non-DTensor args/kwargs could affect the
        cache and whether ``pytree`` is needed to flatten the nested args. ``static_argnum``
        marks the starting index of the non-DTensor args that should be hashed into the
        sharding propagation hash key, and ``static_kwargkey`` marks the keys of the
        non-DTensor kwargs that should be hashed. ``needs_pytree`` should be used when
        the input arg has :class:`list` or :class:`dict` structure.

        For example, ``aten.cat.default`` op has a ``List[Tensor]`` argument ``tensors``
        and an ``int`` argument ``dim``. Because ``dim`` affects the sharding propagation
        result, we want to pass ``RuntimeSchemaInfo(static_argnum=1)`` because the argument
        index of ``dim`` is 1. Besides, we also want to set ``needs_pytree=True`` because
        ``tensors`` needs be flattened in sharding propagation. Another example is
        ``aten.histc.default``. ``histc`` has 4 arguments (self, bins, min, max) and the
        last two would affect sharding propagation along with the :class:`DTensor` argument
        ``self``. Since the argument index of ``min`` is 2, the `schema_info` should be
        `RuntimeSchemaInfo(static_argnum=2)`.
        """
        self.op_strategy_funcs[op_overload] = strategy_func
        if schema_info is not None:
            self.op_to_schema_info[op_overload] = schema_info

    def _propagate_tensor_meta_non_cached(
        self, op_schema: OpSchema
    ) -> TensorMeta | Sequence[TensorMeta | None] | None:
        """
        Propagate the tensor metadata, it could either return a TensorMeta
        or a list/tuple of TensorMetas
        """
        if op_schema.op == aten.equal.default:
            # data dependent ops can't be used for fake propagation
            return None

        # NOTE: We must call the tracing in fake tensor mode so that it avoids
        # materializing memory.
        # NOTE: Use _fake_mode_lock to serialize access when running in
        # multi-threaded tests (lock must be set to threading.Lock()).
        # This is a nullcontext by default.
        # NOTE: disable_proxy_modes_tracing() prevents make_fx from tracing
        # into shard propagation. This op runs purely to derive output
        # metadata (global shape/stride/dtype); it is not part of the real
        # (local) computation. Without disabling the proxy tracer, the
        # gen_fake_args() factory calls and op_schema.op(...) get recorded
        # into the enclosing graph as dead code -- global-shape "shadow" nodes
        # on empty_strided placeholders that no real output consumes.
        from torch.fx.experimental.proxy_tensor import disable_proxy_modes_tracing

        with ShardingPropagator._fake_mode_lock:
            fake_mode = detect_fake_mode() or FakeTensorMode()
            with fake_mode, disable_proxy_modes_tracing():
                fake_args = op_schema.gen_fake_args()
                fake_kwargs = op_schema.gen_fake_kwargs()
                fake_out = op_schema.op(*fake_args, **fake_kwargs)

        if isinstance(fake_out, torch.Tensor):
            return TensorMeta(
                shape=fake_out.shape, stride=fake_out.stride(), dtype=fake_out.dtype
            )

        elif isinstance(fake_out, (tuple, list)):
            tensor_meta_list: list[TensorMeta | None] = []
            for fake_out_item in fake_out:
                if isinstance(fake_out_item, torch.Tensor):
                    tensor_meta_list.append(
                        TensorMeta(
                            shape=fake_out_item.shape,
                            stride=fake_out_item.stride(),
                            dtype=fake_out_item.dtype,
                        )
                    )
                else:
                    tensor_meta_list.append(None)
            return (
                tuple(tensor_meta_list)
                if isinstance(fake_out, tuple)
                else tensor_meta_list
            )
        else:
            # if fake is not a tensor or tuple of tensor, return as none
            return None

    @lru_cache  # noqa: B019
    def _propagate_tensor_meta_cached(
        self, op_schema: OpSchema
    ) -> TensorMeta | Sequence[TensorMeta | None] | None:
        """
        Cached version of _propagate_tensor_meta_non_cached
        Use _propagate_tensor_meta instead to handle dynamic shapes.
        """
        return self._propagate_tensor_meta_non_cached(op_schema)

    def _propagate_tensor_meta(
        self, op_schema: OpSchema
    ) -> TensorMeta | Sequence[TensorMeta | None] | None:
        """
        Propagate the tensor metadata, it could either return a TensorMeta
        or a list/tuple of TensorMetas. Uses the cached version if not
        actively tracing. Use this method instead of _propagate_tensor_meta_non_cached
        """
        if _are_we_tracing():
            return self._propagate_tensor_meta_non_cached(op_schema)
        else:
            return self._propagate_tensor_meta_cached(op_schema)

    def _create_output_spec_with_new_tensor_meta(
        self,
        op: OpOverload,
        output_specs: OutputSpecType,
        output_tensor_meta: TensorMeta | Sequence[TensorMeta | None] | None,
    ) -> OutputSpecType:
        """
        Wrap the output_specs with the tensor metadata from the output.
        """

        if isinstance(output_specs, DTensorSpec):
            if not isinstance(output_tensor_meta, TensorMeta):
                # Either error due to ShardingPropagator or due to incorrect OutputSpec
                if not isinstance(output_tensor_meta, (tuple, list)):
                    raise ValueError(
                        "ShardingPropagator error: output does not have an associated "
                        "TensorMeta"
                    )
                raise ValueError(
                    f"For the op {op.name()}, `output_specs` has 1 output which does "
                    "not equal the "
                    f"number of op outputs: {len(output_tensor_meta)}."
                )
            return output_specs.shallow_copy_with_tensor_meta(output_tensor_meta)
        elif isinstance(output_specs, (tuple, list)):
            new_specs: list[DTensorSpec | None] = []
            if not isinstance(output_tensor_meta, (tuple, list)) or len(
                output_specs
            ) != len(output_tensor_meta):
                raise ValueError(
                    f"For the op {op.name()}, `output_specs` has {len(output_specs)} "
                    "outputs which does not equal the "
                    f"number of op outputs {_length(output_tensor_meta)}."
                )

            # pyrefly: ignore [bad-argument-type]
            for i, spec in enumerate(output_specs):
                if isinstance(spec, DTensorSpec):
                    output_tensor_meta_i = output_tensor_meta[i]
                    if not isinstance(output_tensor_meta_i, TensorMeta):
                        # Some ops (e.g. convolution_backward, native_layer_norm_backward,
                        # _fused_rms_norm_backward) have an output_mask parameter that
                        # controls which outputs are computed. When output_mask[i] is
                        # False, the output at position i is None and has no TensorMeta.
                        if output_tensor_meta_i is None:
                            new_specs.append(None)
                            continue
                        else:
                            raise ValueError(
                                f"ShardingPropagator error: output {i} of {op.name()} "
                                "does not have an associated TensorMeta"
                            )

                    new_specs.append(
                        spec.shallow_copy_with_tensor_meta(output_tensor_meta_i)
                    )
                else:
                    new_specs.append(spec)

            return tuple(new_specs)
        else:
            if output_specs is not None:
                raise AssertionError
            return output_specs

    def _wrap_with_op_strategy(self, op_schema: OpSchema) -> OpSchema:
        """
        wrap a op_schema that contains DTensorSpec to another op_schema that contains
        OpStrategy/TupleStrategy, the returned op_schema is then used for sharding
        strategy propagation on pytorch operators.
        """

        def spec_to_strategy(spec: object) -> object:
            if isinstance(spec, DTensorSpec):
                return OpStrategy([OpSpec(spec)])
            elif isinstance(spec, (list, tuple)) and len(spec) > 0:
                if all(isinstance(s, DTensorSpec) for s in spec):
                    # tensor list create tuple strategy
                    tuple_strategy = [spec_to_strategy(s) for s in spec]
                    tuple_strategy = cast(Sequence[StrategyType], tuple_strategy)
                    return TupleStrategy(
                        tuple(tuple_strategy)
                        if isinstance(spec, tuple)
                        else tuple_strategy
                    )
                elif any(isinstance(s, DTensorSpec) for s in spec):
                    # mixed list (e.g. [DTensorSpec, None, DTensorSpec]) for
                    # ops like aten.index.Tensor; keep as list so pytree
                    # flattening can extract OpStrategy items
                    return [spec_to_strategy(s) for s in spec]
                else:
                    return spec
            else:
                return spec

        args_op_strategy = [spec_to_strategy(i) for i in op_schema.args_schema]

        kwargs_op_strategy = {
            k: spec_to_strategy(v) for k, v in op_schema.kwargs_schema.items()
        }

        return OpSchema(
            op=op_schema.op,
            args_schema=tuple(args_op_strategy),
            kwargs_schema=kwargs_op_strategy,
            schema_info=op_schema.schema_info,
        )

    def propagate(self, op_info: OpInfo) -> None:
        # NB: The logic here is duplicated in _propagate_op_sharding_dispatch_slow_path.
        # Ideally, this function would be deleted, but there are a handful of
        # one off call sites here that aren't cleaned up.

        # NOTE: schema should always be populated when calling this function,
        # as it's only called after unwrap_to_op_info (create_schema=True).
        if op_info.schema is None:
            raise AssertionError(
                "op_info.schema should not be None in propagate. "
                "This function should only be called after unwrap_to_op_info."
            )

        # We cannot use an lru cache if we know that inputs will have dynamic shapes,
        # because SymInts are not hashable.
        # This is generally ok because this only happens during tracing in torch.compile,
        # and tracing does not need to be as fast as eagermode DTensor usages.
        if _are_we_tracing():
            output_sharding = self.propagate_op_sharding_non_cached(op_info.schema)
        else:
            output_sharding = cast(
                OutputSharding, self.propagate_op_sharding(op_info.schema)
            )
        op_info.output_sharding = output_sharding

    def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputSharding:
        """
        Propagate the sharding for an operator given the op_schema.
        """
        # no-op in OSS, logs API usage metrics in meta-internal runs
        torch._C._log_api_usage_once(
            "torch.distributed.tensor._sharding_prop.ShardingPropagator.propogate_op_sharding_non_cached"
        )
        # special case op, we don't need to propagate for local
        # scalar. TODO: figure out a better way to handle this
        if op_schema.op is aten._local_scalar_dense.default:
            return OutputSharding(None, op_schema)

        out_tensor_meta = self._propagate_tensor_meta_non_cached(op_schema)

        single_dim_strategy_info = self.op_single_dim_strategy_funcs.get(op_schema.op)
        op_strategy_func = self.op_strategy_funcs.get(op_schema.op)
        decomp_exception = None
        if single_dim_strategy_info is not None or op_strategy_func is not None:
            # Validate that tensor_meta count matches expected outputs from op schema.
            # This catches bugs in fake tensor propagation early.
            if single_dim_strategy_info is not None:
                _validate_tensor_meta_count(op_schema, out_tensor_meta)
            """
            Given the single_dim_strategy, which is just a minimal set of valid input-output placement specifications
            for the operator over a single mesh dimension,

            And the OpSchema, which includes information about the runtime input tensor placements, and the mesh,

            Combine single_dim_strategies across mesh dims, also expanding placeholders (ShardPlaceholder) to any real
            sharding types in op_schema, and find the lowest cost redistribution of inputs to match a valid strategy
            combination.
            """
            # wrap the op_schema with op strategy for sharding strategy propagation
            strategy_schema = self._wrap_with_op_strategy(op_schema)

            if single_dim_strategy_info is not None:
                mesh = try_find_mesh_from_args(op_schema.op, op_schema.args_schema)
                if not isinstance(mesh, DeviceMesh):
                    raise AssertionError("Expected to find a valid mesh")
                # expand to generate the full set of strategy combinations, each one
                # with a redistribute cost, and then find the min strategy over those costs.
                _expanded_strategy_fn = _expand_single_dim_strategy_to_mesh(
                    mesh, strategy_schema, single_dim_strategy_info, out_tensor_meta
                )
                op_strategy = _expanded_strategy_fn(
                    op_schema.op, strategy_schema.args_meta, strategy_schema.kwargs_meta
                )
            else:
                if op_strategy_func is None:
                    raise AssertionError
                op_strategy = op_strategy_func(strategy_schema)

        else:
            # try operator decomposition path

            op_strategy = None
            if DecompShardingStrategy.has_decomp(op_schema.op):
                # Ensure schema_info is registered for proper cache key computation
                self.decomp_strategy.ensure_schema_info(op_schema.op)
                try:
                    op_strategy = self.decomp_strategy.propagate_strategy(
                        op_schema,
                    )
                except Exception as e:
                    decomp_exception = e

        if op_strategy is not None:
            if isinstance(op_strategy, OpStrategy):
                _propagate_use_strided_shard_flag(op_strategy, op_schema)
                # single Op strategy
                output_strategy = _select_min_cost_strategy(op_strategy, op_schema)

                # check if we need to redistribute the input
                needs_redistribute = False
                # check if we want to use args value from redistribute_schema
                use_val_from_redistribute_schema = False
                expected_input_specs: list[DTensorSpec] = []

                # in case where the op does not specify input_specs and output_specs
                # is a DTensorSpec, we use output_specs as the spec for each DTensor
                # input arg.
                if output_strategy.input_specs is None:
                    if not isinstance(output_strategy.output_specs, DTensorSpec):
                        raise AssertionError

                for idx, input_spec in enumerate(op_schema.args_spec):
                    desired_spec = (
                        output_strategy.output_spec
                        if output_strategy.input_specs is None
                        else output_strategy.input_specs[idx]
                    )
                    expected_input_specs.append(
                        desired_spec.shallow_copy_with_tensor_meta(
                            input_spec.tensor_meta
                        )
                    )
                    if input_spec.placements != desired_spec.placements:
                        needs_redistribute = True

                suggestion_schema = None
                if needs_redistribute:
                    suggestion_schema = OpSchema(
                        op_schema.op, tuple(expected_input_specs), {}
                    )
                    suggestion_schema._inplace_rewrap_schema_suggestion(op_schema)

                # Rewrite shape/stride args from global to local when sharded.
                # Covers view ops, new factory ops, and upsample backward ops
                # whose shape args refer to global tensor dimensions.
                if op_schema.op in self.op_to_shape_and_stride_idx:
                    if not isinstance(output_strategy.output_spec, DTensorSpec):
                        raise AssertionError
                    if any(
                        isinstance(p, Shard | _StridedShard)
                        for p in output_strategy.output_spec.placements
                    ):
                        schema = suggestion_schema or op_schema
                        if not isinstance(out_tensor_meta, TensorMeta):
                            raise AssertionError
                        suggestion_schema = self._adjust_shape_and_stride_args(
                            out_tensor_meta, schema, output_strategy.output_spec
                        )
                        needs_redistribute = True
                        use_val_from_redistribute_schema = True

                # adjust individual scalar shape args (e.g. N, C, HxW in group_norm)
                if op_schema.op in self.op_to_scalar_shape_adjuster:
                    if any(
                        isinstance(p, Shard | _StridedShard)
                        for spec in expected_input_specs
                        for p in spec.placements
                    ):
                        schema = suggestion_schema or op_schema
                        adjuster = self.op_to_scalar_shape_adjuster[op_schema.op]
                        suggestion_schema = adjuster(expected_input_specs, schema)
                        needs_redistribute = True
                        use_val_from_redistribute_schema = True

                # rewrite squeeze to use only globally-singleton dims
                if op_schema.op in self.squeeze_op_to_dims_variant:
                    schema = suggestion_schema or op_schema
                    adjusted = self._adjust_squeeze_to_global_singletons(schema)
                    if adjusted is not None:
                        suggestion_schema = adjusted
                        needs_redistribute = True
                        use_val_from_redistribute_schema = True

                # construct output spec for the op
                if op_schema.return_type_tuple_tensor_like():
                    # for ops that return multiple tensors and the output_specs is not
                    # a tuple, we use a tuple of that single output spec as the new
                    # output_specs
                    output_specs: OutputSpecType = output_strategy.output_specs
                    if isinstance(output_specs, DTensorSpec):
                        output_specs = tuple(
                            # create a new DTensorSpec with the same placement as the
                            # output_specs in output_strategy
                            DTensorSpec(
                                mesh=output_specs.mesh,
                                placements=output_specs.placements,
                                tensor_meta=output_specs.tensor_meta,
                                use_strided_shard_as_shard_order=output_specs.use_strided_shard_as_shard_order,
                            )
                            for _ in range(len(op_schema.op._schema.returns))
                        )
                elif (
                    op_schema.return_type_tensor()
                    or op_schema.return_type_list_tensor_like()
                ):
                    output_specs = output_strategy.output_specs
                else:
                    output_specs = None

                output_sharding = OutputSharding(
                    output_specs,
                    suggestion_schema,
                    needs_redistribute=needs_redistribute,
                    use_val_from_redistribute_schema=use_val_from_redistribute_schema,
                )
            elif isinstance(op_strategy, TupleStrategy):
                # tuple strategy output sharding processing
                # runtime select OpSpec for each TupleStrategy input arg
                selected_strategies: list[OpSpec] = []
                out_spec_list: list[DTensorSpec] = []
                for strategy in op_strategy.children:
                    if not isinstance(strategy, OpStrategy):
                        raise AssertionError
                    _propagate_use_strided_shard_flag(strategy, op_schema)
                    selected_strategy = _select_min_cost_strategy(strategy)
                    selected_strategies.append(selected_strategy)
                    if selected_strategy.output_specs is not None:
                        out_spec_list.append(selected_strategy.output_spec)

                needs_redistribute = False
                suggestion_args: list[object] = []
                tensor_or_list_tensor_arg_idx = 0

                for arg in op_schema.args_schema:
                    if (
                        arg
                        and isinstance(arg, (list, tuple))
                        and isinstance(arg[0], DTensorSpec)
                    ):
                        expected_input_spec_list: list[DTensorSpec] = []
                        for idx, arg_spec in enumerate(arg):
                            expected_input_spec = selected_strategies[idx].input_spec(
                                tensor_or_list_tensor_arg_idx
                            )
                            expected_input_spec = (
                                expected_input_spec.shallow_copy_with_tensor_meta(
                                    arg_spec.tensor_meta
                                )
                            )
                            if arg_spec.placements != expected_input_spec.placements:
                                needs_redistribute = True
                            expected_input_spec_list.append(expected_input_spec)
                        suggestion_args.append(
                            tuple(expected_input_spec_list)
                            if isinstance(arg, tuple)
                            else expected_input_spec_list
                        )
                        tensor_or_list_tensor_arg_idx += 1

                    elif isinstance(arg, DTensorSpec):
                        expected_input_spec = selected_strategies[0].input_spec(
                            tensor_or_list_tensor_arg_idx
                        )
                        expected_input_spec = (
                            expected_input_spec.shallow_copy_with_tensor_meta(
                                arg.tensor_meta
                            )
                        )
                        if arg.placements != expected_input_spec.placements:
                            needs_redistribute = True
                        suggestion_args.append(expected_input_spec)
                        tensor_or_list_tensor_arg_idx += 1
                    else:
                        suggestion_args.append(arg)

                suggestion_schema = None
                if needs_redistribute:
                    suggestion_schema = OpSchema(
                        op_schema.op, tuple(suggestion_args), op_schema.kwargs_schema
                    )

                output_sharding = OutputSharding(
                    tuple(out_spec_list) if out_tensor_meta is not None else None,
                    suggestion_schema,
                    needs_redistribute=needs_redistribute,
                    use_val_from_redistribute_schema=False,
                )
            else:
                raise ValueError("Unsupported op strategy type")

            # associate the output sharding with the output tensor metadata
            new_output_spec = self._create_output_spec_with_new_tensor_meta(
                op_schema.op, output_sharding.output_spec, out_tensor_meta
            )
            output_sharding.output_spec = new_output_spec
            return output_sharding
        elif op_schema.op in self.op_to_rules:
            # propagate the sharding with rule
            sharding_prop_func = self.op_to_rules[op_schema.op]

            # step 1. there's sharding propagation rule, run
            # sharding propagation to get the output sharding
            try:
                output_sharding = sharding_prop_func(op_schema)
            except NotImplementedError as e:
                raise e
            except Exception as e:
                raise RuntimeError(
                    f"Sharding propagation failed on op {op_schema}.\nError: {e}"
                ) from e

            # step 2. if can't get output_spec from sharding
            # propagation (i.e. no rules apply for input
            # placements), we return the output sharding
            # with schema suggestions, which can be used to
            # decide how to do redistribute on inputs
            if output_sharding.output_spec is None:
                if output_sharding.redistribute_schema is None:
                    raise RuntimeError(
                        f"Sharding propagation failed on op {op_schema}!"
                    )
                else:
                    # we do auto redistribute on inputs if necessary
                    # run sharding propagation again with suggested schema
                    propagation_res = sharding_prop_func(
                        output_sharding.redistribute_schema
                    )
                    # we set the output sharding with the new propagation result
                    # so that dispatching know both output_spec and redistribute_schema
                    # exist, which indicates a reshard is needed
                    output_sharding.output_spec = propagation_res.output_spec
                    output_sharding.needs_redistribute = True

            # associate the output sharding with the output tensor metadata
            new_output_spec = self._create_output_spec_with_new_tensor_meta(
                op_schema.op, output_sharding.output_spec, out_tensor_meta
            )
            output_sharding.output_spec = new_output_spec

            return output_sharding
        else:
            raise NotImplementedError(
                f"Operator {op_schema.op} does not have a sharding strategy registered."
            ) from decomp_exception

    def _adjust_shape_and_stride_args(
        self,
        out_tensor_meta: TensorMeta,
        schema: OpSchema,
        spec: DTensorSpec,
    ) -> OpSchema:
        shape_stride_idx = self.op_to_shape_and_stride_idx[schema.op]
        if isinstance(shape_stride_idx, tuple):
            shape_idx, stride_idx = shape_stride_idx
        else:
            shape_idx = shape_stride_idx
            stride_idx = None

        expected_input_schema = list(schema.args_schema)
        # rewrite args[shape_idx] from the global output shape to the local output shape
        local_shape, _ = compute_local_shape_and_global_offset(
            out_tensor_meta.shape, spec.mesh, spec.placements, skip_offset=True
        )
        expected_input_schema[shape_idx] = local_shape

        # adjust the stride arg for aten.new_empty_strided.default
        if stride_idx:
            expected_input_schema[stride_idx] = compute_local_stride(
                out_tensor_meta.stride, local_shape
            )

        return OpSchema(schema.op, tuple(expected_input_schema), schema.kwargs_schema)

    def _adjust_squeeze_to_global_singletons(self, schema: OpSchema) -> OpSchema | None:
        """
        Rewrite squeeze ops to squeeze.dims with only globally-singleton dims.
        Fixes bug where sharded dims with local size 1 get incorrectly squeezed.
        Returns None if no rewrite is needed (already squeeze.dims with correct args).
        """
        from torch.fx.experimental.symbolic_shapes import guard_or_false

        input_spec = cast(DTensorSpec, schema.args_schema[0])
        tensor_meta = input_spec.tensor_meta
        if tensor_meta is None:
            raise RuntimeError("squeeze requires tensor metadata")
        global_shape = tensor_meta.shape
        ndim = len(global_shape)

        def normalize(d: int) -> int:
            return d if d >= 0 else d + ndim

        def is_singleton(d: int) -> bool:
            nd = normalize(d)
            return 0 <= nd < ndim and guard_or_false(global_shape[nd] == 1)

        # guard_or_false: conservatively keep dims when size is symbolic/unknown
        if schema.op in (aten.squeeze.default, aten.squeeze_.default):
            target_dims = tuple(
                i for i, s in enumerate(global_shape) if guard_or_false(s == 1)
            )
        elif schema.op in (aten.squeeze.dim, aten.squeeze_.dim):
            dim = normalize(schema.args_schema[1])  # type: ignore[arg-type]
            target_dims = (dim,) if is_singleton(dim) else ()
        else:
            dims = cast(Sequence[int], schema.args_schema[1])
            target_dims = tuple(  # type: ignore[union-attr]
                normalize(d) for d in dims if is_singleton(d)
            )

        dims_variant = self.squeeze_op_to_dims_variant[schema.op]
        # Skip rewrite if already targeting the right op with the same dims
        if schema.op == dims_variant and len(schema.args_schema) > 1:
            existing_dims = schema.args_schema[1]
            if existing_dims == target_dims:
                return None
        return OpSchema(dims_variant, (input_spec, target_dims), {})
