"""
Virtual ranks: input packs R per-rank tensors along dim 0. Output is the mean over ranks (B, D).

No distributed runtime required; the reference always reduces along dim 0.
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

R = 4
B = 16
D = 48


class Model(nn.Module):
    """x shape (R, B, D) -> mean over R -> (B, D)."""

    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x.mean(dim=0)


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.h>
#include <cuda_runtime.h>

__global__ void mean_reduce_scalar_kernel(const float* __restrict__ input, float* __restrict__ output, 
                                          int R, int B, int D) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = B * D;
    
    if (idx < total) {
        const int b = idx / D;
        const int d = idx % D;
        
        float sum = 0.0f;
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            sum += input[r * B * D + b * D + d];
        }
        output[b * D + d] = sum / R;
    }
}

__global__ void mean_reduce_float4_kernel(const float* __restrict__ input, float* __restrict__ output, 
                                          int R, int B, int D4) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = B * D4;
    
    if (idx < total) {
        const int b = idx / D4;
        const int d4 = idx % D4;
        
        float4 sum = {0.0f, 0.0f, 0.0f, 0.0f};
        
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            const float4 val = reinterpret_cast<const float4*>(input)[r * B * D4 + b * D4 + d4];
            sum.x += val.x;
            sum.y += val.y;
            sum.z += val.z;
            sum.w += val.w;
        }
        
        float invR = 1.0f / R;
        float4* out4 = reinterpret_cast<float4*>(output);
        out4[b * D4 + d4] = {sum.x * invR, sum.y * invR, sum.z * invR, sum.w * invR};
    }
}

torch::Tensor mean_reduce_cuda(torch::Tensor x) {
    const int R = x.size(0);
    const int B = x.size(1);
    const int D = x.size(2);
    
    auto output = torch::empty({B, D}, x.options());
    
    const int threads = 256;
    
    if (D % 4 == 0) {
        const int D4 = D / 4;
        const int total = B * D4;
        const int blocks = (total + threads - 1) / threads;
        
        mean_reduce_float4_kernel<<<blocks, threads>>>(
            x.data_ptr<float>(), output.data_ptr<float>(), R, B, D4
        );
    } else {
        const int total = B * D;
        const int blocks = (total + threads - 1) / threads;
        
        mean_reduce_scalar_kernel<<<blocks, threads>>>(
            x.data_ptr<float>(), output.data_ptr<float>(), R, B, D
        );
    }
    
    return output;
}
"""

cpp_source = "torch::Tensor mean_reduce_cuda(torch::Tensor x);"

mean_extension = load_inline(
    name='mean_reduce_v3',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['mean_reduce_cuda'],
    verbose=True
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return mean_extension.mean_reduce_cuda(x)