willychan21/ParallelKernelBench_Problems

1from typing import List, Optional2 3import torch4import torch.distributed as dist5 6 7def _compute_split_shapes(size: int, num_chunks: int) -> List[int]:8    if num_chunks == 1:9        return [size]10    chunk_size = (size + num_chunks - 1) // num_chunks11    last_chunk_size = max(0, size - chunk_size * (num_chunks - 1))12    if last_chunk_size == 0:13        chunk_size = size // num_chunks14        last_chunk_size = size - chunk_size * (num_chunks - 1)15    return [chunk_size for _ in range(num_chunks - 1)] + [last_chunk_size]16 17 18def _transpose(19    tensor: torch.Tensor,20    dim0: int,21    dim1: int,22    dim1_split_sizes: List[int],23    group: dist.ProcessGroup,24) -> tuple[list[torch.Tensor], list[int]]:25    comm_size = dist.get_world_size(group=group)26    comm_rank = dist.get_rank(group=group)27 28    tsplit = torch.split(tensor, _compute_split_shapes(tensor.shape[dim0], comm_size), dim=dim0)29    x_send = [y.contiguous() for y in tsplit]30    x_send_shapes = [x.shape for x in x_send]31    x_recv = []32    x_shape = list(x_send_shapes[comm_rank])33    for dim1_len in dim1_split_sizes:34        x_shape[dim1] = dim1_len35        x_recv.append(torch.empty(x_shape, dtype=tensor.dtype, device=tensor.device))36 37    dist.all_to_all(x_recv, x_send, group=group)38    dim0_split_sizes = [x[dim0] for x in x_send_shapes]39    return x_recv, dim0_split_sizes40 41 42def _disco_s2_contraction_torch(43    x: torch.Tensor,44    psi: torch.Tensor,45    nlon_out: int,46) -> torch.Tensor:47    psi = psi.to(x.device)48    batch_size, n_chans, nlat_in, nlon_in = x.shape49    kernel_size, nlat_out, _ = psi.shape50    pscale = nlon_in // nlon_out51 52    x = x.reshape(1, batch_size * n_chans, nlat_in, nlon_in).permute(0, 2, 3, 1)53    x = x.expand(kernel_size, -1, -1, -1)54 55    y = torch.zeros(56        nlon_out,57        kernel_size,58        nlat_out,59        batch_size * n_chans,60        device=x.device,61        dtype=x.dtype,62    )63 64    for pout in range(nlon_out):65        y[pout] = torch.bmm(psi, x.reshape(kernel_size, nlat_in * nlon_in, -1))66        x = torch.roll(x, -pscale, dims=2)67 68    y = y.permute(3, 1, 2, 0).reshape(batch_size, n_chans, kernel_size, nlat_out, nlon_out)69    return y70 71 72@torch.no_grad()73def solution(74    x: torch.Tensor,75    psi: torch.Tensor,76    weight: torch.Tensor,77    groups: int,78    nlon_out: int,79    nlon_in: int,80    azimuth_group: Optional[dist.ProcessGroup] = None,81    polar_group: Optional[dist.ProcessGroup] = None,82    bias: Optional[torch.Tensor] = None,83) -> torch.Tensor:84    azimuth_group = azimuth_group or dist.group.WORLD85    polar_group = polar_group or dist.group.WORLD86    azimuth_size = dist.get_world_size(group=azimuth_group)87    polar_size = dist.get_world_size(group=polar_group)88    polar_rank = dist.get_rank(group=polar_group)89 90    lon_in_shapes = _compute_split_shapes(nlon_in, azimuth_size)91    num_chans = x.shape[1]92 93    if azimuth_size > 1:94        xlist, _ = _transpose(x, dim0=1, dim1=-1, dim1_split_sizes=lon_in_shapes, group=azimuth_group)95        x = torch.cat(xlist, dim=-1)96 97    x = _disco_s2_contraction_torch(x, psi, nlon_out)98 99    if polar_size > 1:100        dtype = x.dtype101        xf = x.float().contiguous()102        dist.all_reduce(xf, group=polar_group)103        x = xf.to(dtype)104 105    if polar_size > 1:106        split_shapes = _compute_split_shapes(x.shape[-2], polar_size)107        x = list(torch.split(x, split_shapes, dim=-2))[polar_rank]108 109    if azimuth_size > 1:110        chan_shapes = _compute_split_shapes(num_chans, azimuth_size)111        xlist, _ = _transpose(x, dim0=-1, dim1=1, dim1_split_sizes=chan_shapes, group=azimuth_group)112        x = torch.cat(xlist, dim=1)113 114    B, C, K, H, W = x.shape115    groupsize = C // groups116    x = x.reshape(B, groups, groupsize, K, H, W)117    out = torch.einsum(118        "bgckxy,gock->bgoxy",119        x,120        weight.reshape(groups, -1, weight.shape[1], weight.shape[2]),121    ).contiguous()122    out = out.reshape(out.shape[0], -1, H, W)123 124    if bias is not None:125        out = out + bias.reshape(1, -1, 1, 1)126 127    return out128