```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

# The architecture is:
#   1) Linear (GEMM + bias)                 -> dominated by cuBLAS, keep as-is
#   2) InstanceNorm2d over each row (C=1, H=1, W=out_features, affine=False,
#      track_running_stats=False)           -> per-row mean/var normalization
#   3) residual add of y
#   4) multiply by y
# Steps 2-4 are fused into a single custom CUDA kernel that computes the row
# statistics with a block reduction and writes the final output in one pass,
# eliminating several kernel launches and extra global-memory round trips.

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

#define BLOCK_SIZE 256

__global__ void instnorm_add_mul_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ out,
    const int N,
    const float eps)
{
    const int row = blockIdx.x;
    const float* x_row = x + (size_t)row * N;
    const float* y_row = y + (size_t)row * N;
    float* out_row   = out + (size_t)row * N;

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

    const int vecN = N & ~3;  // largest multiple of 4 <= N

    // Pass 1: accumulate sum and sum of squares (vectorized)
    for (int i = threadIdx.x * 4; i < vecN; i += BLOCK_SIZE * 4) {
        float4 xv = *reinterpret_cast<const float4*>(x_row + i);
        sum   += xv.x + xv.y + xv.z + xv.w;
        sumsq += xv.x * xv.x + xv.y * xv.y + xv.z * xv.z + xv.w * xv.w;
    }
    for (int i = vecN + threadIdx.x; i < N; i += BLOCK_SIZE) {
        float v = x_row[i];
        sum   += v;
        sumsq += v * v;
    }

    // Block reduction of the statistics
    __shared__ float s_sum[BLOCK_SIZE];
    __shared__ float s_sumsq[BLOCK_SIZE];
    s_sum[threadIdx.x]   = sum;
    s_sumsq[threadIdx.x] = sumsq;
    __syncthreads();

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

    const float mean = s_sum[0] / (float)N;
    float var = s_sumsq[0] / (float)N - mean * mean;  // biased variance, like InstanceNorm
    var = var < 0.0f ? 0.0f : var;
    const float rstd = rsqrtf(var + eps);

    // Pass 2: normalize, add residual, multiply by y (vectorized)
    for (int i = threadIdx.x * 4; i < vecN; i += BLOCK_SIZE * 4) {
        float4 xv = *reinterpret_cast<const float4*>(x_row + i);
        float4 yv = *reinterpret_cast<const float4*>(y_row + i);
        float4 ov;
        ov.x = ((xv.x - mean) * rstd + yv.x) * yv.x;
        ov.y = ((xv.y - mean) * rstd + yv.y) * yv.y;
        ov.z = ((xv.z - mean) * rstd + yv.z) * yv.z;
        ov.w = ((xv.w - mean) * rstd + yv.w) * yv.w;
        *reinterpret_cast<float4*>(out_row + i) = ov;
    }
    for (int i = vecN + threadIdx.x; i < N; i += BLOCK_SIZE) {
        float xv = x_row[i];
        float yv = y_row[i];
        out_row[i] = ((xv - mean) * rstd + yv) * yv;
    }
}

torch::Tensor fused_instnorm_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() == torch::kFloat32 &&
                y.scalar_type() == torch::kFloat32, "Inputs must be float32");
    TORCH_CHECK(x.sizes() == y.sizes(), "Input shapes must match");

    auto x_c = x.contiguous();
    auto y_c = y.contiguous();
    auto out = torch::empty_like(x_c);

    const int rows = (int)x_c.size(0);
    const int N    = (int)x_c.size(1);
    if (rows == 0 || N == 0) {
        return out;
    }

    auto stream = at::cuda::getCurrentCUDAStream();
    instnorm_add_mul_kernel<<<rows, BLOCK_SIZE, 0, stream>>>(
        x_c.data_ptr<float>(),
        y_c.data_ptr<float>(),
        out.data_ptr<float>(),
        N,
        (float)eps);

    return out;
}
"""

fused_instnorm_add_mul_cpp_source = (
    "torch::Tensor fused_instnorm_add_mul(torch::Tensor x, torch::Tensor y, double eps);"
)

fused_instnorm_add_mul = load_inline(
    name="fused_instnorm_add_mul_ext",
    cpp_sources=fused_instnorm_add_mul_cpp_source,
    cuda_sources=fused_instnorm_add_mul_cuda_source,
    functions=["fused_instnorm_add_mul"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuBLAS GEMM (nn.Linear) and fuses
    InstanceNorm + residual add + multiplication into one 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)
        # Kept for state_dict compatibility and as a fallback path
        # (affine=False, track_running_stats=False by default, so it has
        #  no parameters or buffers).
        self.instance_norm = nn.InstanceNorm2d(out_features, eps=eps, momentum=momentum)
        self.eps = eps

    def forward(self, x, y):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_features).
            y (torch.Tensor): Input tensor of shape (batch_size, out_features).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_features).
        """
        x = self.bmm(x)
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 2:
            return fused_instnorm_add_mul.fused_instnorm_add_mul(x, y, self.eps)
        # Fallback (CPU / non-float32): original PyTorch ops
        x = self.instance_norm(x.unsqueeze(1).unsqueeze(1)).squeeze(1).squeeze(1)
        x = x + y
        x = x * y
        return x
```