import os
os.environ.setdefault("HCCL_OP_EXPANSION_MODE", "AIV")
import time

import torch
import torch_npu
import torch.distributed as dist

BS, HIDDEN = 128, 2048
N = 40 # 每图迭代数
K = 10 # 间距（原单测 90；调小 => 切流/通信更频繁）
SIDE = 2 # 侧流工作 mul 数
REPLAYS = 100 # 每轮计时 replay 次数
ROUNDS = 3 # off/on 交替轮数

def mul1(x):
x.mul_(1.0000001)

def fj_block(x, alt):
cur = torch.npu.current_stream()
for _ in range(N):
for _ in range(K):
mul1(x)
alt.wait_stream(cur)
with torch.npu.stream(alt):
for _ in range(SIDE):
mul1(x)
cur.wait_stream(alt)
mul1(x)

def ar_space(x):
for _ in range(N):
for _ in range(K):
mul1(x)
dist.all_reduce(x)

def capture(fn):
s = torch.npu.Stream()
s.wait_stream(torch.npu.current_stream())
with torch.npu.stream(s):
for _ in range(3):
fn()
torch.npu.current_stream().wait_stream(s)
torch.npu.synchronize()
dist.barrier()
g = torch.npu.NPUGraph()
try:
ctx = torch.npu.graph(g, stream=s, auto_dispatch_capture=True)
except TypeError:
ctx = torch.npu.graph(g, stream=s)
with ctx:
fn()
dist.barrier()
return g

def time_replay(g, iters):
torch.npu.synchronize()
t0 = time.perf_counter()
for _ in range(iters):
g.replay()
torch.npu.synchronize()
return (time.perf_counter() - t0) / iters

def time_replay_with_profiler(g, iters):
from torch_npu import profiler as npu_profiler
acts = [npu_profiler.ProfilerActivity.CPU, npu_profiler.ProfilerActivity.NPU]
with npu_profiler.profile(activities=acts):
return time_replay(g, iters)

def compare(g, name, rank):
for _ in range(20):
g.replay()
torch.npu.synchronize()
if rank == 0:
print(f"# {name}: N={N} K={K} SIDE={SIDE} replays={REPLAYS} rounds={ROUNDS}", flush=True)
print("case,round,off_us_per_iter,on_us_per_iter,delta_us,on/off", flush=True)
offs, ons = [], []
for r in range(ROUNDS):
off = time_replay(g, REPLAYS)
on = time_replay_with_profiler(g, REPLAYS)
offs.append(off)
ons.append(on)
if rank == 0:
print(f"{name},{r},{off / N * 1e6:.2f},{on / N * 1e6:.2f},"
f"{(on - off) / N * 1e6:+.2f},{on / off:.4f}", flush=True)
if rank == 0:
off_avg = sum(offs) / len(offs)
on_avg = sum(ons) / len(ons)
print(f"{name},avg,{off_avg / N * 1e6:.2f},{on_avg / N * 1e6:.2f},"
f"{(on_avg - off_avg) / N * 1e6:+.2f},{on_avg / off_avg:.4f}", flush=True)

def main():
rank = int(os.environ.get("RANK", "0"))
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
torch.npu.set_device(f"npu:{local_rank}")
dist.init_process_group(backend="hccl")
x = torch.full((BS, HIDDEN), 1.0, dtype=torch.bfloat16, device=f"npu:{local_rank}")
alt = torch.npu.Stream()
compare(capture(lambda: fj_block(x, alt)), "fj_block_K10", rank)
compare(capture(lambda: ar_space(x)), "ar_K10", rank)
dist.destroy_process_group()

if name == "main":
main()