```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:
#   InstanceNorm (over the feature dim) -> residual add (y) -> multiply (y)
#
# The Linear (matmul) part is kept as cuBLAS (F.linear), since it dominates
# the FLOPs and cuBLAS is already near-optimal. Everything after the GEMM is
# fused into a single memory-bound kernel (one thread block per batch row,
# double-precision block reduction for the mean/variance statistics, and
# float4-vectorized loads/stores).
# ---------------------------------------------------------------------------

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

#define BLOCK_SIZE 256

__inline__ __device__ double warpReduceSum(double val) {
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffffu, val, offset);
    }
    return val;
}

// One block handles one row (batch element) of length F.
__global__ void fused_in_add_mul_kernel(
    const float* __restrict__ z,
    const float* __restrict__ y,
    float* __restrict__ out,
    const int F,
    const float eps)
{
    const int row = blockIdx.x;
    const long long base = (long long)row * (long long)F;
    const float* __restrict__ zr = z + base;
    const float* __restrict__ yr = y + base;
    float* __restrict__ outr = out + base;

    const int F4 = F >> 2;
    const float4* __restrict__ z4 = reinterpret_cast<const float4*>(zr);

    // ---- Pass 1: compute sum / sum-of-squares in double precision ----
    double sum = 0.0;
    double sumsq = 0.0;
    for (int i = threadIdx.x; i < F4; i += blockDim.x) {
        const float4 v = z4[i];
        sum   += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
        sumsq += (double)v.x * (double)v.x + (double)v.y * (double)v.y
               + (double)v.z * (double)v.z + (double)v.w * (double)v.w;
    }

    __shared__ double s_sum[BLOCK_SIZE / 32];
    __shared__ double s_sumsq[BLOCK_SIZE / 32];
    __shared__ double s_mean;
    __shared__ double s_rstd;

    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;

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

    if (wid == 0) {
        double a = (lane < BLOCK_SIZE / 32) ? s_sum[lane] : 0.0;
        double b = (lane < BLOCK_SIZE / 32) ? s_sumsq[lane] : 0.0;
        a = warpReduceSum(a);
        b = warpReduceSum(b);
        if (lane == 0) {
            const double mean = a / (double)F;
            double var = b / (double)F - mean * mean;
            var = var < 0.0 ? 0.0 : var;   // guard against fp roundoff
            s_mean = mean;
            s_rstd = 1.0 / sqrt(var + (double)eps);
        }
    }
    __syncthreads();

    const float mean = (float)s_mean;
    const float rstd = (float)s_rstd;

    // ---- Pass 2: normalize -> add y -> multiply by y ----
    const float4* __restrict__ y4 = reinterpret_cast<const float4*>(yr);
    float4* __restrict__ o4 = reinterpret_cast<float4*>(outr);
    for (int i = threadIdx.x; i < F4; i += blockDim.x) {
        const float4 zv = z4[i];
        const float4 yv = y4[i];
        float4 ov;
        ov.x = ((zv.x - mean) * rstd + yv.x) * yv.x;
        ov.y = ((zv.y - mean) * rstd + yv.y) * yv.y;
        ov.z = ((zv.z - mean) * rstd + yv.z) * yv.z;
        ov.w = ((zv.w - mean) * rstd + yv.w) * yv.w;
        o4[i] = ov;
    }
}

torch::Tensor fused_in_add_mul(torch::Tensor z, torch::Tensor y, double eps) {
    TORCH_CHECK(z.is_cuda() && y.is_cuda(), "Inputs must be CUDA tensors.");
    TORCH_CHECK(z.scalar_type() == at::kFloat && y.scalar_type() == at::kFloat,
                "Inputs must be float32.");
    TORCH_CHECK(z.dim() == 2 && y.dim() == 2, "Inputs must be 2-D.");
    TORCH_CHECK(z.size(0) == y.size(0) && z.size(1) == y.size(1),
                "z and y must have the same shape.");
    auto zc = z.contiguous();
    auto yc = y.contiguous();
    auto out = at::empty_like(zc);

    const int B = (int)zc.size(0);
    const int Fdim = (int)zc.size(1);
    TORCH_CHECK(Fdim % 4 == 0, "Feature dimension must be divisible by 4.");

    dim3 grid(B);
    dim3 block(BLOCK_SIZE);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
    fused_in_add_mul_kernel<<<grid, block, 0, stream>>>(
        zc.data_ptr<float>(), yc.data_ptr<float>(), out.data_ptr<float>(),
        Fdim, (float)eps);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_in_add_mul_cpp_source = (
    "torch::Tensor fused_in_add_mul(torch::Tensor z, torch::Tensor y, double eps);"
)

# Compile the inline CUDA extension
fused_ops = load_inline(
    name="fused_in_add_mul",
    cpp_sources=fused_in_add_mul_cpp_source,
    cuda_sources=fused_in_add_mul_cuda_source,
    functions=["fused_in_add_mul"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model:
      - Linear (GEMM + bias) executed by cuBLAS via F.linear.
      - InstanceNorm (per-row normalization) + residual add + elementwise
        multiply 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_ops = fused_ops

    def forward(self, x, y):
        # GEMM + bias (cuBLAS)
        z = F.linear(x, self.bmm.weight, self.bmm.bias)

        # Fused instance-norm + add + mul path (default InstanceNorm2d has
        # affine=False / track_running_stats=False, so it is a plain
        # per-row normalization with the given eps).
        if (z.is_cuda and z.dtype == torch.float32 and z.size(1) % 4 == 0
                and self.instance_norm.weight is None):
            return self.fused_ops.fused_in_add_mul(z, y, self.eps)

        # Fallback to the reference computation (e.g. CPU tensors).
        out = self.instance_norm(z.unsqueeze(1).unsqueeze(1)).squeeze(1).squeeze(1)
        out = out + y
        out = out * y
        return out
```