"""
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_runtime.h>

__device__ __forceinline__ float erf_approx_fast(float x) {
    float s = copysignf(1.0f, x);
    float ax = fabsf(x);
    float t = 1.0f / (1.0f + 0.3275911f * ax);
    float y = (((((1.061405429f * t - 1.453152027f) * t) + 1.421413741f) * t - 0.284496736f) * t + 0.254829592f) * t;
    y = 1.0f - y * expf(-ax * ax);
    return s * y;
}

__device__ __forceinline__ float gelu_fast(float a) {
    return 0.5f * a * (1.0f + erf_approx_fast(a * 0.7071067811865475f));
}

__global__ void gelu_sum_kernel(const float* __restrict__ x, float* __restrict__ out) {
    constexpr int N = 560;
    constexpr int TILE = 128;
    __shared__ float sm[4 * TILE];
    int t = threadIdx.x;
    int local = t & (TILE - 1);
    int r = t >> 7;
    int idx = blockIdx.x * TILE + local;
    float val = 0.0f;
    if (idx < N) val = gelu_fast(x[r * N + idx]);
    sm[t] = val;
    __syncthreads();
    if (t < TILE && idx < N) {
        out[idx] = sm[local] + sm[TILE + local] + sm[2 * TILE + local] + sm[3 * TILE + local];
    }
}

torch::Tensor gelu_sum_forward(torch::Tensor x) {
    auto out = torch::empty({14, 40}, x.options());
    gelu_sum_kernel<<<5, 512, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
'''

_ext = load_inline(
    name='gelu_sum_ext_final',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['gelu_sum_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    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)