import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a convolution, applies Instance Normalization, and divides by a constant.
    """
    def __init__(self, in_channels, out_channels, kernel_size, divide_by):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.instance_norm = nn.InstanceNorm2d(out_channels)
        self.divide_by = divide_by

    def forward(self, x):
        x = self.conv(x)
        x = self.instance_norm(x)
        x = x / self.divide_by
        return x

batch_size = 128
in_channels  = 64  
out_channels = 128  
height = width = 128  
kernel_size = 3
divide_by = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, divide_by]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

torch::Tensor instance_norm_div(torch::Tensor input, double divide_by);
"""

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

#ifndef BLOCK_SIZE
#define BLOCK_SIZE 256
#endif

__global__ void inst_norm_div_vec4_kernel(const float* __restrict__ x,
                                          float* __restrict__ y,
                                          int m4,
                                          float div_by) {
    const int nc = blockIdx.x;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x + ((long long)nc) * ((long long)m4) * 4LL);
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y + ((long long)nc) * ((long long)m4) * 4LL);

    float sum = 0.0f;
    float sumsq = 0.0f;

    for (int i = threadIdx.x; i < m4; i += BLOCK_SIZE) {
        float4 v = x4[i];
        sum += v.x + v.y + v.z + v.w;
        sumsq = fmaf(v.x, v.x, sumsq);
        sumsq = fmaf(v.y, v.y, sumsq);
        sumsq = fmaf(v.z, v.z, sumsq);
        sumsq = fmaf(v.w, v.w, sumsq);
    }

    __shared__ float s_sum[BLOCK_SIZE];
    __shared__ float s_sumsq[BLOCK_SIZE];
    s_sum[threadIdx.x] = sum;
    s_sumsq[threadIdx.x] = sumsq;
    __syncthreads();

    #pragma unroll
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            s_sum[threadIdx.x] += s_sum[threadIdx.x + stride];
            s_sumsq[threadIdx.x] += s_sumsq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    const float inv_m = 1.0f / (float)(m4 * 4);
    const float mean = s_sum[0] * inv_m;
    float var = s_sumsq[0] * inv_m - mean * mean;
    var = var > 0.0f ? var : 0.0f;
    const float scale = rsqrtf(var + 1.0e-5f) / div_by;

    for (int i = threadIdx.x; i < m4; i += BLOCK_SIZE) {
        float4 v = x4[i];
        v.x = (v.x - mean) * scale;
        v.y = (v.y - mean) * scale;
        v.z = (v.z - mean) * scale;
        v.w = (v.w - mean) * scale;
        y4[i] = v;
    }
}

__global__ void inst_norm_div_scalar_kernel(const float* __restrict__ x,
                                            float* __restrict__ y,
                                            int m,
                                            float div_by) {
    const int nc = blockIdx.x;
    const long long base = ((long long)nc) * (long long)m;

    float sum = 0.0f;
    float sumsq = 0.0f;
    for (int i = threadIdx.x; i < m; i += BLOCK_SIZE) {
        float v = x[base + i];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }

    __shared__ float s_sum[BLOCK_SIZE];
    __shared__ float s_sumsq[BLOCK_SIZE];
    s_sum[threadIdx.x] = sum;
    s_sumsq[threadIdx.x] = sumsq;
    __syncthreads();

    #pragma unroll
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            s_sum[threadIdx.x] += s_sum[threadIdx.x + stride];
            s_sumsq[threadIdx.x] += s_sumsq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    const float inv_m = 1.0f / (float)m;
    const float mean = s_sum[0] * inv_m;
    float var = s_sumsq[0] * inv_m - mean * mean;
    var = var > 0.0f ? var : 0.0f;
    const float scale = rsqrtf(var + 1.0e-5f) / div_by;

    for (int i = threadIdx.x; i < m; i += BLOCK_SIZE) {
        y[base + i] = (x[base + i] - mean) * scale;
    }
}

torch::Tensor instance_norm_div(torch::Tensor input, double divide_by) {
    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(input.dim() == 4, "expected NCHW tensor");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");

    auto output = torch::empty_like(input);
    const int n = (int)input.size(0);
    const int c = (int)input.size(1);
    const int h = (int)input.size(2);
    const int w = (int)input.size(3);
    const int m = h * w;
    const int blocks = n * c;
    const float div = (float)divide_by;

    if (blocks > 0 && m > 0) {
        const float* x = input.data_ptr<float>();
        float* y = output.data_ptr<float>();
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
        if ((m & 3) == 0) {
            inst_norm_div_vec4_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(x, y, m >> 2, div);
        } else {
            inst_norm_div_scalar_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(x, y, m, div);
        }
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return output;
}
"""

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


class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, divide_by):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.instance_norm = nn.InstanceNorm2d(out_channels)
        self.divide_by = float(divide_by)

    def forward(self, x):
        y = self.conv(x)
        return _ext.instance_norm_div(y, self.divide_by)