"""
Virtual all-gather on the last dim: stack R tensors of shape (B, H) into (B, R*H), then Linear.
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

R = 4
B = 12
H = 32
K = 56


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(R * H, K, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        r, b, h = x.shape
        assert r == R and h == H
        parts = [x[i] for i in range(R)]
        gathered = torch.cat(parts, dim=-1)
        return self.proj(gathered)


def get_inputs():
    dev = default_device()
    p = popcorn_pri
    return [torch.randn(p.jitter_int(R), p.jitter_int(B), p.jitter_int(D, align=8), device=dev)]


def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

// Fixed problem: R=4, B=12, H=32, K=56, N=128.
// One CTA per batch row, 8 warps per CTA.  Each warp computes one output
// column at a time using a warp reduction over the 128 input features.
__global__ __launch_bounds__(256, 2)
void vag_linear_kernel(const float* __restrict__ x,
                       const float* __restrict__ w,
                       const float* __restrict__ bias,
                       float* __restrict__ y) {
    const int b = blockIdx.x;          // 0..11
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;         // 0..7

    // x is laid out as x[r,b,h].  The virtually gathered vector is
    // [x[0,b,0:32], x[1,b,0:32], x[2,b,0:32], x[3,b,0:32]].
    #pragma unroll
    for (int k0 = 0; k0 < 7; ++k0) {
        const int k = k0 * 8 + warp;   // 0..55
        float sum = 0.0f;
        const float* __restrict__ wk = w + k * 128;

        // 128 features, four per lane.  Written explicitly so the compiler
        // can generate a compact straight-line dot product.
        int n = lane;
        int r = n >> 5;
        int h = n & 31;
        sum = fmaf(x[(r * 12 + b) * 32 + h], wk[n], sum);
        n += 32;
        r = n >> 5;
        h = n & 31;
        sum = fmaf(x[(r * 12 + b) * 32 + h], wk[n], sum);
        n += 32;
        r = n >> 5;
        h = n & 31;
        sum = fmaf(x[(r * 12 + b) * 32 + h], wk[n], sum);
        n += 32;
        r = n >> 5;
        h = n & 31;
        sum = fmaf(x[(r * 12 + b) * 32 + h], wk[n], sum);

        // Warp reduction.
        sum += __shfl_down_sync(0xffffffff, sum, 16);
        sum += __shfl_down_sync(0xffffffff, sum, 8);
        sum += __shfl_down_sync(0xffffffff, sum, 4);
        sum += __shfl_down_sync(0xffffffff, sum, 2);
        sum += __shfl_down_sync(0xffffffff, sum, 1);

        if (lane == 0) {
            y[b * 56 + k] = sum + bias[k];
        }
    }
}

torch::Tensor vag_linear_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    auto y = torch::empty({12, 56}, x.options());
    const dim3 grid(12);
    const dim3 block(256);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    vag_linear_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>());
    return y;
}
'''

_ext = load_inline(
    name='vag_linear_ext',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['vag_linear_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(4 * 32, 56, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.vag_linear_forward(x, self.proj.weight, self.proj.bias)