willychan21/ParallelKernelBench_Problems

1from typing import Optional, Sequence2 3import torch4import torch.distributed as dist5import torch.nn.functional as F6 7 8def _all_to_all_transpose(9    tensor: torch.Tensor,10    split_dim: int,11    group: dist.ProcessGroup,12) -> list[torch.Tensor]:13    world_size = dist.get_world_size(group)14    chunk = tensor.shape[split_dim] // world_size15    send = [x.contiguous() for x in torch.split(tensor, chunk, dim=split_dim)]16    recv = [torch.empty_like(send[0]) for _ in range(world_size)]17    dist.all_to_all(recv, send, group=group)18    return recv19 20 21def _pad_zero(tensor: torch.Tensor, dim: int, size: int) -> torch.Tensor:22    dim = dim % tensor.ndim23    pad = [0] * (2 * (tensor.ndim - dim))24    pad[1] = size - tensor.shape[dim]25    return F.pad(tensor, pad, mode="constant", value=0.0)26 27 28def _gather_dim(29    tensor: torch.Tensor,30    dim: int,31    group: dist.ProcessGroup,32) -> torch.Tensor:33    world_size = dist.get_world_size(group)34    chunks = [torch.empty_like(tensor) for _ in range(world_size)]35    dist.all_gather(chunks, tensor, group=group)36    return torch.cat(chunks, dim=dim).contiguous()37 38 39def _scatter_dim(40    tensor: torch.Tensor,41    dim: int,42    group: dist.ProcessGroup,43) -> torch.Tensor:44    world_size = dist.get_world_size(group)45    rank = dist.get_rank(group)46    chunks = torch.split(tensor, tensor.shape[dim] // world_size, dim=dim)47    return chunks[rank].contiguous()48 49 50def _conj_pad_2d(51    tensor: torch.Tensor,52    pad_dim: int,53    other_dim: int,54    size: int,55    group: dist.ProcessGroup,56) -> torch.Tensor:57    pad_dim = pad_dim % tensor.ndim58    other_dim = other_dim % tensor.ndim59    orig_size = tensor.shape[pad_dim]60 61    tensor_pad = _pad_zero(tensor, pad_dim, size)62    lhs_slice = [slice(0, s) for s in tensor.shape]63    lhs_slice[pad_dim] = slice(orig_size, size)64    rhs_slice = [slice(0, s) for s in tensor.shape]65    rhs_slice[pad_dim] = slice(1, size - orig_size + 1)66    tensor_pad[tuple(lhs_slice)] = torch.flip(torch.conj(tensor_pad[tuple(rhs_slice)]), dims=[pad_dim])67 68    tensor_pad = _gather_dim(tensor_pad, other_dim, group)69    flip_slice = [slice(0, s) for s in tensor_pad.shape]70    flip_slice[pad_dim] = slice(orig_size, size)71    flip_slice[other_dim] = slice(1, tensor_pad.shape[other_dim])72    tensor_pad[tuple(flip_slice)] = torch.flip(tensor_pad[tuple(flip_slice)], dims=[other_dim])73    return _scatter_dim(tensor_pad, other_dim, group)74 75 76@torch.no_grad()77def solution(78    x: torch.Tensor,79    s: Optional[Sequence[int]],80    dim: Sequence[int],81    norm: str = "ortho",82    group: Optional[dist.ProcessGroup] = None,83) -> torch.Tensor:84    group = group or dist.group.WORLD85    dim0, dim1 = int(dim[0]), int(dim[1])86    if s is not None:87        first_dim_size = int(s[0])88        last_dim_size = int(s[1])89    else:90        first_dim_size = int(x.shape[dim0])91        last_dim_size = int(2 * (x.shape[dim1] - 1))92 93    x_pad = _conj_pad_2d(x, pad_dim=dim1, other_dim=dim0, size=last_dim_size, group=group)94 95    x1 = torch.fft.ifft(x_pad, n=last_dim_size, dim=dim1, norm=norm)96 97    x1_recv = _all_to_all_transpose(x1, split_dim=dim1, group=group)98    x1_tran = torch.cat(x1_recv, dim=dim0)99 100    x2 = torch.fft.ifft(x1_tran, n=first_dim_size, dim=dim0, norm=norm)101    return torch.real(x2).contiguous()