```
# 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 that replaces:
#   1) InstanceNorm2d over the (1, out_features) plane  == per-row normalization
#      over the feature dimension (biased variance, no affine, batch statistics
#      since track_running_stats=False).
#   2) residual addition  (x + y)
#   3) elementwise multiply ((x + y) * y)
#
# Each thread block handles one row of the (batch_size, out_features) matrix.
# The row is staged in shared memory so global memory is read exactly once
# for the linear output and once for y, with a single write of the result.
fused_inorm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#define FULL_MASK 0xffffffffu

__inline__ __device__ float warp_reduce_sum(float v) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(FULL_MASK, v, offset);
    }
    return v;
}

__global__ void fused_inorm_add_mul_kernel(
    const float* __restrict__ h,
    const float* __restrict__ y,
    float* __restrict__ out,
    int L,
    float eps)
{
    extern __shared__ __align__(16) float sh[];
    __shared__ float red[32];

    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)L;
    const float* __restrict__ hr = h + base;
    const float* __restrict__ yr = y + base;
    float* __restrict__ orr = out + base;

    const float4* __restrict__ hr4 = reinterpret_cast<const float4*>(hr);
    float4* __restrict__ sh4 = reinterpret_cast<float4*>(sh);
    // Vector path is only valid when every row stays 16-byte aligned.
    const int L4 = ((L & 3) == 0) ? (L >> 2) : 0;

    // Pass 1: load row into shared memory and accumulate the sum.
    float sum = 0.0f;
    for (int j = threadIdx.x; j < L4; j += blockDim.x) {
        float4 v = hr4[j];
        sh4[j] = v;
        sum += (v.x + v.y) + (v.z + v.w);
    }
    for (int j = (L4 << 2) + threadIdx.x; j < L; j += blockDim.x) {
        float v = hr[j];
        sh[j] = v;
        sum += v;
    }

    // Block reduce the sum.
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    const int nwarps = (blockDim.x + 31) >> 5;

    sum = warp_reduce_sum(sum);
    if (lane == 0) red[wid] = sum;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < nwarps) ? red[lane] : 0.0f;
        v = warp_reduce_sum(v);
        if (lane == 0) red[0] = v;
    }
    __syncthreads();
    const float mean = red[0] / (float)L;
    __syncthreads();  // red[] is reused below

    // Pass 2: two-pass variance (sum of squared deviations) for numerical
    // stability, matching the reference instance normalization.
    float sq = 0.0f;
    for (int j = threadIdx.x; j < L4; j += blockDim.x) {
        float4 v = sh4[j];
        float a = v.x - mean;
        float b = v.y - mean;
        float c = v.z - mean;
        float d = v.w - mean;
        sq += (a * a + b * b) + (c * c + d * d);
    }
    for (int j = (L4 << 2) + threadIdx.x; j < L; j += blockDim.x) {
        float d = sh[j] - mean;
        sq += d * d;
    }

    sq = warp_reduce_sum(sq);
    if (lane == 0) red[wid] = sq;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < nwarps) ? red[lane] : 0.0f;
        v = warp_reduce_sum(v);
        if (lane == 0) red[0] = v;
    }
    __syncthreads();
    const float var = red[0] / (float)L;
    const float inv_std = rsqrtf(var + eps);

    // Pass 3: normalize + residual add + multiply, write result.
    const float4* __restrict__ yr4 = reinterpret_cast<const float4*>(yr);
    float4* __restrict__ orr4 = reinterpret_cast<float4*>(orr);
    for (int j = threadIdx.x; j < L4; j += blockDim.x) {
        float4 v = sh4[j];
        float4 t = yr4[j];
        float4 o;
        o.x = ((v.x - mean) * inv_std + t.x) * t.x;
        o.y = ((v.y - mean) * inv_std + t.y) * t.y;
        o.z = ((v.z - mean) * inv_std + t.z) * t.z;
        o.w = ((v.w - mean) * inv_std + t.w) * t.w;
        orr4[j] = o;
    }
    for (int j = (L4 << 2) + threadIdx.x; j < L; j += blockDim.x) {
        float t = yr[j];
        orr[j] = ((sh[j] - mean) * inv_std + t) * t;
    }
}

torch::Tensor fused_inorm_add_mul(torch::Tensor h, torch::Tensor y, double eps) {
    TORCH_CHECK(h.is_cuda() && y.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(h.scalar_type() == at::kFloat && y.scalar_type() == at::kFloat,
                "inputs must be float32");
    TORCH_CHECK(h.dim() == 2 && y.dim() == 2, "inputs must be 2D");
    TORCH_CHECK(h.size(0) == y.size(0) && h.size(1) == y.size(1),
                "shape mismatch");

    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(h));

    auto hc = h.contiguous();
    auto yc = y.contiguous();
    auto out = at::empty_like(hc);

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

    const size_t shmem = (size_t)L * sizeof(float);

    // Allow large dynamic shared-memory segments (needed when L*4 > 48KB).
    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(fused_inorm_add_mul_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             200 * 1024);
        attr_set = true;
    }
    TORCH_CHECK(shmem <= 200 * 1024, "feature dimension too large");

    const int threads = 256;
    const dim3 grid(rows);
    auto stream = at::cuda::getCurrentCUDAStream();

    fused_inorm_add_mul_kernel<<<grid, threads, shmem, stream>>>(
        hc.data_ptr<float>(),
        yc.data_ptr<float>(),
        out.data_ptr<float>(),
        L,
        (float)eps);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_inorm_cpp_source = (
    "torch::Tensor fused_inorm_add_mul(torch::Tensor h, torch::Tensor y, double eps);"
)

fused_inorm = load_inline(
    name="fused_inorm_add_mul",
    cpp_sources=fused_inorm_cpp_source,
    cuda_sources=fused_inorm_cuda_source,
    functions=["fused_inorm_add_mul"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The Linear (batched matmul) is kept as-is since cuBLAS is already optimal
    for it. The instance normalization (which, for the (B, 1, 1, F) view used
    here, simply normalizes each row over the feature dimension with biased
    variance and no affine parameters), the residual addition, and the final
    multiplication are fused into a single custom CUDA kernel, eliminating
    several kernel launches and multiple full passes over the data.
    """
    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

    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).
        """
        h = self.bmm(x)
        return self.fused.fused_inorm_add_mul(h, y, self.eps)
```