willychan21/ParallelKernelBench_Problems

1from typing import List, Optional, Tuple2 3import torch4import torch.distributed as dist5import torch.nn.functional as F6 7 8def _index_undot(index: int, loop_size: List[int]) -> List[int]:9    out: List[int] = []10    for size in reversed(loop_size):11        out.append(index % size)12        index //= size13    return list(reversed(out))14 15 16def _index_dot(index: List[int], loop_size: List[int]) -> int:17    value = 018    for dim, size in zip(index, loop_size):19        value = value * size + dim20    return value21 22 23def _split_tiles(24    tile_numels: List[int],25    group: Optional[dist.ProcessGroup],26) -> Tuple[List[int], List[int]]:27    if group is None:28        tile_indices = list(range(len(tile_numels)))29        return tile_indices, tile_indices30 31    world_size = dist.get_world_size(group=group)32    rank = dist.get_rank(group=group)33    sorted_tiles = sorted(34        range(len(tile_numels)),35        key=lambda idx: tile_numels[idx],36        reverse=True,37    )38    per_rank = [sorted_tiles[r::world_size] for r in range(world_size)]39    global_order = [idx for shard in per_rank for idx in shard]40    return per_rank[rank], global_order41 42 43def _gather_tiles(44    tiles: List[torch.Tensor],45    global_order: List[int],46    template: torch.Tensor,47    group: Optional[dist.ProcessGroup],48) -> List[torch.Tensor]:49    if group is None:50        return tiles51 52    world_size = dist.get_world_size(group=group)53    local_shapes = [tuple(tile.shape) for tile in tiles]54    all_shapes: List[List[Tuple[int, ...]]] = [[] for _ in range(world_size)]55    dist.all_gather_object(all_shapes, local_shapes, group=group)56 57    local_flat = (58        torch.cat([tile.reshape(-1).contiguous() for tile in tiles], dim=0)59        if tiles60        else template.new_empty(0)61    )62    local_size = int(local_flat.numel())63    rank_sizes: List[int] = []64    for shapes in all_shapes:65        total = 066        for shape in shapes:67            numel = 168            for size in shape:69                numel *= size70            total += numel71        rank_sizes.append(total)72    send = local_flat.repeat(world_size)73    recv = template.new_empty(sum(rank_sizes))74    dist.all_to_all_single(75        recv,76        send,77        output_split_sizes=rank_sizes,78        input_split_sizes=[local_size] * world_size,79        group=group,80    )81 82    gathered: List[torch.Tensor] = []83    offset = 084    for shapes, total in zip(all_shapes, rank_sizes):85        rank_buf = recv[offset : offset + total]86        rank_offset = 087        for shape in shapes:88            numel = 189            for size in shape:90                numel *= size91            gathered.append(rank_buf[rank_offset : rank_offset + numel].view(shape))92            rank_offset += numel93        offset += total94 95    by_index = {tile_idx: tile for tile_idx, tile in zip(global_order, gathered)}96    return [by_index[idx] for idx in sorted(by_index)]97 98 99def _blend_t(prev: torch.Tensor, cur: torch.Tensor, extent: int) -> torch.Tensor:100    extent = min(prev.shape[2], cur.shape[2], extent)101    for idx in range(extent):102        ratio = idx / extent103        cur[:, :, idx] = prev[:, :, -extent + idx] * (1.0 - ratio) + cur[:, :, idx] * ratio104    return cur105 106 107def _blend_h(prev: torch.Tensor, cur: torch.Tensor, extent: int) -> torch.Tensor:108    extent = min(prev.shape[3], cur.shape[3], extent)109    for idx in range(extent):110        ratio = idx / extent111        cur[:, :, :, idx] = prev[:, :, :, -extent + idx] * (1.0 - ratio) + cur[:, :, :, idx] * ratio112    return cur113 114 115def _blend_w(prev: torch.Tensor, cur: torch.Tensor, extent: int) -> torch.Tensor:116    extent = min(prev.shape[4], cur.shape[4], extent)117    for idx in range(extent):118        ratio = idx / extent119        cur[:, :, :, :, idx] = prev[:, :, :, :, -extent + idx] * (1.0 - ratio) + cur[:, :, :, :, idx] * ratio120    return cur121 122 123def _decode_tile(tile: torch.Tensor, spatial_upsample: int, temporal_upsample: int) -> torch.Tensor:124    decoded = F.interpolate(125        tile.float(),126        scale_factor=(temporal_upsample, spatial_upsample, spatial_upsample),127        mode="trilinear",128        align_corners=False,129    )130    if decoded.shape[1] < 3:131        repeats = (3 + decoded.shape[1] - 1) // decoded.shape[1]132        decoded = decoded.repeat(1, repeats, 1, 1, 1)133    return decoded[:, :3].to(torch.bfloat16)134 135 136@torch.no_grad()137def solution(138    z: torch.Tensor,139    tile_latent_min_length: int,140    tile_latent_min_height: int,141    tile_latent_min_width: int,142    spatial_tile_overlap_factor: float,143    temporal_tile_overlap_factor: float,144    spatial_upsample: int,145    temporal_upsample: int,146    sr_ratio: int = 1,147    first_frame_as_image: bool = False,148    group: Optional[dist.ProcessGroup] = None,149) -> torch.Tensor:150    if dist.is_available() and dist.is_initialized():151        group = group or dist.group.WORLD152    else:153        group = None154    tile_latent_min_length = tile_latent_min_length + int(first_frame_as_image)155    spatial_upsample = spatial_upsample * sr_ratio156    stride_h = int(tile_latent_min_height * (1.0 - spatial_tile_overlap_factor))157    stride_w = int(tile_latent_min_width * (1.0 - spatial_tile_overlap_factor))158    stride_t = int(tile_latent_min_length * (1.0 - temporal_tile_overlap_factor))159    if min(stride_t, stride_h, stride_w) <= 0:160        raise ValueError("tile overlap factors must leave a positive stride")161 162    real_t = tile_latent_min_length * temporal_upsample163    real_h = tile_latent_min_height * spatial_upsample164    real_w = tile_latent_min_width * spatial_upsample165    blend_t = int(real_t * temporal_tile_overlap_factor)166    blend_h = int(real_h * spatial_tile_overlap_factor)167    blend_w = int(real_w * spatial_tile_overlap_factor)168    keep_t = real_t - blend_t169    keep_h = real_h - blend_h170    keep_w = real_w - blend_w171 172    tiles_t = (z.shape[2] + stride_t - 1) // stride_t173    tiles_h = (z.shape[3] + stride_h - 1) // stride_h174    tiles_w = (z.shape[4] + stride_w - 1) // stride_w175    loop_size = [tiles_t, tiles_h, tiles_w]176    total_tiles = tiles_t * tiles_h * tiles_w177 178    latent_tiles: List[torch.Tensor] = []179    tile_numels: List[int] = []180    for tile_idx in range(total_tiles):181        t_idx, h_idx, w_idx = _index_undot(tile_idx, loop_size)182        t0 = t_idx * stride_t183        h0 = h_idx * stride_h184        w0 = w_idx * stride_w185        tile = z[186            :,187            :,188            t0 : t0 + tile_latent_min_length,189            h0 : h0 + tile_latent_min_height,190            w0 : w0 + tile_latent_min_width,191        ]192        latent_tiles.append(tile)193        tile_numels.append(int(tile.numel()))194 195    local_indices, global_order = _split_tiles(tile_numels, group)196    decoded = [197        _decode_tile(latent_tiles[idx], spatial_upsample, temporal_upsample)198        for idx in local_indices199    ]200    template = decoded[0] if decoded else _decode_tile(latent_tiles[0], spatial_upsample, temporal_upsample)201    decoded_all = _gather_tiles(decoded, global_order, template, group)202 203    blended: List[torch.Tensor] = []204    for tile_idx in local_indices:205        t_idx, h_idx, w_idx = _index_undot(tile_idx, loop_size)206        tile = decoded_all[tile_idx].clone()207        if t_idx > 0:208            prev_idx = _index_dot([t_idx - 1, h_idx, w_idx], loop_size)209            tile = _blend_t(decoded_all[prev_idx], tile, blend_t)210        if h_idx > 0:211            prev_idx = _index_dot([t_idx, h_idx - 1, w_idx], loop_size)212            tile = _blend_h(decoded_all[prev_idx], tile, blend_h)213        if w_idx > 0:214            prev_idx = _index_dot([t_idx, h_idx, w_idx - 1], loop_size)215            tile = _blend_w(decoded_all[prev_idx], tile, blend_w)216        blended.append(tile[:, :, :keep_t, :keep_h, :keep_w].contiguous())217 218    blended_all = _gather_tiles(blended, global_order, template, group)219    frames_t: List[torch.Tensor] = []220    for t_idx in range(tiles_t):221        rows: List[torch.Tensor] = []222        for h_idx in range(tiles_h):223            row: List[torch.Tensor] = []224            for w_idx in range(tiles_w):225                row.append(blended_all[_index_dot([t_idx, h_idx, w_idx], loop_size)])226            rows.append(torch.cat(row, dim=4))227        frames_t.append(torch.cat(rows, dim=3))228    return torch.cat(frames_t, dim=2)229