"""
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_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void fused_gather_linear_kernel(const float* x, const float* weight, const float* bias, float* output, int R, int B, int H, int K) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    if (b >= B || k >= K) return;
    
    float sum = bias[k];
    for (int r = 0; r < R; r++) {
        for (int h = 0; h < H; h++) {
            int x_idx = r * B * H + b * H + h;
            int w_idx = k * (R * H) + r * H + h;
            sum += x[x_idx] * weight[w_idx];
        }
    }
    output[b * K + k] = sum;
}

torch::Tensor fused_gather_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    int R = x.size(0);
    int B = x.size(1);
    int H = x.size(2);
    int K = weight.size(0);
    
    auto output = torch::empty({B, K}, x.options());
    
    dim3 block(16, 16);
    dim3 grid((B + 15) / 16, (K + 15) / 16);
    
    fused_gather_linear_kernel<<<grid, block>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        output.data_ptr<float>(), R, B, H, K
    );
    
    return output;
}
'''

cpp_source = '''
torch::Tensor fused_gather_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
'''

module = load_inline(
    name='fused_gather_linear',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_gather_linear'],
    verbose=True
)

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:
        weight = self.proj.weight
        bias = self.proj.bias
        return module.fused_gather_linear(x, weight, bias)