```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
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernel that fuses the entire post-GEMM epilogue:
#   out = min(x @ W^T + bias, constant) - constant
# The GEMM itself is handled by cuBLAS (torch.mm), and this kernel
# fuses bias-add + clamp-min + subtraction into a single vectorized pass.
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void bias_min_sub_vec4_kernel(
    const float4* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ c_ptr,
    float4* __restrict__ out,
    int cols,
    long total4)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total4) {
        const float c = c_ptr[0];
        float4 xv = x[idx];
        int col = (int)((idx * 4) % cols);   // multiple of 4 -> 16B aligned
        float4 bv = *reinterpret_cast<const float4*>(bias + col);
        float4 o;
        o.x = fminf(xv.x + bv.x, c) - c;
        o.y = fminf(xv.y + bv.y, c) - c;
        o.z = fminf(xv.z + bv.z, c) - c;
        o.w = fminf(xv.w + bv.w, c) - c;
        out[idx] = o;
    }
}

__global__ void bias_min_sub_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ c_ptr,
    float* __restrict__ out,
    int cols,
    long total)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        const float c = c_ptr[0];
        int col = (int)(idx % cols);
        float v = x[idx] + bias[col];
        out[idx] = fminf(v, c) - c;
    }
}

torch::Tensor fused_bias_min_sub(torch::Tensor x, torch::Tensor bias, torch::Tensor c) {
    TORCH_CHECK(x.is_cuda() && bias.is_cuda() && c.is_cuda(),
                "All tensors must be CUDA tensors");
    auto xc = x.contiguous();
    auto bc = bias.contiguous();
    auto cc = c.contiguous();
    auto out = torch::empty_like(xc);

    long total = (long)xc.numel();
    int cols = (int)xc.size(xc.dim() - 1);
    const int block = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (total == 0) {
        return out;
    }

    if (cols % 4 == 0 && total % 4 == 0) {
        long total4 = total / 4;
        unsigned int grid = (unsigned int)((total4 + block - 1) / block);
        bias_min_sub_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            bc.data_ptr<float>(),
            cc.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            cols, total4);
    } else {
        unsigned int grid = (unsigned int)((total + block - 1) / block);
        bias_min_sub_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(),
            bc.data_ptr<float>(),
            cc.data_ptr<float>(),
            out.data_ptr<float>(),
            cols, total);
    }
    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_bias_min_sub(torch::Tensor x, torch::Tensor bias, torch::Tensor c);"
)

fused_epilogue = load_inline(
    name="fused_bias_min_sub_epilogue",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["fused_bias_min_sub"],
    verbose=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version: cuBLAS GEMM (torch.mm) followed by a single fused
    CUDA kernel performing bias-add + torch.min(constant) + subtraction.
    """
    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):
        if x.is_cuda and x.dtype == torch.float32:
            # Plain GEMM via cuBLAS; bias/min/sub all fused in custom kernel.
            x = torch.mm(x, self.linear.weight.t())
            return fused_epilogue.fused_bias_min_sub(
                x, self.linear.bias, self.constant.detach())
        # Fallback reference path (CPU or non-fp32 inputs).
        x = self.linear(x)
        x = torch.min(x, self.constant)
        x = x - self.constant
        return x
```