"""
Order-sensitive virtual collective: out = sum_r GELU(x_r), same shape as one rank slice (B, D).

Note: this is not equal to GELU(sum_r x_r).
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

from kernelbench.distributed_collectives import default_device

R = 4
B = 14
D = 40


class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.gelu(x).sum(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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor gelu_sum_forward(torch::Tensor x);
"""

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

__device__ __forceinline__ float gelu_exact_f32(float v) {
    // PyTorch F.gelu default (approximate='none'): 0.5*x*(1+erf(x/sqrt(2))).
    return 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));
}

__global__ void gelu_sum_kernel_4(const float* __restrict__ x, float* __restrict__ y) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < 560) {  // B*D = 14*40 for this benchmark
        float s = 0.0f;
        s += gelu_exact_f32(x[idx]);
        s += gelu_exact_f32(x[560 + idx]);
        s += gelu_exact_f32(x[1120 + idx]);
        s += gelu_exact_f32(x[1680 + idx]);
        y[idx] = s;
    }
}

__global__ void gelu_sum_kernel_generic(const float* __restrict__ x, float* __restrict__ y,
                                        int R, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float s = 0.0f;
        for (int r = 0; r < R; ++r) {
            s += gelu_exact_f32(x[r * N + idx]);
        }
        y[idx] = s;
    }
}

torch::Tensor gelu_sum_forward(torch::Tensor x) {
    const int64_t R = x.size(0);
    const int64_t B = x.size(1);
    const int64_t D = x.size(2);
    auto y = torch::empty({B, D}, x.options());
    const int N = (int)(B * D);
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (R == 4 && N == 560 && x.scalar_type() == at::ScalarType::Float) {
        gelu_sum_kernel_4<<<3, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>());
    } else {
        int blocks = (N + threads - 1) / threads;
        gelu_sum_kernel_generic<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), (int)R, N);
    }
    return y;
}
"""

_ext = load_inline(
    name="gelu_sum_r4_b14_d40_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["gelu_sum_forward"],
    extra_cuda_cflags=["-O3"],
    extra_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.gelu_sum_forward(x)