willychan21/ParallelKernelBench_Problems

1import math2from typing import List, Optional, Tuple, Union3 4import torch5import torch.distributed as dist6import torch.nn.functional as F7 8_CONV3D_NUMEL_LIMIT = 2**319 10 11def _to_3tuple(value: Union[int, Tuple[int, int, int]]) -> Tuple[int, int, int]:12    return (value, value, value) if isinstance(value, int) else value13 14 15def _ceil_to_divisible(n: int, dividend: int) -> int:16    return math.ceil(dividend / (dividend // n))17 18 19def _output_shape(20    input_shape: torch.Size,21    out_channels: int,22    kernel_size: Tuple[int, int, int],23    stride: Tuple[int, int, int],24    padding: Tuple[int, int, int],25    dilation: Tuple[int, int, int],26) -> List[int]:27    shape = [input_shape[0], out_channels]28    for idx, size in enumerate(input_shape[-3:]):29        out = (size + 2 * padding[idx] - dilation[idx] * (kernel_size[idx] - 1) - 1)30        shape.append(math.floor(out / stride[idx] + 1))31    return shape32 33 34def _chunk_count(numel: int, channels: int, limit: int) -> int:35    chunks = math.ceil(numel / limit)36    return _ceil_to_divisible(chunks, channels)37 38 39def _channel_chunk_conv3d(40    x: torch.Tensor,41    weight: torch.Tensor,42    bias: Optional[torch.Tensor],43    stride: Tuple[int, int, int],44    padding: Tuple[int, int, int],45    dilation: Tuple[int, int, int],46    groups: int,47    numel_limit: int,48) -> torch.Tensor:49    out_channels, in_channels = weight.shape[:2]50    output_shape = _output_shape(51        x.shape,52        out_channels,53        tuple(weight.shape[2:]),54        stride,55        padding,56        dilation,57    )58    in_chunks = _chunk_count(x.numel(), in_channels, numel_limit)59    out_chunks = _chunk_count(math.prod(output_shape), out_channels, numel_limit)60    if in_chunks == 1 and out_chunks == 1:61        return F.conv3d(x, weight, bias, stride, padding, dilation, groups)62 63    x_chunks = x.chunk(in_chunks, dim=1)64    weight_out_chunks = weight.chunk(out_chunks, dim=0)65    bias_chunks = bias.chunk(out_chunks) if bias is not None else [None] * out_chunks66    outputs: List[torch.Tensor] = []67    for weight_chunk, bias_chunk in zip(weight_out_chunks, bias_chunks):68        partial_sum: Optional[torch.Tensor] = None69        for x_chunk, w_chunk in zip(x_chunks, weight_chunk.chunk(in_chunks, dim=1)):70            partial = F.conv3d(71                x_chunk,72                w_chunk,73                None,74                stride,75                padding,76                dilation,77                groups,78            ).float()79            partial_sum = partial if partial_sum is None else partial_sum + partial80        if partial_sum is None:81            raise RuntimeError("conv3d chunking produced no partial outputs")82        out = partial_sum.to(dtype=x.dtype)83        if bias_chunk is not None:84            out = out + bias_chunk.view(1, -1, 1, 1, 1)85        outputs.append(out)86    return torch.cat(outputs, dim=1)87 88 89@torch.no_grad()90def solution(91    input: torch.Tensor,92    weight: torch.Tensor,93    bias: Optional[torch.Tensor],94    stride: Union[int, Tuple[int, int, int]],95    padding: Union[int, Tuple[int, int, int]],96    dilation: Union[int, Tuple[int, int, int]],97    groups: int = 1,98    group: Optional[dist.ProcessGroup] = None,99) -> torch.Tensor:100    group = group or dist.group.WORLD101    out = _channel_chunk_conv3d(102        input,103        weight,104        None,105        _to_3tuple(stride),106        _to_3tuple(padding),107        _to_3tuple(dilation),108        groups,109        _CONV3D_NUMEL_LIMIT,110    )111    dist.all_reduce(out, op=dist.ReduceOp.SUM, group=group)112    if bias is not None:113        out = out + bias.view(1, -1, 1, 1, 1)114    return out115