mscclpp/executor_test.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import argparse
from mscclpp import (
    DataType,
    Executor,
    ExecutionPlan,
    PacketType,
    npkit,
    env,
)
from mscclpp import CommGroup, GpuBuffer
from mscclpp.utils import KernelBuilder, pack
import os
import struct

import cupy as cp
from mpi4py import MPI


def parse_dtype(dtype_str):
    dtype_str = dtype_str.strip().lower()
    if dtype_str == "float16":
        return cp.float16
    elif dtype_str == "float32":
        return cp.float32
    elif dtype_str == "int32":
        return cp.int32
    else:
        raise ValueError(f"Unknown data type: {dtype_str}")


def parse_size(size_str):
    size_str = size_str.strip()
    if not size_str:
        raise ValueError("Size string can not be empty")
    units = {"K": 1024, "M": 1024**2, "G": 1024**3}
    if size_str[-1].upper() in units:
        return int(size_str[:-1]) * units[size_str[-1].upper()]
    else:
        return int(size_str)


def dtype_to_mscclpp_dtype(dtype):
    if dtype == cp.float16:
        return DataType.float16
    elif dtype == cp.float32:
        return DataType.float32
    elif dtype == cp.int32:
        return DataType.int32
    else:
        raise ValueError(f"Unknown data type: {dtype}")


def bench_time(n_iters: int, n_graph_iters: int, func_iter):
    """
    Capture CUDA graph for n_iters launches. func_iter(stream, i) must vary slot by i.
    """
    stream = cp.cuda.Stream(non_blocking=True)
    with stream:
        stream.begin_capture()
        for i in range(n_iters):
            func_iter(stream, i)
        graph = stream.end_capture()

    # warmup
    graph.launch(stream)

    start = cp.cuda.Event()
    end = cp.cuda.Event()

    start.record(stream)
    for _ in range(n_graph_iters):
        graph.launch(stream)
    end.record(stream)
    end.synchronize()

    # us per iteration
    return cp.cuda.get_elapsed_time(start, end) / n_iters * 1000.0 / n_graph_iters


def bench_correctness(
    collective: str,
    input_slot: cp.ndarray,
    result_slot: cp.ndarray,
    test_buf: cp.ndarray,
    dtype_str: str,
    rank: int,
    num_ranks: int,
    n_iters: int,
    func_iter,
):
    """
    Correctness check on ONE per-iteration slot view (input_slot/result_slot change per i via func_iter).
    We pass the per-iteration element count to verifier kernels.
    """
    type_size = cp.dtype(parse_dtype(dtype_str)).itemsize
    nelems_per_iter = input_slot.nbytes // type_size

    print("collective: ", collective)

    fill_data_kernel_name = "fill_data_%s" % dtype_str
    if "allgather" in collective:
        coll = "all_gather"
    elif "reducescatter" in collective:
        coll = "reduce_scatter"
    elif "allreduce" in collective:
        coll = "all_reduce"
    else:
        coll = "sendrecv"
    test_data_kernel_name = "test_data_%s_%s" % (coll, dtype_str)

    file_dir = os.path.dirname(os.path.abspath(__file__))
    fill_data_kernel = KernelBuilder(
        file="executor_test_verifier.cu", kernel_name=fill_data_kernel_name, file_dir=file_dir
    ).get_compiled_kernel()
    test_data_kernel = KernelBuilder(
        file="executor_test_verifier.cu", kernel_name=test_data_kernel_name, file_dir=file_dir
    ).get_compiled_kernel()

    nblocks = 64
    nthreads = 1024

    stream = cp.cuda.Stream(non_blocking=True)
    with stream:
        stream.begin_capture()
        for i in range(n_iters):
            # WARNING: input_slot/result_slot variables are placeholders; actual slot views are chosen inside func_iter.
            # We only use these kernels with the CURRENT slot views computed below for this iteration.
            func_iter(stream, i, do_verify=True, fill_kernel=fill_data_kernel, test_kernel=test_data_kernel,
                      nblocks=nblocks, nthreads=nthreads, nelems_per_iter=nelems_per_iter,
                      test_buf=test_buf, rank=rank, num_ranks=num_ranks)
        graph = stream.end_capture()

    graph.launch(stream)
    stream.synchronize()


def build_bufs_sendrecv_ring(size_bytes: int, slots: int, dtype: cp.dtype):
    """
    Build ring buffers for sendrecv:
      - per-iteration message bytes = size_bytes
      - total allocated bytes per buffer = slots * size_bytes
    """
    type_size = cp.dtype(dtype).itemsize
    assert (size_bytes % type_size) == 0, "size not multiple of dtype size"

    nelems_per_iter = size_bytes // type_size
    total_nelems = nelems_per_iter * slots

    input_buf = GpuBuffer(total_nelems, dtype=dtype)
    result_buf = GpuBuffer(total_nelems, dtype=dtype)
    test_buf = cp.zeros(nelems_per_iter, dtype=dtype)  # expected for one iteration

    return input_buf, result_buf, test_buf, nelems_per_iter


