willychan21/ParallelKernelBench_Problems

1from typing import List, Optional, Tuple2 3import torch4import torch.distributed as dist5 6 7def _topk_with_labels(8    similarity: torch.Tensor,9    labels: torch.Tensor,10    k: int,11) -> Tuple[torch.Tensor, torch.Tensor]:12    topk_sims, indices = similarity.topk(k, dim=1, largest=True, sorted=True)13    topk_labels = torch.gather(labels.expand(similarity.shape[0], -1), 1, indices)14    return topk_sims, topk_labels15 16 17def _broadcast_queries(18    queries: torch.Tensor,19    source: int,20    rank: int,21    group: dist.ProcessGroup,22) -> torch.Tensor:23    shape = torch.tensor(queries.shape, dtype=torch.long, device=queries.device)24    dist.broadcast(shape, src=source, group=group)25    if rank == source:26        out = queries.contiguous()27    else:28        out = queries.new_empty(tuple(int(v) for v in shape.tolist()))29    dist.broadcast(out, src=source, group=group)30    return out31 32 33def _local_candidates(34    queries: torch.Tensor,35    train_features_t: torch.Tensor,36    train_labels: torch.Tensor,37    k: int,38) -> Tuple[torch.Tensor, torch.Tensor]:39    similarity = queries @ train_features_t40    return _topk_with_labels(similarity, train_labels, k)41 42 43def _merge_on_owner(44    topk_sims: torch.Tensor,45    topk_labels: torch.Tensor,46    owner: int,47    rank: int,48    world_size: int,49    k: int,50    group: dist.ProcessGroup,51) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:52    gathered_sims: Optional[List[torch.Tensor]] = None53    gathered_labels: Optional[List[torch.Tensor]] = None54    if rank == owner:55        gathered_sims = [torch.empty_like(topk_sims) for _ in range(world_size)]56        gathered_labels = [torch.empty_like(topk_labels) for _ in range(world_size)]57 58    dist.gather(topk_sims, gather_list=gathered_sims, dst=owner, group=group)59    dist.gather(topk_labels, gather_list=gathered_labels, dst=owner, group=group)60    if rank != owner:61        return None62 63    all_sims = torch.cat(gathered_sims, dim=1)64    all_labels = torch.cat(gathered_labels, dim=1)65    return _topk_with_labels(all_sims, all_labels, k)66 67 68@torch.no_grad()69def solution(70    test_features_rank: torch.Tensor,71    train_features_rank_T: torch.Tensor,72    train_labels_rank: torch.Tensor,73    max_k: int,74    group: Optional[dist.ProcessGroup] = None,75) -> Tuple[torch.Tensor, torch.Tensor]:76    group = group or dist.group.WORLD77    rank = dist.get_rank(group=group)78    world_size = dist.get_world_size(group=group)79    if max_k > train_features_rank_T.shape[1]:80        raise ValueError("max_k must not exceed the local train shard size")81 82    result: Optional[Tuple[torch.Tensor, torch.Tensor]] = None83    for owner in range(world_size):84        queries = _broadcast_queries(test_features_rank, owner, rank, group)85        topk_sims, topk_labels = _local_candidates(86            queries,87            train_features_rank_T,88            train_labels_rank,89            max_k,90        )91        merged = _merge_on_owner(92            topk_sims,93            topk_labels,94            owner,95            rank,96            world_size,97            max_k,98            group,99        )100        if merged is not None:101            result = merged102 103    if result is None:104        raise RuntimeError("k-NN ring did not produce local results")105    return result106