willychan21/ParallelKernelBench_Problems

1from typing import Optional2 3import torch4import torch.distributed as dist5import torch.nn.functional as F6 7 8def _gather_boundaries(9    x: torch.Tensor,10    boundary: int,11    group: dist.ProcessGroup,12) -> list[torch.Tensor]:13    if boundary == 0:14        empty = x[:, :, :0, :]15        return [torch.stack([empty, empty], dim=0)]16 17    local = torch.stack([x[:, :, :boundary, :], x[:, :, -boundary:, :]], dim=0)18    gathered = [torch.empty_like(local) for _ in range(dist.get_world_size(group))]19    dist.all_gather(gathered, local.contiguous(), group=group)20    return gathered21 22 23@torch.no_grad()24def solution(25    x: torch.Tensor,26    weight: torch.Tensor,27    bias: Optional[torch.Tensor],28    stride: int = 1,29    padding: int = 1,30    group: Optional[dist.ProcessGroup] = None,31) -> torch.Tensor:32    group = group or dist.group.WORLD33    world_size = dist.get_world_size(group)34    rank = dist.get_rank(group)35    boundary = int(padding)36 37    if boundary == 0 or world_size == 1:38        return F.conv2d(x, weight, bias, stride=stride, padding=padding)39 40    boundaries = _gather_boundaries(x, boundary, group)41    pieces = []42    if rank == 0:43        pieces.append(x.new_zeros(*x.shape[:2], boundary, x.shape[-1]))44    else:45        pieces.append(boundaries[rank - 1][1])46    pieces.append(x)47    if rank == world_size - 1:48        pieces.append(x.new_zeros(*x.shape[:2], boundary, x.shape[-1]))49    else:50        pieces.append(boundaries[rank + 1][0])51 52    padded_x = torch.cat(pieces, dim=2)53    return F.conv2d(padded_x, weight, bias, stride=stride, padding=(0, padding))54