willychan21/ParallelKernelBench_Problems

1from typing import List, Optional, Tuple2 3import numpy as np4import torch5import torch.distributed as dist6 7 8def _sample_one_hop_csc_dist(9    input_nodes: torch.Tensor,10    k: int,11    colptr: torch.Tensor,12    row: torch.Tensor,13    replace: bool = False,14) -> Tuple[torch.Tensor, torch.Tensor, List[int]]:15    n = input_nodes.numel()16    sampled_nodes = []17    sampled_edges = []18    cumsum = [n]19 20    for i in range(n):21        v = int(input_nodes[i].item())22        start = int(colptr[v].item())23        end = int(colptr[v + 1].item())24        deg = end - start25        take = min(k, deg) if k >= 0 else deg26 27        if take > 0:28            if replace:29                perm = torch.randint(deg, (take,), device=input_nodes.device)30            else:31                perm = torch.randperm(deg, device=input_nodes.device)[:take]32            sampled_nodes.append(row[start:end].index_select(0, perm))33            sampled_edges.append(torch.arange(start, end, device=input_nodes.device).index_select(0, perm))34 35        cumsum.append(cumsum[-1] + take)36 37    nbr_tensor = (38        torch.cat(sampled_nodes)39        if sampled_nodes40        else torch.empty(0, dtype=torch.long, device=input_nodes.device)41    )42    eid_tensor = (43        torch.cat(sampled_edges)44        if sampled_edges45        else torch.empty(0, dtype=torch.long, device=input_nodes.device)46    )47    return torch.cat([input_nodes, nbr_tensor]), eid_tensor, cumsum48 49 50def _remove_duplicates(out_node: torch.Tensor, node: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:51    num_nodes = node.numel()52    node_combined = torch.cat([node, out_node])53    _, idx = np.unique(node_combined.cpu().numpy(), return_index=True)54    idx = torch.from_numpy(idx).to(node.device).sort().values55    node = node_combined[idx]56    src = node[num_nodes:]57    return src, node58 59 60def _relabel_neighborhood(61    node: torch.Tensor,62    dst_with_dupl: torch.Tensor,63    node_with_dupl: torch.Tensor,64) -> Tuple[torch.Tensor, torch.Tensor]:65    if node_with_dupl.numel() == 0:66        return node.new_empty(0), node.new_empty(0)67 68    assoc = torch.full(69        (int(node.max().item()) + 1,),70        -1,71        dtype=torch.long,72        device=node.device,73    )74    assoc[node] = torch.arange(node.numel(), device=node.device)75    row = assoc[node_with_dupl]76    col = assoc[dst_with_dupl]77    return row, col78 79 80def _exchange_nodes(81    send_nodes_list: List[torch.Tensor],82    group: dist.ProcessGroup,83) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:84    world_size = dist.get_world_size(group)85    device = send_nodes_list[0].device86    send_counts = torch.tensor([x.numel() for x in send_nodes_list], dtype=torch.long, device=device)87    recv_counts = torch.empty_like(send_counts)88    dist.all_to_all_single(recv_counts, send_counts, group=group)89 90    send_nodes = torch.cat(send_nodes_list) if send_nodes_list else torch.empty(0, dtype=torch.long, device=device)91    recv_nodes = torch.empty(int(recv_counts.sum().item()), dtype=torch.long, device=device)92    dist.all_to_all_single(93        recv_nodes,94        send_nodes,95        input_split_sizes=send_counts.cpu().tolist(),96        output_split_sizes=recv_counts.cpu().tolist(),97        group=group,98    )99    return recv_nodes, send_counts, recv_counts100 101 102def _exchange_replies(103    sampled_nodes: torch.Tensor,104    sampled_edges: torch.Tensor,105    sampled_counts: torch.Tensor,106    recv_counts: torch.Tensor,107    group: dist.ProcessGroup,108) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:109    world_size = dist.get_world_size(group)110    device = sampled_nodes.device111    recv_splits = recv_counts.cpu().tolist()112    send_node_counts = torch.empty(world_size, dtype=torch.long, device=device)113    offset = 0114    for r, count in enumerate(recv_splits):115        send_node_counts[r] = sampled_counts[offset : offset + count].sum()116        offset += count117 118    reply_node_counts = torch.empty_like(send_node_counts)119    dist.all_to_all_single(reply_node_counts, send_node_counts, group=group)120 121    reply_count_counts = torch.empty_like(recv_counts)122    dist.all_to_all_single(reply_count_counts, recv_counts, group=group)123 124    reply_nodes = torch.empty(int(reply_node_counts.sum().item()), dtype=torch.long, device=device)125    reply_edges = torch.empty_like(reply_nodes)126    reply_counts = torch.empty(int(reply_count_counts.sum().item()), dtype=torch.long, device=device)127 128    dist.all_to_all_single(129        reply_nodes,130        sampled_nodes,131        input_split_sizes=send_node_counts.cpu().tolist(),132        output_split_sizes=reply_node_counts.cpu().tolist(),133        group=group,134    )135    dist.all_to_all_single(136        reply_edges,137        sampled_edges,138        input_split_sizes=send_node_counts.cpu().tolist(),139        output_split_sizes=reply_node_counts.cpu().tolist(),140        group=group,141    )142    dist.all_to_all_single(143        reply_counts,144        sampled_counts,145        input_split_sizes=recv_splits,146        output_split_sizes=reply_count_counts.cpu().tolist(),147        group=group,148    )149    return reply_nodes, reply_edges, reply_counts150 151 152@torch.no_grad()153def solution(154    seed_nodes: torch.Tensor,155    fanouts: List[int],156    local_adj_row_ptr: torch.Tensor,157    local_adj_col: torch.Tensor,158    node_to_rank: torch.Tensor,159    group: Optional[dist.ProcessGroup] = None,160    replace: bool = False,161) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:162    group = group or dist.group.WORLD163    world_size = dist.get_world_size(group)164    device = seed_nodes.device165 166    seed = seed_nodes.to(dtype=torch.long, device=device)167    src = seed.clone()168    node = src.clone()169    node_with_dupl = [seed.new_empty(0)]170    dst_with_dupl = [seed.new_empty(0)]171    edge = [seed.new_empty(0)]172 173    for fanout in fanouts:174        if src.numel() == 0:175            break176 177        partition_ids = node_to_rank[src].to(torch.long)178        partition_orders = torch.empty_like(partition_ids)179        send_nodes_list = []180        send_pos_list = []181        for r in range(world_size):182            pos = (partition_ids == r).nonzero(as_tuple=False).flatten()183            partition_orders[pos] = torch.arange(pos.numel(), dtype=torch.long, device=device)184            send_nodes_list.append(src[pos])185            send_pos_list.append(pos)186 187        recv_nodes, send_counts, recv_counts = _exchange_nodes(send_nodes_list, group)188        node_out, edge_out, cumsum = _sample_one_hop_csc_dist(189            recv_nodes, int(fanout), local_adj_row_ptr, local_adj_col, replace190        )191 192        seed_size = recv_nodes.numel()193        sampled_nodes = node_out[seed_size:]194        sampled_counts = torch.tensor(195            np.subtract(np.array(cumsum[1:]), np.array(cumsum[:-1])),196            dtype=torch.long,197            device=device,198        )199 200        reply_nodes, reply_edges, reply_counts = _exchange_replies(201            sampled_nodes, edge_out, sampled_counts, recv_counts, group202        )203 204        rank_offsets = torch.cat(205            [send_counts.new_zeros(1), torch.cumsum(send_counts, dim=0)[:-1]]206        )207        grouped_index = rank_offsets[partition_ids] + partition_orders208        node_chunks = list(torch.split(reply_nodes, reply_counts.cpu().tolist()))209        edge_chunks = list(torch.split(reply_edges, reply_counts.cpu().tolist()))210 211        ordered_nodes = []212        ordered_edges = []213        ordered_dst = []214        for idx in grouped_index.tolist():215            ordered_nodes.append(node_chunks[idx])216            ordered_edges.append(edge_chunks[idx])217        for dst_node, count in zip(src, reply_counts[grouped_index]):218            ordered_dst.append(dst_node.repeat(int(count.item())))219 220        out_node = torch.cat(ordered_nodes) if ordered_nodes else seed.new_empty(0)221        out_edge = torch.cat(ordered_edges) if ordered_edges else seed.new_empty(0)222        out_dst = torch.cat(ordered_dst) if ordered_dst else seed.new_empty(0)223        if out_node.numel() == 0:224            break225 226        src, node = _remove_duplicates(out_node, node)227        node_with_dupl.append(out_node)228        dst_with_dupl.append(out_dst)229        edge.append(out_edge)230 231    node_dupl = torch.cat(node_with_dupl)232    dst_dupl = torch.cat(dst_with_dupl)233    row, col = _relabel_neighborhood(node, dst_dupl, node_dupl)234    return node, row, col, torch.cat(edge)