"""Lab 5: split-K with a "last CTA reduces" semaphore, and the atomic variant, simulated.

Run:   python3 lab5_splitk.py
Needs: Python 3 only.

When a matrix has few rows (the router has 256 = only 16 row blocks of 16), one CTA per row block
would leave most of the 132 SMs idle. Split-K cuts the K dimension into S pieces, one CTA each, and
adds the S partial results at the end. This lab simulates the two protocols in submission.py:

  serial   (gemv/gemvz lines 178-194, Triton _epilogue with ATOM=False): every CTA writes its partial
           to a workspace, then atomically increments a per-tile counter. The CTA that sees S-1 (it is
           the last to arrive) sums all S partials, writes the output and resets the counter to 0.
  atomic   (Triton _epilogue with ATOM=True, used for the router at m=128/256): every CTA atomically
           adds its partial into a zeroed accumulator; the last CTA copies it out and re-zeroes it.

CTAs finish in a random order on a real GPU, so each "replay" shuffles the order. Replays run the same
CUDA graph again and again, so whatever state a call leaves behind is what the next call starts with.

Things to try
  - Set RESET = False: the counter keeps counting, and on the second replay nobody sees S-1.
  - Set REZERO = False: the atomic accumulator keeps the previous call's sum.
"""
import random

RESET = True    # serial protocol: the last CTA writes sem[tile] = 0     (gemv line 193)
REZERO = True   # atomic protocol: the last CTA stores zeros into the accumulator (Triton line 494)
S = 4           # K-splits = CTAs per output tile


class Gpu:
    """Global memory that survives between replays (like state['ws'], state['sem'], state['racc'])."""
    def __init__(self):
        self.sem = 0
        self.ws = [None] * S
        self.racc = 0.0
        self.out = None


def atomic_add(gpu, field, v):
    """Returns the OLD value, like CUDA's atomicAdd."""
    old = getattr(gpu, field)
    setattr(gpu, field, old + v)
    return old


def replay_serial(gpu, parts, order):
    for sidx in order:                          # CTAs finish in this order
        gpu.ws[sidx] = parts[sidx]              # ws[(rb*S + sidx)*...] = r            (line 179)
        # __threadfence(): make the partial visible GPU-wide BEFORE the counter moves (line 180)
        last = atomic_add(gpu, "sem", 1) == S - 1   # s_last = atomicAdd(&sem[rb], 1) == S-1 (182)
        if not last:
            continue                            # if (!s_last) return;                  (line 184)
        gpu.out = sum(gpu.ws)                   # sum the S partials (__ldcg reads, L2)  (186-191)
        if RESET:
            gpu.sem = 0                         # sem[rb] = 0 for the next call         (line 193)


def replay_atomic(gpu, parts, order):
    for sidx in order:
        atomic_add(gpu, "racc", parts[sidx])    # tl.atomic_add(a_ptrs, acc, sem="relaxed") (484)
        last = atomic_add(gpu, "sem", 1) == S - 1   # tl.atomic_add(sem_ptr + tile, 1, sem="acq_rel")
        if last:
            gpu.out = gpu.racc                  # load the finished sum                   (488)
            if REZERO:
                gpu.racc = 0.0                  # tl.store(a_ptrs, zeros)                 (494)
            gpu.sem = 0                         # tl.atomic_xchg(sem_ptr + tile, 0)       (495)


def main():
    random.seed(1)
    for name, fn in (("serial", replay_serial), ("atomic", replay_atomic)):
        gpu = Gpu()
        print(f"--- {name} protocol, S = {S} ---")
        for rep in range(3):
            parts = [random.randint(1, 9) for _ in range(S)]   # fresh inputs every call
            order = random.sample(range(S), S)
            gpu.out = None
            fn(gpu, parts, order)
            ok = gpu.out == sum(parts)
            print(f"  replay {rep}: partials {parts}, finish order {order} -> out {gpu.out}, "
                  f"expected {sum(parts)}  {'OK' if ok else 'WRONG'}")
    print("\nWhy the fences matter (not visible in this simulation): on a GPU, another CTA could see the\n"
          "counter increment before it sees the partial written just before it. __threadfence() (CUDA) and\n"
          "sem='acq_rel' (Triton) forbid that reordering, and __ldcg / cache_modifier='.cg' make the last\n"
          "CTA read the partials from L2 instead of a possibly stale L1 copy.")


if __name__ == "__main__":
    main()