import torch
from torch.utils.cpp_extension import load_inline

source = """
#include <torch/extension.h>

__global__ void group_softmax_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int N, int C, int H, int W, int group) {
    int tid = threadIdx.x;
    int stride = blockDim.x;
    int total = N * H * W;
    extern __shared__ float sh[];
    for (int pos = blockIdx.x; pos < total; pos += gridDim.x) {
        int n = pos / (H * W);
        int hw = pos % (H * W);
        int h = hw / W;
        int w = hw % W;
        int G = (C + group - 1) / group;
        float* sX = sh;
        float* sZ = sh + C;
        float* sExp = sh + 2 * C;
        float* sM = sh + 3 * C;
        float* sDen = sM + G;
        for (int c = tid; c < C; c += stride) {
            long long idx = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W + w;
            float xv = x[idx];
            sX[c] = xv;
            sZ[c] = xv * scale[c] + bias[c];
        }
        __syncthreads();
        if (tid == 0) {
            for (int gi = 0; gi < G; ++gi) {
                int s = gi * group;
                int e = s + group; if (e > C) e = C;
                float m = -INFINITY;
                for (int c = s; c < e; ++c) {
                    float z = sZ[c];
                    if (z > m) m = z;
                }
                sM[gi] = m;
            }
        }
        __syncthreads();
        for (int c = tid; c < C; c += stride) {
            int gi = c / group;
            float e = expf(sZ[c] - sM[gi]);
            sExp[c] = e;
        }
        __syncthreads();
        if (tid == 0) {
            for (int gi = 0; gi < G; ++gi) {
                int s = gi * group;
                int e = s + group; if (e > C) e = C;
                float den = 0.0f;
                for (int c = s; c < e; ++c) den += sExp[c];
                sDen[gi] = den;
            }
        }
        __syncthreads();
        for (int c = tid; c < C; c += stride) {
            int gi = c / group;
            long long idx = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W + w;
            float g = sExp[c] / sDen[gi];
            y[idx] = g * sX[c];
        }
        __syncthreads();
    }
}

torch::Tensor group_softmax_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor group) {
    auto xc = x.contiguous();
    auto sc = scale.contiguous();
    auto bc = bias.contiguous();
    auto y = torch::empty_like(xc);
    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int H = (int)xc.size(2);
    int W = (int)xc.size(3);
    int g = group.item<int>();
    int block = C <= 1024 ? C : 1024;
    int grid = (N * H * W);
    if (grid > 65535) grid = 65535;
    int G = (C + g - 1) / g;
    size_t shmem = (size_t)(3 * C + G) * sizeof(float) + (size_t)block * sizeof(float);
    group_softmax_affine_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, g);
    return y;
}
"""

cpp_source = """
torch::Tensor group_softmax_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor group);
"""

ops = load_inline(
    name="group_softmax_affine",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["group_softmax_affine_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, group: int):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("scale", scale)
        self.register_buffer("bias", bias)
        self.register_buffer("group", torch.tensor(int(group), dtype=torch.int32))

    def forward(self, x: torch.Tensor):
        return self.ops.group_softmax_affine_cuda(x, self.scale, self.bias, self.group)
