import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, applies minimum, and subtracts a constant.
    """
    def __init__(self, in_features, out_features, constant):
        super(Model, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.constant = nn.Parameter(torch.tensor(constant))

    def forward(self, x):
        x = self.linear(x)
        x = torch.min(x, self.constant)
        x = x - self.constant
        return x

batch_size = 128
in_features = 16384
out_features = 16384
constant = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_features)]

def get_init_inputs():
    return [in_features, out_features, constant]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor linear_min_sub(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor constant);
"""

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

__global__ void min_sub_kernel(float* __restrict__ y, const float* __restrict__ cptr, long long n) {
    const float c = __ldg(cptr);
    long long i4 = (long long)(blockIdx.x) * blockDim.x + threadIdx.x;
    long long n4 = n >> 2;
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y);
    for (long long i = i4; i < n4; i += (long long)blockDim.x * gridDim.x) {
        float4 v = y4[i];
        v.x = fminf(v.x, c) - c;
        v.y = fminf(v.y, c) - c;
        v.z = fminf(v.z, c) - c;
        v.w = fminf(v.w, c) - c;
        y4[i] = v;
    }
    long long base = n4 << 2;
    for (long long j = base + i4; j < n; j += (long long)blockDim.x * gridDim.x) {
        float v = y[j];
        y[j] = fminf(v, c) - c;
    }
}

torch::Tensor linear_min_sub(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor constant) {
    c10::cuda::CUDAGuard device_guard(x.device());
    // Equivalent to torch.nn.functional.linear(x, weight, bias).  weight.t() is a view;
    // addmm dispatches to cuBLAS/cuBLASLt for the large GEMM.
    torch::Tensor out = at::addmm(bias, x, weight.t());

    const long long n = out.numel();
    if (n > 0) {
        const int threads = 256;
        long long work = (n + 3) >> 2;
        int blocks = (int)((work + threads - 1) / threads);
        if (blocks > 4096) blocks = 4096;
        min_sub_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            out.data_ptr<float>(), constant.data_ptr<float>(), n);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return out;
}
"""

_ext = load_inline(
    name="linear_min_sub_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["linear_min_sub"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    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):
        return _ext.linear_min_sub(x, self.linear.weight, self.linear.bias, self.constant)