willychan21/ParallelKernelBench_Problems

1from typing import List, Optional2 3import torch4import torch.distributed as dist5 6 7def _shift(chunks: List[torch.Tensor], group: dist.ProcessGroup) -> List[torch.Tensor]:8    cutoff = len(chunks) - dist.get_rank(group)9    return chunks[cutoff:] + chunks[:cutoff]10 11 12def _all_to_all(13    outputs: List[torch.Tensor],14    inputs: List[torch.Tensor],15    group: dist.ProcessGroup,16) -> None:17    outputs = _shift(list(outputs), group)18    inputs = _shift(list(inputs), group)19    if outputs and outputs[0].is_cuda:20        dist.all_to_all(outputs, inputs, group=group)21        return22 23    output_splits = [out.size(0) for out in outputs]24    input_splits = [inp.size(0) for inp in inputs]25    flat_out = torch.cat(outputs) if outputs else torch.empty(0)26    flat_in = torch.cat(inputs) if inputs else torch.empty(0)27    dist.all_to_all_single(28        flat_out,29        flat_in,30        output_split_sizes=output_splits,31        input_split_sizes=input_splits,32        group=group,33    )34    for out, temp in zip(outputs, flat_out.split(output_splits)):35        out.copy_(temp)36 37 38@torch.no_grad()39def solution(40    local_features: torch.Tensor,41    seed_inverse_ids: torch.Tensor,42    counts_sent: List[int],43    counts_received: List[int],44    group: Optional[dist.ProcessGroup] = None,45) -> torch.Tensor:46    group = group or dist.group.WORLD47    gathered = local_features[seed_inverse_ids]48    out = local_features.new_empty((sum(counts_sent),) + local_features.shape[1:])49    _all_to_all(50        list(torch.split(out, counts_sent)),51        list(torch.split(gathered, counts_received)),52        group,53    )54    return out