```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:
#   1) per-row instance normalization  (equivalent to InstanceNorm2d applied
#      on a (B, 1, 1, F) tensor, i.e. normalize each row with biased variance)
#   2) residual addition  (+ y)
#   3) elementwise multiplication  (* y)
#
# The expensive (B, in) x (in, out) GEMM is left to cuBLAS via F.linear,
# everything after it is executed in a single kernel launch.
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

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 float* xr = x + (long)row * Fdim;
    const float* yr = y + (long)row * Fdim;
    float* orow = out + (long)row * Fdim;

    float sum = 0.f;
    float sumsq = 0.f;

    // ---- pass 1: accumulate row statistics ----
    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 v = xr4[i];
            sum   += v.x + v.y + v.z + v.w;
            sumsq += 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;
            sumsq += v * v;
        }
    }

    // ---- block reduction (warp shuffles + shared memory) ----
    __shared__ float s_sum[32];
    __shared__ float s_sq[32];
    const unsigned int mask = 0xffffffffu;
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(mask, sum, off);
        sumsq += __shfl_down_sync(mask, sumsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        s_sum[wid] = sum;
        s_sq[wid]  = sumsq;
    }
    __syncthreads();

    const int nwarps = blockDim.x >> 5;
    if (wid == 0) {
        float a = (lane < nwarps) ? s_sum[lane] : 0.f;
        float b = (lane < nwarps) ? s_sq[lane]  : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            a += __shfl_down_sync(mask, a, off);
            b += __shfl_down_sync(mask, b, off);
        }
        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;  // biased, like InstanceNorm
    const float inv  = rsqrtf(var + eps);

    // ---- pass 2: normalize + add + multiply ----
    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        const float4* yr4 = reinterpret_cast<const float4*>(yr);
        float4* or4 = reinterpret_cast<float4*>(orow);
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 xv = xr4[i];
            float4 yv = yr4[i];
            float4 ov;
            ov.x = ((xv.x - mean) * inv + yv.x) * yv.x;
            ov.y = ((xv.y - mean) * inv + yv.y) * yv.y;
            ov.z = ((xv.z - mean) * inv + yv.z) * yv.z;
            ov.w = ((xv.w - mean) * inv + yv.w) * yv.w;
            or4[i] = ov;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float xv = xr[i];
            float yv = yr[i];
            orow[i] = ((xv - mean) * inv + yv) * yv;
        }
    }
}

torch::Tensor fused_inorm_add_mul(torch::Tensor x, torch::Tensor y, double eps) {
    TORCH_CHECK(x.is_cuda() && y.is_cuda(), "x and y must be CUDA tensors");
    TORCH_CHECK(x.dim() == 2 && y.dim() == 2, "x and y must be 2-D");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");

    auto xc = x.contiguous();
    auto yc = y.contiguous();
    auto out = torch::empty_like(xc);

    const int B = xc.size(0);
    const int Fdim = xc.size(1);
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((Fdim & 3) == 0) {
        fused_inorm_add_mul_kernel<true><<<B, threads, 0, stream>>>(
            xc.data_ptr<float>(), yc.data_ptr<float>(),
            out.data_ptr<float>(), Fdim, (float)eps);
    } else {
        fused_inorm_add_mul_kernel<false><<<B, threads, 0, stream>>>(
            xc.data_ptr<float>(), yc.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_inorm_add_mul = 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 (GEMM) runs on cuBLAS, while the instance
    normalization + residual add + elementwise multiply 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.eps = eps
        self.fused = fused_inorm_add_mul

    def forward(self, x, y):
        z = F.linear(x, self.bmm.weight, self.bmm.bias)
        if z.is_cuda and z.dtype == torch.float32:
            return self.fused.fused_inorm_add_mul(z, y, self.eps)
        # fallback path (CPU / non-float32) replicating the original ops
        mean = z.mean(dim=1, keepdim=True)
        var = z.var(dim=1, unbiased=False, keepdim=True)
        z = (z - mean) / torch.sqrt(var + self.eps)
        return (z + y) * y
```