def main(
    execution_plan_path: str,
    size: int,                 # per-iteration bytes
    in_place: bool = True,
    dtype_str: str = "float16",
    packet_type: PacketType = PacketType.LL16,
    n_iters: int = 10,
    n_graph_iters: int = 10,
    slots: int = 4,            # ring buffer depth
):
    mscclpp_group = CommGroup(MPI.COMM_WORLD)
    cp.cuda.Device(mscclpp_group.my_rank % mscclpp_group.nranks_per_node).use()

    executor = Executor(mscclpp_group.communicator)

    npkit_dump_dir = env().npkit_dump_dir
    if npkit_dump_dir != "":
        npkit.init(mscclpp_group.my_rank)

    execution_plan = ExecutionPlan(execution_plan_path, mscclpp_group.my_rank)
    collective = execution_plan.collective

    dtype = parse_dtype(dtype_str)

    # We only change allocation/behavior for sendrecv
    if "sendrecv" in collective.lower():
        input_buf, result_buf, test_buf, nelems_per_iter = build_bufs_sendrecv_ring(size, slots, dtype)
        type_size = cp.dtype(dtype).itemsize
        bytes_per_iter = nelems_per_iter * type_size

        def slot_view(buf, slot_idx):
            start = slot_idx * nelems_per_iter
            end = start + nelems_per_iter
            return buf[start:end]

        # Iteration-aware executor call (rotates slot each iteration)
        def executor_func_iter(stream, i, do_verify=False, **vk):
            slot = i % slots
            in_slot = slot_view(input_buf, slot)
            out_slot = slot_view(result_buf, slot)

            if do_verify:
                # Fill per-iteration input slot with unique (rank, i) pattern
                fill_data_kernel = vk["fill_kernel"]
                test_data_kernel = vk["test_kernel"]
                nblocks = vk["nblocks"]
                nthreads = vk["nthreads"]
                nelems = vk["nelems_per_iter"]
                test_buf_local = vk["test_buf"]
                rank = vk["rank"]
                num_ranks = vk["num_ranks"]

                fill_params = pack(in_slot) + struct.pack("Q", nelems) + pack(rank, i)
                fill_data_kernel.launch_kernel(fill_params, nblocks, nthreads, 0, stream)

            # Execute exactly one per-iteration message: bytes_per_iter == user --size
            executor.execute(
                mscclpp_group.my_rank,
                in_slot.data.ptr,
                out_slot.data.ptr,
                in_slot.nbytes,
                out_slot.nbytes,
                dtype_to_mscclpp_dtype(dtype),
                execution_plan,
                stream.ptr,
                packet_type,
            )

            if do_verify:
                # Validate the output slot for this iteration i
                test_params = (
                    pack(out_slot, test_buf_local)
                    + struct.pack("Q", nelems)
                    + pack(num_ranks, rank, i)
                )
                test_data_kernel.launch_kernel(test_params, nblocks, nthreads, 0, stream)

        # One-shot sentinel check (slot 0)
        mscclpp_group.barrier()
        print("per-iter size= ", bytes_per_iter, "bytes, slots=", slots, "total buffer bytes=", input_buf.nbytes)

        # Fill whole result with sentinel then run ONE iter (i=0)
        result_buf.fill(cp.asarray(123.0, dtype=dtype))
        cp.cuda.runtime.deviceSynchronize()

        stream = cp.cuda.Stream(non_blocking=True)
        with stream:
            executor_func_iter(stream, 0)
        stream.synchronize()

        # Count changes only in slot 0 region
        out0 = slot_view(result_buf, 0)
        changed = cp.count_nonzero(out0 != cp.asarray(123.0, dtype=dtype)).item()
        print("changed elements in slot0:", changed, "out of", out0.size)

        cp.cuda.runtime.deviceSynchronize()
        mscclpp_group.barrier()

        # Correctness: fills + executes + tests with unique i and rotating slots
        bench_correctness(
            collective,
            slot_view(input_buf, 0),   # placeholder; real slot chosen per i
            slot_view(result_buf, 0),  # placeholder; real slot chosen per i
            test_buf,
            dtype_str,
            mscclpp_group.my_rank,
            mscclpp_group.nranks,
            n_iters,
            executor_func_iter,
        )

        mscclpp_group.barrier()

        # Timing (CUDA graph captures n_iters launches with varying slot pointers)
        execution_time = bench_time(n_iters, n_graph_iters, executor_func_iter)

        if npkit_dump_dir is not None:
            npkit.dump(npkit_dump_dir)
            npkit.shutdown()

        print(
            f"Rank: {mscclpp_group.my_rank} Execution time: {execution_time} us, "
            f"per-iter data size: {bytes_per_iter} bytes dtype: {dtype().dtype.name} "
            f"bandwidth: {bytes_per_iter / (execution_time * 1e-6) / (1024**3):.2f} GB/s, "
            f"packet type: {packet_type}, slots: {slots}"
        )

    else:
        raise RuntimeError(
            f"This rewritten executor_test.py currently specializes sendrecv. "
            f"Plan collective was: {collective}"
        )

    executor = None
    mscclpp_group = None


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-path", "--execution_plan_path", type=str, required=True)
    parser.add_argument("--size", type=str, required=True, help="PER-ITERATION bytes (e.g., 1K, 4M, 1G)")
    parser.add_argument("--in_place", action="store_true", help="flag to define an in-place operation")
    parser.add_argument("--dtype", type=str, default="float16", help="Choose from float16, float32, int32")
    parser.add_argument("--packet_type", type=str, default="LL16", help="Choose from LL8, LL16")
    parser.add_argument("--n_iters", type=int, default=10)
    parser.add_argument("--n_graph_iters", type=int, default=10)
    parser.add_argument("--slots", type=int, default=4, help="ring buffer depth; rotates slot = iter % slots")
    args = parser.parse_args()

    packet_type = PacketType.LL16
    if args.packet_type == "LL8":
        packet_type = PacketType.LL8

    per_iter_size = parse_size(args.size)

    main(
        args.execution_plan_path,
        per_iter_size,
        args.in_place,
        args.dtype,
        packet_type,
        args.n_iters,
        args.n_graph_iters,
        args.slots,
    )