# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""Minimal repro for torch.compile scheduling of async D2H split-size reads.

Run on a single CUDA device:

    TORCH_LOGS=output_code python scripts/repro_compile_async_d2h_splits.py

By default the compiled function uses the problematic pattern:

    split_sizes_gpu.to("cpu", non_blocking=True).tolist()

Inspect the generated code for ``copy_(..., True)`` immediately followed by
``.item()`` reads. Pass ``--blocking`` to compare with ``non_blocking=False``.
"""

from __future__ import annotations

import argparse

import torch
import torch.distributed as dist
from torch.distributed._functional_collectives import (
    all_to_all_single,
    all_to_all_single_autograd,
)
from torch.testing._internal.distributed.fake_pg import FakeStore

FAKE_WORLD_SIZE = 4
NUM_LOCAL_EXPERTS = 4
HIDDEN_SIZE = 8


def repro_step(
    routed_input: torch.Tensor,
    num_tokens_per_expert: torch.Tensor,
    ep_degree: int,
    non_blocking_input_copy: bool,
) -> torch.Tensor:
    num_tokens_per_expert_group = all_to_all_single(
        num_tokens_per_expert,
        None,
        None,
        group=dist.group.WORLD,
    )
    num_tokens_per_expert_group = torch.ops._c10d_functional.wait_tensor(
        num_tokens_per_expert_group
    )

    input_splits = (
        num_tokens_per_expert.view(ep_degree, -1)
        .sum(dim=1)
        .to(torch.device("cpu"), non_blocking=non_blocking_input_copy)
    )
    output_splits = (
        num_tokens_per_expert_group.view(ep_degree, -1)
        .sum(dim=1)
        .to(torch.device("cpu"), non_blocking=False)
    )

    input_splits_list = input_splits.tolist()
    output_splits_list = output_splits.tolist()

    return all_to_all_single_autograd(
        routed_input,
        output_splits_list,
        input_splits_list,
        group=dist.group.WORLD,
    )


parser = argparse.ArgumentParser()
parser.add_argument(
    "--blocking",
    action="store_true",
    help="Use non_blocking=False for the input_splits D2H copy for comparison.",
)
args = parser.parse_args()

if not torch.cuda.is_available():
    raise RuntimeError("This repro requires CUDA.")

torch.cuda.set_device(0)
dist.init_process_group(
    "fake",
    rank=0,
    world_size=FAKE_WORLD_SIZE,
    store=FakeStore(),
)
rank = dist.get_rank()
world_size = dist.get_world_size()

# Capture .tolist()/.item() as unbacked SymInts instead of graph-breaking.
torch._dynamo.config.capture_scalar_outputs = True

# Layout matches AllToAllTokenDispatcher: rows are grouped by destination EP
# rank, and each group contains this rank's token counts for that rank's
# local experts.
counts = [
    ((rank + 1) * (dst_rank + 2) + expert_idx) % 5 + 1
    for dst_rank in range(world_size)
    for expert_idx in range(NUM_LOCAL_EXPERTS)
]
num_tokens_per_expert = torch.tensor(counts, dtype=torch.int64, device="cuda")
routed_input = torch.randn(
    sum(counts),
    HIDDEN_SIZE,
    device="cuda",
    dtype=torch.float32,
)

compiled_step = torch.compile(
    repro_step,
    backend="inductor",
    fullgraph=True,
    dynamic=True,
)
out = compiled_step(
    routed_input,
    num_tokens_per_expert,
    world_size,
    not args.blocking,
)
torch.cuda.synchronize()

dist.destroy_process_group()

---

triton_red_fused_sum_view_1.run(arg1_1, buf19, ...)
  buf20.copy_(buf19, True)

  u0 = reinterpret_tensor(buf20, (), (), 0).item()
  u1 = reinterpret_tensor(buf20, (), (), 1).item()
  u2 = reinterpret_tensor(buf20, (), (), 2).item()
  u3 = reinterpret_tensor(buf20, (), (), 3).item()

  buf25 = torch.ops._c10d_functional.all_to_all_single.default(
      arg5_1,
      [u4, u5, u6, u7],  # output_splits
      [u0, u1, u2, u3],  # input_splits
      "0",
  )

---

triton_red_fused_sum_view_1.run(arg1_1, buf19, ...)
  buf20.copy_(buf19, False)

  u0 = reinterpret_tensor(buf20, (), (), 0).item()
  u1 = reinterpret_tensor(buf20, (), (), 1).item()
  u2 = reinterpret_tensor(buf20, (), (), 2).item()
  u3 = reinterpret_tensor(buf20, (), (), 3).item()

  buf25 = torch.ops._c10d_functional.all_to_all_single.default(
      arg5_1,
      [u4, u5, u6, u7],
      [u0, u1, u2, u3],
      "0",
  )