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

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float gelu(float x) {
    return x * 0.5f * (1.0f + erff(x * 0.7071067811865475f));
}

__global__ void gelu_sum_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * D;
    if (idx >= total) return;
    float sum = 0.0f;
    for (int r = 0; r < R; r++) {
        float val = x[r * total + idx];
        sum += gelu(val);
    }
    out[idx] = sum;
}

torch::Tensor gelu_sum_cuda(torch::Tensor x) {
    int R = x.size(0);
    int B = x.size(1);
    int D = x.size(2);
    auto out = torch::empty({B, D}, x.options());
    int total = B * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    gelu_sum_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), R, B, D);
    return out;
}
"""

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

module = load_inline(
    name="gelu_sum",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["gelu_sum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"]
)

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.module.gelu_sum_cuda(x)