willychan21/ParallelKernelBench_Problems huggingface dataset:
from typing import List, Optional, Tuple

import torch
import torch.distributed as dist

_NO_OBJ_LOGIT = -10.0


def _mask_iou(lhs: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor:
    lhs_flat = lhs.flatten(1).float()
    rhs_flat = rhs.flatten(1).float()
    intersection = lhs_flat @ rhs_flat.T
    lhs_area = lhs_flat.sum(dim=1)
    rhs_area = rhs_flat.sum(dim=1)
    union = lhs_area[:, None] + rhs_area[None, :] - intersection
    return intersection / union.clamp_min(1.0)


def _all_gather_variable(
    tensor: torch.Tensor,
    counts: List[int],
    group: dist.ProcessGroup,
) -> torch.Tensor:
    recv = [tensor.new_empty((count, *tensor.shape[1:])) for count in counts]
    dist.all_gather(recv, tensor.contiguous(), group=group)
    return torch.cat(recv, dim=0)


def _suppression_mask(
    masks: torch.Tensor,
    last_occluded: torch.Tensor,
    iou_threshold: float,
    reverse: bool,
) -> torch.Tensor:
    num_objects = masks.shape[0]
    suppress = torch.zeros(num_objects, dtype=torch.bool, device=masks.device)
    if num_objects <= 1:
        return suppress

    overlaps = torch.triu(_mask_iou(masks, masks) >= iou_threshold, diagonal=1)
    last_i = last_occluded.view(num_objects, 1)
    last_j = last_occluded.view(1, num_objects)
    cmp = torch.lt if reverse else torch.gt

    suppress_i = overlaps & cmp(last_i, last_j) & (last_j > -1)
    suppress_j = overlaps & cmp(last_j, last_i) & (last_i > -1)
    return suppress_i.any(dim=1) | suppress_j.any(dim=0)


@torch.no_grad()
def solution(
    low_res_masks_local: torch.Tensor,
    obj_scores_local: torch.Tensor,
    num_obj_per_gpu: List[int],
    last_occluded: torch.Tensor,
    iou_threshold: float = 0.7,
    reverse: bool = False,
    group: Optional[dist.ProcessGroup] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    group = group or dist.group.WORLD
    rank = dist.get_rank(group=group)
    expected = int(num_obj_per_gpu[rank])
    if low_res_masks_local.shape[0] != expected:
        raise ValueError("local mask count does not match num_obj_per_gpu")
    if obj_scores_local.shape[0] != expected:
        raise ValueError("local score count does not match num_obj_per_gpu")

    masks_local = low_res_masks_local.float().contiguous()
    scores_local = obj_scores_local.float().contiguous()
    masks_global = _all_gather_variable(masks_local, num_obj_per_gpu, group)
    scores_global = _all_gather_variable(scores_local, num_obj_per_gpu, group)

    last_occluded = last_occluded.to(device=masks_global.device, dtype=torch.long)
    binary_masks = masks_global > 0
    to_suppress = _suppression_mask(
        binary_masks, last_occluded, iou_threshold=iou_threshold, reverse=reverse
    )
    masks_global[to_suppress] = _NO_OBJ_LOGIT
    return masks_global, scores_global, to_suppress
from typing import Optional, Sequence

import torch
import torch.distributed as dist


_COEFFICIENTS: dict[str, Sequence[tuple[float, float, float]]] = {
    "simple": ((3.4445, -4.7750, 2.0315),),
    "quintic": (
        (4.0848, -6.8946, 2.9270),
        (3.9505, -6.3029, 2.6377),
        (3.7418, -5.5913, 2.3037),
        (2.8769, -3.1427, 1.2046),
        (2.8366, -3.0525, 1.2012),
    ),
    "polar_express": (
        (8.2051, -22.9019, 16.4607),
        (4.0664, -2.8612, 0.5184),
        (3.9096, -2.8234, 0.5250),
        (3.2856, -2.4647, 0.5074),
        (2.2779, -1.6447, 0.4162),
        (1.8726, -1.2307, 0.3585),
        (1.8564, -1.2132, 0.3568),
        (1.8750, -1.2500, 0.3750),
    ),
    "aol": (
        (4.0098, -7.0585, 2.4635),
        (3.4585, -5.5479, 2.5959),
        (2.7573, -3.2939, 1.4254),
        (2.7215, -3.0494, 1.3169),
    ),
}


def _coefficient_at(
    coefficients: Sequence[tuple[float, float, float]],
    step: int,
) -> tuple[float, float, float]:
    return coefficients[step % len(coefficients)]


def _distributed_normalize(
    x: torch.Tensor,
    group: dist.ProcessGroup,
    eps: float = 1e-7,
) -> torch.Tensor:
    norm_sq = (x * x).sum()
    dist.all_reduce(norm_sq, op=dist.ReduceOp.SUM, group=group)
    return x / torch.sqrt(norm_sq).clamp_min(eps)


def _newton_schulz_step(
    x: torch.Tensor,
    a: float,
    b: float,
    c: float,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    gram = x @ x.mT
    dist.all_reduce(gram, op=dist.ReduceOp.SUM, group=group)
    update = torch.addmm(gram, gram, gram, alpha=c, beta=b)
    return torch.addmm(x, update, x, alpha=1.0, beta=a)


@torch.no_grad()
def solution(
    x: torch.Tensor,
    steps: int = 5,
    coefficient_type: str = "quintic",
    partition_dim: int = 1,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    assert x.ndim == 2
    assert coefficient_type in _COEFFICIENTS
    coefficients = _COEFFICIENTS[coefficient_type]
    assert steps % len(coefficients) == 0

    in_dtype = x.dtype

    if partition_dim == 0:
        x_work = x.mT.contiguous()
    elif partition_dim == 1:
        x_work = x
    else:
        raise AssertionError("invalid partition_dim")

    x_work = x_work.to(torch.float32)
    x_work = _distributed_normalize(x_work, group)

    for step in range(steps):
        a, b, c = _coefficient_at(coefficients, step)
        x_work = _newton_schulz_step(x_work, a, b, c, group)

    x_work = x_work.to(in_dtype)
    if partition_dim == 0:
        return x_work.mT.contiguous()
    return x_work.contiguous()
from typing import Optional

import torch
import torch.distributed as dist
import torch.nn.functional as F


def _gather_boundaries(
    x: torch.Tensor,
    boundary: int,
    group: dist.ProcessGroup,
) -> list[torch.Tensor]:
    if boundary == 0:
        empty = x[:, :, :0, :]
        return [torch.stack([empty, empty], dim=0)]

    local = torch.stack([x[:, :, :boundary, :], x[:, :, -boundary:, :]], dim=0)
    gathered = [torch.empty_like(local) for _ in range(dist.get_world_size(group))]
    dist.all_gather(gathered, local.contiguous(), group=group)
    return gathered


@torch.no_grad()
def solution(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: Optional[torch.Tensor],
    stride: int = 1,
    padding: int = 1,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    rank = dist.get_rank(group)
    boundary = int(padding)

    if boundary == 0 or world_size == 1:
        return F.conv2d(x, weight, bias, stride=stride, padding=padding)

    boundaries = _gather_boundaries(x, boundary, group)
    pieces = []
    if rank == 0:
        pieces.append(x.new_zeros(*x.shape[:2], boundary, x.shape[-1]))
    else:
        pieces.append(boundaries[rank - 1][1])
    pieces.append(x)
    if rank == world_size - 1:
        pieces.append(x.new_zeros(*x.shape[:2], boundary, x.shape[-1]))
    else:
        pieces.append(boundaries[rank + 1][0])

    padded_x = torch.cat(pieces, dim=2)
    return F.conv2d(padded_x, weight, bias, stride=stride, padding=(0, padding))