import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void temp_softmax_channel_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int N, int C, int H, int W, float temp) {
    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;
        float* sX = sh;
        float* sZ = sh + C;
        float* sExp = sh + 2 * C;
        float* sRed = sh + 3 * C;
        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;
            float z = xv * scale[c] + bias[c];
            sZ[c] = z / temp;
        }
        __syncthreads();
        float local_max = -INFINITY;
        for (int c = tid; c < C; c += stride) {
            float z = sZ[c];
            if (z > local_max) local_max = z;
        }
        sRed[tid] = local_max;
        __syncthreads();
        for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
            if (tid < s && sRed[tid + s] > sRed[tid]) sRed[tid] = sRed[tid + s];
            __syncthreads();
        }
        float m = sRed[0];
        float local_sum = 0.0f;
        for (int c = tid; c < C; c += stride) {
            float e = expf(sZ[c] - m);
            sExp[c] = e;
            local_sum += e;
        }
        sRed[tid] = local_sum;
        __syncthreads();
        for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
            if (tid < s) sRed[tid] += sRed[tid + s];
            __syncthreads();
        }
        float denom = sRed[0];
        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 g = sExp[c] / denom;
            y[idx] = g * sX[c];
        }
        __syncthreads();
    }
}

torch::Tensor temp_softmax_channel_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor temp) {
    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);
    float t = temp.item<float>();
    int block = C <= 1024 ? C : 1024;
    int grid = (N * H * W);
    if (grid > 65535) grid = 65535;
    size_t shmem = (size_t)(3 * C + block) * sizeof(float);
    temp_softmax_channel_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, t);
    return y;
}
"""

cpp_source = """
torch::Tensor temp_softmax_channel_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor temp);
"""

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

class ModelNew(torch.nn.Module):
    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, temp: float):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("scale", scale)
        self.register_buffer("bias", bias)
        self.register_buffer("temp", torch.tensor(float(temp), dtype=torch.float32))

    def forward(self, x: torch.Tensor):
        return self.ops.temp_softmax_channel_affine_cuda(x, self.scale, self.bias, self.temp)
