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

# Custom fused CUDA implementation:
#   out = relu((x @ W^T + bias - subtract_value) * multiply_value)
#
# Fast path: a single custom tensor-core (TF32 WMMA) GEMM kernel with the
# bias + subtract + multiply + ReLU epilogue fused in, so the whole operator
# is one kernel launch. Fallback path: cuBLAS GEMM + a custom fused
# elementwise kernel (subtract + multiply + ReLU in one pass).
_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <mutex>
#include <algorithm>

using namespace nvcuda;

constexpr int BM = 128;
constexpr int BN = 128;
constexpr int BK = 32;
constexpr int LDA = BK + 8;  // padded leading dim for A tile (stored [BM][LDA])
constexpr int LDB = BK + 8;  // padded leading dim for B tile (stored [BN][LDB])
constexpr int LDS = BN + 4;  // padded leading dim for epilogue staging tile
constexpr int SMEM_TILE_BYTES  = (BM * LDA + BN * LDB) * (int)sizeof(float);
constexpr int SMEM_STAGE_BYTES = BM * LDS * (int)sizeof(float);
constexpr int SMEM_BYTES = SMEM_TILE_BYTES > SMEM_STAGE_BYTES ? SMEM_TILE_BYTES : SMEM_STAGE_BYTES;

// Fused GEMM + bias + subtract + multiply + ReLU.
//   A: [M, K] row-major (activations x)
//   W: [N, K] row-major (linear weight);  W[n][k] == (W^T)[k][n]
//   Out: [M, N] row-major
// Requires M % BM == 0, N % BN == 0, K % BK == 0 and sm_80+ (TF32 tensor cores).
__global__ void __launch_bounds__(256) gemm_tf32_bias_affine_relu_kernel(
    const float* __restrict__ A,
    const float* __restrict__ W,
    const float* __restrict__ bias,
    float* __restrict__ Out,
    int M, int N, int K,
    float sub, float mul)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    extern __shared__ float smem[];
    float* As = smem;             // [BM][LDA], row-major A tile
    float* Bs = smem + BM * LDA;  // [BN][LDB], W tile kept as [n][k]

    const int tid   = threadIdx.x;
    const int warp  = tid >> 5;
    const int warpM = warp >> 1;  // 0..3  (4 warps along M)
    const int warpN = warp & 1;   // 0..1  (2 warps along N)

    const int m0 = blockIdx.y * BM;
    const int n0 = blockIdx.x * BN;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[2][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            wmma::fill_fragment(acc[i][j], 0.0f);
        }
    }

    const float4* A4 = reinterpret_cast<const float4*>(A);
    const float4* W4 = reinterpret_cast<const float4*>(W);

    const int ldRow = tid >> 3;  // 0..31
    const int ldK4  = tid & 7;   // float4 index within a 32-float row

    for (int k0 = 0; k0 < K; k0 += BK) {
        // Coalesced float4 loads of the A tile (BM x BK)
        #pragma unroll
        for (int rr = 0; rr < BM; rr += 32) {
            const int r = ldRow + rr;
            float4 v = A4[((size_t)(m0 + r) * K + k0) / 4 + ldK4];
            *reinterpret_cast<float4*>(&As[r * LDA + ldK4 * 4]) = v;
        }
        // Coalesced float4 loads of the W tile (BN x BK), kept as [n][k]
        #pragma unroll
        for (int rr = 0; rr < BN; rr += 32) {
            const int r = ldRow + rr;
            float4 v = W4[((size_t)(n0 + r) * K + k0) / 4 + ldK4];
            *reinterpret_cast<float4*>(&Bs[r * LDB + ldK4 * 4]) = v;
        }
        __syncthreads();

        #pragma unroll
        for (int kf = 0; kf < BK / 8; ++kf) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::col_major> b_frag[4];
            wmma::load_matrix_sync(a_frag[0], &As[(warpM * 32)      * LDA + kf * 8], LDA);
            wmma::load_matrix_sync(a_frag[1], &As[(warpM * 32 + 16) * LDA + kf * 8], LDA);
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                // Bs[n][k] viewed as col-major (k, n) fragment == W^T tile
                wmma::load_matrix_sync(b_frag[j], &Bs[(warpN * 64 + j * 16) * LDB + kf * 8], LDB);
            }
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    wmma::mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]);
                }
            }
        }
        __syncthreads();
    }

    // Fused epilogue: stage accumulators in shared, then apply bias/affine/ReLU.
    float* Stage = smem;  // reuse shared memory
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            wmma::store_matrix_sync(&Stage[(warpM * 32 + i * 16) * LDS + warpN * 64 + j * 16],
                                    acc[i][j], LDS, wmma::mem_row_major);
        }
    }
    __syncthreads();

    for (int idx = tid; idx < BM * BN; idx += 256) {
        const int r = idx >> 7;
        const int c = idx & (BN - 1);
        float v = Stage[r * LDS + c];
        v = (v + bias[n0 + c] - sub) * mul;
        Out[(size_t)(m0 + r) * N + n0 + c] = fmaxf(v, 0.0f);
    }
