```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# ---------------------------------------------------------------------------
# Fused CUDA kernel: per-row instance normalization + residual add + multiply
#
# The original graph is:
#   h = Linear(x)                                    (cuBLAS, keep as-is)
#   h = InstanceNorm2d(h.unsqueeze(1).unsqueeze(1))  (normalize over features)
#   h = h + y
#   h = h * y
#
# The instance-norm over a (batch, 1, 1, F) tensor is exactly a per-row
# normalization over the F features, so we fuse norm + add + mul into a
# single kernel launch: one thread block per row computes mean/variance with
# a warp-shuffle block reduction, then writes ((h-mean)/sqrt(var+eps)+y)*y
# using vectorized float4 accesses.
# ---------------------------------------------------------------------------

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

__device__ __forceinline__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffffu, val, offset);
    return val;
}

template <bool VEC>
__global__ void fused_inorm_add_mul_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ out,
    const int Fdim,
    const float eps)
{
    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)Fdim;
    const float* xr = x + base;
    const float* yr = y + base;
    float* orow = out + base;

    // ---- pass 1: accumulate sum and sum of squares -----------------------
    float sum = 0.0f, sq = 0.0f;
    if (VEC) {
        const int n4 = Fdim >> 2;
        const float4* x4 = reinterpret_cast<const float4*>(xr);
        for (int i = threadIdx.x; i < n4; i += blockDim.x) {
            float4 v = x4[i];
            sum += (v.x + v.y) + (v.z + v.w);
            sq  += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float v = xr[i];
            sum += v;
            sq  += v * v;
        }
    }

    // ---- block reduction (sum, sumsq) ------------------------------------
    __shared__ float s_sum[32];
    __shared__ float s_sq[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;

    sum = warpReduceSum(sum);
    sq  = warpReduceSum(sq);
    if (lane == 0) { s_sum[wid] = sum; s_sq[wid] = sq; }
    __syncthreads();

    const int nwarps = (int)((blockDim.x + 31) >> 5);
    if (wid == 0) {
        float a = (lane < nwarps) ? s_sum[lane] : 0.0f;
        float b = (lane < nwarps) ? s_sq[lane]  : 0.0f;
        a = warpReduceSum(a);
        b = warpReduceSum(b);
        if (lane == 0) { s_sum[0] = a; s_sq[0] = b; }
    }
    __syncthreads();

    const float mean = s_sum[0] / (float)Fdim;
    const float var  = s_sq[0] / (float)Fdim - mean * mean;
    const float inv  = rsqrtf(var + eps);

    // ---- pass 2: normalize + residual add + multiply ----------------------
    if (VEC) {
        const int n4 = Fdim >> 2;
        const float4* x4 = reinterpret_cast<const float4*>(xr);
        const float4* y4 = reinterpret_cast<const float4*>(yr);
        float4* o4 = reinterpret_cast<float4*>(orow);
        for (int i = threadIdx.x; i < n4; i += blockDim.x) {
            float4 v = x4[i];
            float4 u = y4[i];
            float4 r;
            r.x = ((v.x - mean) * inv + u.x) * u.x;
            r.y = ((v.y - mean) * inv + u.y) * u.y;
            r.z = ((v.z - mean) * inv + u.z) * u.z;
            r.w = ((v.w - mean) * inv + u.w) * u.w;
            o4[i] = r;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float v = xr[i];
            float u = yr[i];
            orow[i] = ((v - mean) * inv + u) * u;
        }
    }
}

torch::Tensor fused_inorm_add_mul(torch::Tensor x, torch::Tensor y, double eps) {
    TORCH_CHECK(x.is_cuda() && y.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && y.scalar_type() == at::kFloat,
                "only float32 tensors are supported");
    TORCH_CHECK(x.is_contiguous() && y.is_contiguous(),
                "inputs must be contiguous");

    const int B = (int)x.size(0);
    const int Fdim = (int)x.size(1);
    auto out = at::empty_like(x);

    const int threads = 256;
    auto stream = at::cuda::getCurrentCUDAStream();
    if ((Fdim & 3) == 0) {
        fused_inorm_add_mul_kernel<true><<<B, threads, 0, stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(),
            out.data_ptr<float>(), Fdim, (float)eps);
    } else {
        fused_inorm_add_mul_kernel<false><<<B, threads, 0, stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(),
            out.data_ptr<float>(), Fdim, (float)eps);
    }
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_inorm_add_mul(torch::Tensor x, torch::Tensor y, double eps);"
)

fused_op = load_inline(
    name="fused_inorm_add_mul",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_inorm_add_mul"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the Linear layer stays on cuBLAS (F.linear), while the
    instance normalization, residual addition and final multiplication are
    fused into a single custom CUDA kernel.
    """
    def __init__(self, in_features, out_features, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.bmm = nn.Linear(in_features, out_features)
        self.instance_norm = nn.InstanceNorm2d(out_features, eps=eps, momentum=momentum)
        self.eps = eps
        self.fused = fused_op

    def forward(self, x, y):
        h = F.linear(x, self.bmm.weight, self.bmm.bias)
        if h.is_cuda and h.dtype == torch.float32:
            return self.fused.fused_inorm_add_mul(h.contiguous(), y.contiguous(), self.eps)
        # Fallback (CPU / non-float32): original eager path.
        z = self.instance_norm(h.unsqueeze(1).unsqueeze(1)).squeeze(1).squeeze(1)
        z = z + y
        return z * y
```