"""
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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor vag_linear_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
'''

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

constexpr int R_ = 4;
constexpr int B_ = 12;
constexpr int H_ = 32;
constexpr int K_ = 56;
constexpr int N_ = R_ * H_;
constexpr int WARPS_PER_BLOCK = 7;
constexpr int KTILE = 8;

__global__ __launch_bounds__(WARPS_PER_BLOCK * 32, 4)
void vag_linear_kernel8(const float* __restrict__ x,
                        const float* __restrict__ w,
                        const float* __restrict__ bias,
                        float* __restrict__ y) {
    int b = blockIdx.x;
    int warp = threadIdx.x >> 5;
    int lane = threadIdx.x & 31;
    int k0 = warp * KTILE;

    float s0=0.0f, s1=0.0f, s2=0.0f, s3=0.0f;
    float s4=0.0f, s5=0.0f, s6=0.0f, s7=0.0f;
#pragma unroll
    for (int rr = 0; rr < R_; ++rr) {
        int n = rr * H_ + lane;
        float xv = x[(rr * B_ + b) * H_ + lane];
        const float* wp = w + k0 * N_ + n;
        s0 = fmaf(xv, wp[0 * N_], s0);
        s1 = fmaf(xv, wp[1 * N_], s1);
        s2 = fmaf(xv, wp[2 * N_], s2);
        s3 = fmaf(xv, wp[3 * N_], s3);
        s4 = fmaf(xv, wp[4 * N_], s4);
        s5 = fmaf(xv, wp[5 * N_], s5);
        s6 = fmaf(xv, wp[6 * N_], s6);
        s7 = fmaf(xv, wp[7 * N_], s7);
    }
#pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        s0 += __shfl_down_sync(0xffffffff, s0, off);
        s1 += __shfl_down_sync(0xffffffff, s1, off);
        s2 += __shfl_down_sync(0xffffffff, s2, off);
        s3 += __shfl_down_sync(0xffffffff, s3, off);
        s4 += __shfl_down_sync(0xffffffff, s4, off);
        s5 += __shfl_down_sync(0xffffffff, s5, off);
        s6 += __shfl_down_sync(0xffffffff, s6, off);
        s7 += __shfl_down_sync(0xffffffff, s7, off);
    }
    if (lane == 0) {
        float* yp = y + b * K_ + k0;
        const float* bp = bias + k0;
        yp[0] = s0 + bp[0];
        yp[1] = s1 + bp[1];
        yp[2] = s2 + bp[2];
        yp[3] = s3 + bp[3];
        yp[4] = s4 + bp[4];
        yp[5] = s5 + bp[5];
        yp[6] = s6 + bp[6];
        yp[7] = s7 + bp[7];
    }
}

torch::Tensor vag_linear_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    auto y = torch::empty({B_, K_}, x.options());
    vag_linear_kernel8<<<B_, WARPS_PER_BLOCK * 32>>>(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_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['vag_linear_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    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)