willychan21/ParallelKernelBench_Problems
ParallelKernelBench (benchmark) Reference problems for ParallelKernelBench: a benchmark for LLM-generated multi-GPU CUDA kernels. This dataset contains 87 reference implementations in reference/ and the input tensor specification in utils/input_output_tensors.py. Files Path Description data/problems.parquet One row per problem (tabular access) reference/*.py Reference solution() implementations utils/input_output_tensors.py Input/output tensor… See the full description on the dataset page: https://huggingface.co/datasets/willychan21/ParallelKernelBench_Problems.
0266
1from typing import List, Optional, Tuple2 3import torch4import torch.distributed as dist5 6 7def solution(8 x: torch.Tensor,9 gather_dim: int,10 group: Optional[dist.ProcessGroup] = None,11) -> torch.Tensor:12 group = group or dist.group.WORLD13 world_size = dist.get_world_size(group)14 if world_size == 1:15 return x.contiguous()16 17 device = x.device18 dtype = x.dtype19 x = x.contiguous()20 x_size = torch.tensor(x.size(), dtype=torch.int64, device=device)21 size_list = [torch.zeros(x_size.size(), dtype=torch.int64, device=device) for _ in range(world_size)]22 dist.all_gather(size_list, x_size, group=group)23 tensor_list = [24 torch.empty(torch.Size(size_list[i].tolist()), dtype=dtype, device=device)25 for i in range(world_size)26 ]27 dist.all_gather(tensor_list, x, group=group)28 return torch.cat(tensor_list, dim=gather_dim).contiguous()29