```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 kernel.
#
# The original model performs, after the linear layer:
#     x = torch.min(x, constant)
#     x = x - constant
# which are two separate element-wise CUDA kernels, each doing a full
# read + write pass over the [batch, out_features] activation.
#
# We fuse them into a single kernel:
#     out = fminf(x, c) - c
# cutting the element-wise memory traffic (and kernel launches) in half.
# The kernel uses vectorized float4 loads/stores when the tensor size
# permits, plus a grid-stride loop for generality.
# ---------------------------------------------------------------------------
fused_min_sub_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 256

__global__ void fused_min_sub_vec4_kernel(const float4* __restrict__ x,
                                          const float* __restrict__ c,
                                          float4* __restrict__ out,
                                          long n_vec) {
    const float cv = __ldg(c);
    long idx = (long)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    const long stride = (long)gridDim.x * BLOCK_SIZE;
    for (; idx < n_vec; idx += stride) {
        float4 v = x[idx];
        v.x = fminf(v.x, cv) - cv;
        v.y = fminf(v.y, cv) - cv;
        v.z = fminf(v.z, cv) - cv;
        v.w = fminf(v.w, cv) - cv;
        out[idx] = v;
    }
}

__global__ void fused_min_sub_kernel(const float* __restrict__ x,
                                     const float* __restrict__ c,
                                     float* __restrict__ out,
                                     long n) {
    const float cv = __ldg(c);
    long idx = (long)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    const long stride = (long)gridDim.x * BLOCK_SIZE;
    for (; idx < n; idx += stride) {
        out[idx] = fminf(x[idx], cv) - cv;
    }
}

torch::Tensor fused_min_sub_cuda(torch::Tensor x, torch::Tensor c) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(c.is_cuda(), "c must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(c.scalar_type() == at::kFloat, "c must be float32");

    auto x_contig = x.contiguous();
    auto out = at::empty_like(x_contig);
    const long n = x_contig.numel();
    if (n == 0) {
        return out;
    }

    const float* x_ptr = x_contig.data_ptr<float>();
    const float* c_ptr = c.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();

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

    if (n % 4 == 0) {
        const long n_vec = n >> 2;
        long blocks = (n_vec + BLOCK_SIZE - 1) / BLOCK_SIZE;
        if (blocks > 8192) blocks = 8192;
        fused_min_sub_vec4_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(
            reinterpret_cast<const float4*>(x_ptr), c_ptr,
            reinterpret_cast<float4*>(out_ptr), n_vec);
    } else {
        long blocks = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
        if (blocks > 8192) blocks = 8192;
        fused_min_sub_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(
            x_ptr, c_ptr, out_ptr, n);
    }
    return out;
}
"""

fused_min_sub_cpp_source = (
    "torch::Tensor fused_min_sub_cuda(torch::Tensor x, torch::Tensor c);"
)

# Compile the inline CUDA extension.
fused_min_sub_ext = load_inline(
    name="fused_min_sub_ext",
    cpp_sources=fused_min_sub_cpp_source,
    cuda_sources=fused_min_sub_cuda_source,
    functions=["fused_min_sub_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    The dominant cost is the [128, 16384] x [16384, 16384] GEMM inside the
    linear layer, which is kept on the highly tuned cuBLAS path via
    F.linear (its bias add is already fused into the GEMM epilogue).

    The two trailing element-wise ops (min with the constant, then
    subtracting the constant) are fused into a single custom CUDA kernel,
    eliminating one kernel launch and one full read+write pass over the
    activation tensor.
    """
    def __init__(self, in_features, out_features, constant):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.constant = nn.Parameter(torch.tensor(constant))

    def forward(self, x):
        x = F.linear(x, self.linear.weight, self.linear.bias)
        if (x.is_cuda and x.dtype == torch.float32
                and self.constant.is_cuda
                and self.constant.dtype == torch.float32):
            x = fused_min_sub_ext.fused_min_sub_cuda(x, self.constant)
        else:
            # Fallback (e.g. CPU tensors) preserving original semantics.
            x = torch.min(x, self.constant) - self.constant
        return x
```