```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 CUDA implementation:
#   1) The GEMM (x @ W^T) is dispatched to the highly-tuned cuBLAS path via
#      torch::mm (writing a custom SGEMM that beats cuBLAS at this size is
#      impractical).
#   2) The bias-add, torch.min(x, constant) and the subtraction of the
#      constant are ALL fused into a single custom CUDA epilogue kernel,
#      eliminating multiple full read/write passes over the (batch, out)
#      activation tensor that PyTorch would otherwise perform. A float4
#      vectorized version is used when out_features % 4 == 0.

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

// Vectorized (float4) fused epilogue: out = min(mm_out + bias, c) - c
__global__ void bias_min_sub_vec4_kernel(
    float4* __restrict__ data,
    const float4* __restrict__ bias,
    const float* __restrict__ c_ptr,
    const int total4,
    const int n4)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total4) {
        const float c = *c_ptr;
        const int col = idx % n4;
        const float4 b = bias[col];
        float4 v = data[idx];
        v.x = fminf(v.x + b.x, c) - c;
        v.y = fminf(v.y + b.y, c) - c;
        v.z = fminf(v.z + b.z, c) - c;
        v.w = fminf(v.w + b.w, c) - c;
        data[idx] = v;
    }
}

// Scalar fallback fused epilogue for arbitrary shapes
__global__ void bias_min_sub_scalar_kernel(
    float* __restrict__ data,
    const float* __restrict__ bias,
    const float* __restrict__ c_ptr,
    const int total,
    const int N)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        const float c = *c_ptr;
        const int col = idx % N;
        const float v = data[idx] + bias[col];
        data[idx] = fminf(v, c) - c;
    }
}

torch::Tensor fused_linear_min_sub(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor constant)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda() && constant.is_cuda(),
                "All tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                weight.scalar_type() == torch::kFloat32 &&
                bias.scalar_type() == torch::kFloat32 &&
                constant.scalar_type() == torch::kFloat32,
                "All tensors must be float32");

    auto xc = x.contiguous();

    // cuBLAS GEMM: out = x @ W^T  (transpose handled natively, no copy)
    auto out = torch::mm(xc, weight.t());

    const int64_t M = out.size(0);
    const int64_t N = out.size(1);
    const int64_t total = M * N;
    if (total == 0) {
        return out;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int block = 256;

    if ((N % 4) == 0) {
        const int total4 = static_cast<int>(total / 4);
        const int n4 = static_cast<int>(N / 4);
        const int grid = (total4 + block - 1) / block;
        bias_min_sub_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias.data_ptr<float>()),
            constant.data_ptr<float>(),
            total4,
            n4);
    } else {
        const int grid = static_cast<int>((total + block - 1) / block);
        bias_min_sub_scalar_kernel<<<grid, block, 0, stream>>>(
            out.data_ptr<float>(),
            bias.data_ptr<float>(),
            constant.data_ptr<float>(),
            static_cast<int>(total),
            static_cast<int>(N));
    }

    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_linear_min_sub(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, torch::Tensor constant);"
)

fused_op = load_inline(
    name="fused_linear_min_sub_ext",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_linear_min_sub"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM + a single fused CUDA kernel performing
    bias-add, min with the constant, and subtraction of the constant in one
    pass over the output 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):
        if (x.is_cuda and x.dtype == torch.float32 and
                self.linear.weight.is_cuda and self.constant.is_cuda):
            return fused_op.fused_linear_min_sub(
                x, self.linear.weight, self.linear.bias, self.constant)
        # Fallback for non-CUDA / non-float32 inputs
        out = F.linear(x, self.linear.weight, self.linear.bias)
        return torch.min(out, self.constant) - self.constant
```