import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Instance Normalization.
    """
    def __init__(self, num_features: int):
        """
        Initializes the InstanceNorm layer.

        Args:
            num_features (int): Number of features in the input tensor.
        """
        super(Model, self).__init__()
        self.inorm = nn.InstanceNorm2d(num_features=num_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Instance Normalization to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, num_features, height, width).

        Returns:
            torch.Tensor: Output tensor with Instance Normalization applied, same shape as input.
        """
        return self.inorm(x)

batch_size = 112  # heavier workload
features = 64
dim1 = 512
dim2 = 512

def get_inputs():
    x = torch.rand(batch_size, features, dim1, dim2)
    return [x]

def get_init_inputs():
    return [features]

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_forward(torch::Tensor x);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

__inline__ __device__ float warp_sum(float v) {
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}

__inline__ __device__ void block_reduce_pair(float &a, float &b) {
    __shared__ float sh_a[32];
    __shared__ float sh_b[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    a = warp_sum(a);
    b = warp_sum(b);
    if (lane == 0) { sh_a[wid] = a; sh_b[wid] = b; }
    __syncthreads();
    a = (threadIdx.x < (blockDim.x >> 5)) ? sh_a[lane] : 0.0f;
    b = (threadIdx.x < (blockDim.x >> 5)) ? sh_b[lane] : 0.0f;
    if (wid == 0) { a = warp_sum(a); b = warp_sum(b); }
    if (threadIdx.x == 0) { sh_a[0] = a; sh_b[0] = b; }
    __syncthreads();
    a = sh_a[0]; b = sh_b[0];
}

__global__ __launch_bounds__(256, 2)
void inst_norm_kernel(const float* __restrict__ x, float* __restrict__ y,
                      int instances, int spatial) {
    int inst = blockIdx.x;
    int offset = inst * spatial;
    const float* base = x + offset;
    float* out = y + offset;
    float sum = 0.0f, sumsq = 0.0f;
    int n4 = spatial >> 2;
    const float4* __restrict__ b4 = reinterpret_cast<const float4*>(base);
    for (int i = threadIdx.x; i < n4; i += blockDim.x) {
        float4 v = b4[i];
        float s = (v.x + v.y) + (v.z + v.w);
        sum += s;
        sumsq += fmaf(v.x, v.x, fmaf(v.y, v.y, fmaf(v.z, v.z, v.w*v.w)));
    }
    for (int i = (n4 << 2) + threadIdx.x; i < spatial; i += blockDim.x) {
        float v = base[i];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }
    block_reduce_pair(sum, sumsq);
    float mean = sum / (float)spatial;
    float var = sumsq / (float)spatial - mean * mean;
    float invstd = rsqrtf(fmaxf(var, 0.0f) + 1.0e-5f);
    float4* __restrict__ o4 = reinterpret_cast<float4*>(out);
    for (int i = threadIdx.x; i < n4; i += blockDim.x) {
        float4 v = b4[i];
        o4[i] = make_float4((v.x-mean)*invstd, (v.y-mean)*invstd, (v.z-mean)*invstd, (v.w-mean)*invstd);
    }
    for (int i = (n4 << 2) + threadIdx.x; i < spatial; i += blockDim.x) {
        out[i] = (base[i]-mean)*invstd;
    }
}

torch::Tensor instance_norm_forward(torch::Tensor x) {
    CHECK_CUDA(x); CHECK_CONTIGUOUS(x); CHECK_FLOAT(x);
    TORCH_CHECK(x.dim()==4, "expected NCHW 4D tensor");
    auto y = torch::empty_like(x);
    int64_t inst64 = x.size(0)*x.size(1), sp64 = x.size(2)*x.size(3);
    if (inst64 == 0 || sp64 == 0) return y;
    TORCH_CHECK(inst64 <= INT_MAX && sp64 <= INT_MAX && x.numel() <= (int64_t)INT_MAX * 4LL, "unsupported size");
    inst_norm_kernel<<<(unsigned int)inst64, 256, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), (int)inst64, (int)sp64);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''
_ext = load_inline(name='instancenorm2d_custom_ext_final', cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=['instance_norm_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self, num_features: int):
        super(ModelNew, self).__init__()
        self.num_features = num_features
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.instance_norm_forward(x)