#endif
}

// Fused subtract + multiply + ReLU (vectorized), used by the fallback path.
__global__ void affine_relu_vec4_kernel(const float4* __restrict__ in,
                                        float4* __restrict__ out,
                                        float sub, float mul, long n4) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long stride = (long)gridDim.x * blockDim.x;
    for (; i < n4; i += stride) {
        float4 v = in[i];
        v.x = fmaxf((v.x - sub) * mul, 0.0f);
        v.y = fmaxf((v.y - sub) * mul, 0.0f);
        v.z = fmaxf((v.z - sub) * mul, 0.0f);
        v.w = fmaxf((v.w - sub) * mul, 0.0f);
        out[i] = v;
    }
}

__global__ void affine_relu_tail_kernel(const float* __restrict__ in,
                                        float* __restrict__ out,
                                        float sub, float mul, long start, long n) {
    long i = start + blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float v = (in[i] - sub) * mul;
        out[i] = fmaxf(v, 0.0f);
    }
}

static bool tf32_supported() {
    static int cached = -1;
    if (cached < 0) {
        int dev = 0;
        cudaGetDevice(&dev);
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, dev);
        cached = (prop.major >= 8) ? 1 : 0;
    }
    return cached == 1;
}

torch::Tensor fused_linear_affine_relu(torch::Tensor x, torch::Tensor weight,
                                       torch::Tensor bias, double sub, double mul) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "CUDA tensors required");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "float32 required");
    TORCH_CHECK(x.dim() == 2 && weight.dim() == 2, "2D tensors expected");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();

    const int64_t M = xc.size(0);
    const int64_t K = xc.size(1);
    const int64_t N = wc.size(0);
    TORCH_CHECK(wc.size(1) == K && bc.numel() == N, "shape mismatch");

    auto stream = at::cuda::getCurrentCUDAStream();

    if ((M % BM == 0) && (N % BN == 0) && (K % BK == 0) && tf32_supported()) {
        auto out = torch::empty({M, N}, xc.options());
        static std::once_flag once;
        std::call_once(once, []() {
            cudaFuncSetAttribute(gemm_tf32_bias_affine_relu_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 SMEM_BYTES);
        });
        dim3 grid((unsigned)(N / BN), (unsigned)(M / BM));
        gemm_tf32_bias_affine_relu_kernel<<<grid, 256, SMEM_BYTES, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            out.data_ptr<float>(), (int)M, (int)N, (int)K,
            (float)sub, (float)mul);
        cudaError_t err = cudaGetLastError();
        TORCH_CHECK(err == cudaSuccess, "kernel launch failed: ", cudaGetErrorString(err));
        return out;
    }

    // Fallback: cuBLAS GEMM, then one fused subtract+multiply+ReLU pass.
    auto tmp = torch::addmm(bc, xc, wc.t());
    const int64_t n = tmp.numel();
    float* p = tmp.data_ptr<float>();
    const int64_t n4 = n >> 2;
    if (n4 > 0) {
        const int threads = 256;
        const int64_t want = (n4 + threads - 1) / threads;
        const int blocks = (int)std::min<int64_t>(want, 4096);
        affine_relu_vec4_kernel<<<blocks, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(p), reinterpret_cast<float4*>(p),
            (float)sub, (float)mul, n4);
    }
    if (n & 3) {
        affine_relu_tail_kernel<<<1, 8, 0, stream>>>(p, p, (float)sub, (float)mul, n4 << 2, n);
    }
    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "kernel launch failed: ", cudaGetErrorString(err));
    return tmp;
}
"""

_fused_cpp_src = (
    "torch::Tensor fused_linear_affine_relu(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, double sub, double mul);"
)

_fused_ext = load_inline(
    name="fused_linear_affine_relu_ext",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_linear_affine_relu"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


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

    The matrix multiplication (nn.Linear), scalar subtraction, scalar
    multiplication and ReLU are all fused into a single custom CUDA kernel
    (tensor-core GEMM with a fused epilogue). Falls back to cuBLAS + a custom
    fused elementwise kernel when the fast path is not applicable.
    """
    def __init__(self, in_features, out_features, subtract_value, multiply_value):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.subtract_value = float(subtract_value)
        self.multiply_value = float(multiply_value)

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and self.linear.bias is not None:
            return _fused_ext.fused_linear_affine_relu(
                x,
                self.linear.weight,
                self.linear.bias,
                self.subtract_value,
                self.multiply_value,
            )
        out = F.linear(x, self.linear.weight, self.linear.bias)
        out = (out - self.subtract_value) * self.multiply_value
        return torch.relu(out)
```