togethercomputer/ParallelKernelBench_Problems huggingface dataset:
# 71_hyena_conv1d_boundary_exchange
from typing import Optional, Tuple

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


def _zigzag_get_overlapping_patches(
    data: torch.Tensor,
    seq_dim: int,
    overlap_size: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    shape = list(data.shape)
    shape[seq_dim : seq_dim + 1] = [2, data.shape[seq_dim] // 2]
    chunks = data.reshape(shape)

    order = list(range(chunks.dim()))
    order.insert(0, order.pop(seq_dim))
    chunks = chunks.permute(order)

    chunk_len = chunks.shape[seq_dim + 1]
    overlaps = chunks.narrow(seq_dim + 1, chunk_len - overlap_size, overlap_size)
    return overlaps[0], overlaps[1]


@torch.no_grad()
def solution(
    x: torch.Tensor,
    weight: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    group_ranks = dist.get_process_group_ranks(group)
    group_rank = dist.get_rank(group)
    group_world_size = len(group_ranks)

    batch, hidden, local_seq = x.shape
    chunk_len = local_seq // 2
    pad_size = weight.shape[-1] - 1
    chunk_a, chunk_b = _zigzag_get_overlapping_patches(x, 2, pad_size)

    ops = []
    recv_prev_a = None
    recv_next_b = None

    if group_rank > 0:
        recv_prev_a = torch.empty_like(chunk_a)
        ops.append(
            dist.P2POp(dist.irecv, recv_prev_a, group_ranks[group_rank - 1], group)
        )
    if group_rank < group_world_size - 1:
        ops.append(
            dist.P2POp(
                dist.isend,
                chunk_a.contiguous(),
                group_ranks[group_rank + 1],
                group,
            )
        )

    if group_rank < group_world_size - 1:
        recv_next_b = torch.empty_like(chunk_b)
        ops.append(
            dist.P2POp(dist.irecv, recv_next_b, group_ranks[group_rank + 1], group)
        )
    if group_rank > 0:
        ops.append(
            dist.P2POp(
                dist.isend,
                chunk_b.contiguous(),
                group_ranks[group_rank - 1],
                group,
            )
        )

    for request in dist.batch_isend_irecv(ops):
        request.wait()

    if recv_prev_a is None:
        recv_prev_a = torch.zeros_like(chunk_a)
    if recv_next_b is None:
        recv_next_b = chunk_a.clone().contiguous()

    # Move the two zigzag chunks into batch so both use the same grouped conv.
    x_chunks = x.reshape(batch, hidden, 2, chunk_len).permute(2, 0, 1, 3)
    x_chunks = x_chunks.reshape(2 * batch, hidden, chunk_len)
    padding = torch.cat([recv_prev_a, recv_next_b], dim=0)
    x_padded = torch.cat([padding, x_chunks], dim=-1)

    y = F.conv1d(x_padded, weight, bias=None, stride=1, padding=0, groups=hidden)
    return (
        y.reshape(2, batch, hidden, chunk_len)
        .permute(1, 2, 0, 3)
        .reshape(batch, hidden, local_seq)
    )
# 73_vocab_parallel_cross_entropy_loss
from typing import Optional, Tuple

import torch
import torch.distributed as dist


def _vocab_range(partition_vocab_size: int, rank: int) -> Tuple[int, int]:
    start = rank * partition_vocab_size
    return start, start + partition_vocab_size


@torch.no_grad()
def solution(
    vocab_parallel_logits: torch.Tensor,
    target: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    rank = dist.get_rank(group=group)

    logits_max = torch.max(vocab_parallel_logits, dim=-1).values
    dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=group)
    vocab_parallel_logits.sub_(logits_max.unsqueeze(dim=-1))

    partition_vocab_size = vocab_parallel_logits.shape[-1]
    vocab_start, vocab_end = _vocab_range(partition_vocab_size, rank)
    target_mask = (target < vocab_start) | (target >= vocab_end)
    masked_target = target - vocab_start
    masked_target = masked_target.masked_fill(target_mask, 0)

    logits_2d = vocab_parallel_logits.reshape(-1, partition_vocab_size)
    target_1d = masked_target.reshape(-1)
    row_ids = torch.arange(logits_2d.shape[0], device=logits_2d.device)
    predicted_logits = logits_2d[row_ids, target_1d].clone().reshape_as(target)
    predicted_logits = predicted_logits.masked_fill(target_mask, 0.0)
    dist.all_reduce(predicted_logits, op=dist.ReduceOp.SUM, group=group)

    exp_logits = torch.exp(vocab_parallel_logits)
    sum_exp_logits = exp_logits.sum(dim=-1)
    dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=group)

    return torch.log(sum_exp_logits) - predicted_logits
# 80_dinov2_distributed_sinkhorn_knopp
from typing import Optional

import torch
import torch.distributed as dist


@torch.no_grad()
def solution(
    teacher_output: torch.Tensor,
    teacher_temp: float,
    n_masked_patches_tensor: torch.Tensor,
    n_iterations: int = 3,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    q = torch.exp(teacher_output.float() / teacher_temp).T
    total_batch = n_masked_patches_tensor.to(device=q.device, dtype=q.dtype).clone()
    dist.all_reduce(total_batch, group=group)

    num_prototypes = q.shape[0]
    total_mass = q.sum()
    dist.all_reduce(total_mass, group=group)
    q /= total_mass

    for _ in range(n_iterations):
        row_sum = q.sum(dim=1, keepdim=True)
        dist.all_reduce(row_sum, group=group)
        q /= row_sum
        q /= num_prototypes

        q /= q.sum(dim=0, keepdim=True)
        q /= total_batch

    q *= total_batch
    return q.T.contiguous()
# 81_sam3_allgather_iou_suppression
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
# 86_tp_muon_orthogonalization
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()
# 87_conv2d_boundary_exchange
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))