willychan21/ParallelKernelBench_Problems

1from typing import List, Optional, Tuple2 3import torch4import torch.distributed as dist5 6 7def _local_sizes(group: dist.ProcessGroup, device: torch.device, local_n: int) -> List[int]:8    world_size = dist.get_world_size(group=group)9    size = torch.tensor([local_n], dtype=torch.long, device=device)10    gathered = [torch.empty_like(size) for _ in range(world_size)]11    dist.all_gather(gathered, size, group=group)12    return [int(item.item()) for item in gathered]13 14 15def _active_rank_info(rank: int, sizes: List[int]) -> Tuple[List[int], int]:16    active = [idx for idx, size in enumerate(sizes) if size > 0]17    sort_rank = active.index(rank) if rank in active else -118    return active, sort_rank19 20 21def _extract_samples(22    sorted_local: torch.Tensor,23    sort_rank: int,24    n_samples: int,25) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:26    if sort_rank < 0 or sorted_local.numel() == 0:27        values = sorted_local.new_full((n_samples,), float("inf"))28        ranks = torch.full((n_samples,), -1, dtype=torch.long, device=sorted_local.device)29        positions = torch.full_like(ranks, -1)30        return values, ranks, positions31 32    local_n = sorted_local.numel()33    sample_idx = torch.arange(n_samples, dtype=torch.long, device=sorted_local.device)34    valid_count = min(n_samples, local_n)35    values = sorted_local.new_full((n_samples,), float("inf"))36    ranks = torch.full((n_samples,), -1, dtype=torch.long, device=sorted_local.device)37    positions = torch.full_like(ranks, -1)38    if n_samples < local_n:39        valid_positions = ((sample_idx + 1) * local_n).div(n_samples, rounding_mode="floor") - 140    else:41        valid_positions = sample_idx[:valid_count]42    values[:valid_count] = sorted_local[valid_positions[:valid_count]]43    ranks[:valid_count] = sort_rank44    positions[:valid_count] = valid_positions[:valid_count]45    return values, ranks, positions46 47 48def _gather_splitters(49    sample_values: torch.Tensor,50    sample_ranks: torch.Tensor,51    sample_positions: torch.Tensor,52    active_count: int,53    group: dist.ProcessGroup,54) -> List[Tuple[float, int, int]]:55    world_size = dist.get_world_size(group=group)56    value_parts = [torch.empty_like(sample_values) for _ in range(world_size)]57    rank_parts = [torch.empty_like(sample_ranks) for _ in range(world_size)]58    pos_parts = [torch.empty_like(sample_positions) for _ in range(world_size)]59    dist.all_gather(value_parts, sample_values, group=group)60    dist.all_gather(rank_parts, sample_ranks, group=group)61    dist.all_gather(pos_parts, sample_positions, group=group)62 63    values = torch.cat(value_parts).detach().cpu().tolist()64    ranks = torch.cat(rank_parts).detach().cpu().tolist()65    positions = torch.cat(pos_parts).detach().cpu().tolist()66    samples = [67        (float(value), int(sample_rank), int(position))68        for value, sample_rank, position in zip(values, ranks, positions)69        if int(sample_rank) >= 070    ]71    samples.sort(key=lambda item: (item[0], item[1], item[2]))72 73    splitters: List[Tuple[float, int, int]] = []74    usable = len(samples)75    for sort_rank in range(active_count - 1):76        index = (sort_rank + 1) * usable // active_count - 177        splitters.append(samples[max(0, min(index, usable - 1))])78    return splitters79 80 81def _split_positions(82    sorted_local: torch.Tensor,83    splitters: List[Tuple[float, int, int]],84    sort_rank: int,85) -> List[int]:86    if sort_rank < 0:87        return [0] * (len(splitters) + 2)88 89    boundaries = [0]90    for value, splitter_rank, splitter_position in splitters:91        probe = torch.tensor(value, dtype=sorted_local.dtype, device=sorted_local.device)92        if sort_rank > splitter_rank:93            end = int(torch.searchsorted(sorted_local, probe, right=False).item())94        elif sort_rank < splitter_rank:95            end = int(torch.searchsorted(sorted_local, probe, right=True).item())96        else:97            end = int(splitter_position) + 198        boundaries.append(max(boundaries[-1], min(end, sorted_local.numel())))99    boundaries.append(sorted_local.numel())100    return boundaries101 102 103def _variable_all_to_all(104    send_chunks: List[torch.Tensor],105    group: dist.ProcessGroup,106) -> List[torch.Tensor]:107    device = send_chunks[0].device108    dtype = send_chunks[0].dtype109    send_counts = torch.tensor(110        [chunk.numel() for chunk in send_chunks], dtype=torch.long, device=device111    )112    recv_counts = torch.empty_like(send_counts)113    dist.all_to_all_single(recv_counts, send_counts, group=group)114 115    send = (116        torch.cat(send_chunks, dim=0)117        if int(send_counts.sum().item()) > 0118        else torch.empty(0, dtype=dtype, device=device)119    )120    recv = torch.empty(int(recv_counts.sum().item()), dtype=dtype, device=device)121    dist.all_to_all_single(122        recv,123        send,124        output_split_sizes=recv_counts.cpu().tolist(),125        input_split_sizes=send_counts.cpu().tolist(),126        group=group,127    )128 129    outputs: List[torch.Tensor] = []130    offset = 0131    for count in recv_counts.cpu().tolist():132        next_offset = offset + int(count)133        outputs.append(recv[offset:next_offset])134        offset = next_offset135    return outputs136 137 138def _merge_sorted(chunks: List[torch.Tensor], like: torch.Tensor) -> torch.Tensor:139    chunks = [chunk for chunk in chunks if chunk.numel() > 0]140    if not chunks:141        return like.new_empty(0)142    return torch.cat(chunks, dim=0).sort().values143 144 145def _target_range(rank: int, world_size: int, total: int) -> Tuple[int, int]:146    base = total // world_size147    extra = total % world_size148    start = rank * base + min(rank, extra)149    end = start + base + (1 if rank < extra else 0)150    return start, end151 152 153def _redistribute_exact(merged: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:154    world_size = dist.get_world_size(group=group)155    rank = dist.get_rank(group=group)156    sizes = _local_sizes(group, merged.device, merged.numel())157    total = sum(sizes)158 159    bucket_start = sum(sizes[:rank])160    bucket_end = bucket_start + merged.numel()161    send_chunks: List[torch.Tensor] = []162    for dest in range(world_size):163        target_start, target_end = _target_range(dest, world_size, total)164        start = max(bucket_start, target_start)165        end = min(bucket_end, target_end)166        if start < end:167            send_chunks.append(merged[start - bucket_start : end - bucket_start])168        else:169            send_chunks.append(merged.new_empty(0))170    return torch.cat(_variable_all_to_all(send_chunks, group), dim=0)171 172 173@torch.no_grad()174def solution(local_shard: torch.Tensor, group: Optional[dist.ProcessGroup] = None) -> torch.Tensor:175    group = group or dist.group.WORLD176    rank = dist.get_rank(group=group)177    world_size = dist.get_world_size(group=group)178    sorted_local = local_shard.sort().values179 180    initial_sizes = _local_sizes(group, local_shard.device, local_shard.numel())181    active_ranks, sort_rank = _active_rank_info(rank, initial_sizes)182    active_count = len(active_ranks)183    if active_count == 0:184        return local_shard.new_empty(0)185 186    sample_values, sample_ranks, sample_positions = _extract_samples(187        sorted_local, sort_rank, active_count188    )189    splitters = _gather_splitters(190        sample_values, sample_ranks, sample_positions, active_count, group191    )192    boundaries = _split_positions(sorted_local, splitters, sort_rank)193 194    send_chunks = [sorted_local.new_empty(0) for _ in range(world_size)]195    for bucket, dest_rank in enumerate(active_ranks):196        send_chunks[dest_rank] = sorted_local[boundaries[bucket] : boundaries[bucket + 1]].contiguous()197 198    received = _variable_all_to_all(send_chunks, group)199    merged = _merge_sorted(received, sorted_local)200    return _redistribute_exact(merged, group)201