import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a matrix multiplication, subtraction, multiplication, and ReLU activation.
    """
    def __init__(self, in_features, out_features, subtract_value, multiply_value):
        super(Model, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.subtract_value = subtract_value
        self.multiply_value = multiply_value

    def forward(self, x):
        x = self.linear(x)
        x = x - self.subtract_value
        x = x * self.multiply_value
        x = torch.relu(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
subtract_value = 2.0
multiply_value = 1.5

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

def get_init_inputs():
    return [in_features, out_features, subtract_value, multiply_value]

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

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

void fused_relu_scale_cuda(torch::Tensor y, double subtract_value, double multiply_value);

torch::Tensor fused_linear_act(torch::Tensor x,
                              torch::Tensor weight,
                              torch::Tensor bias,
                              double subtract_value,
                              double multiply_value) {
    // Use the same highly optimized GEMM path as torch.nn.Linear, then fuse the
    // three following pointwise ops into one in-place CUDA kernel.
    auto y = at::linear(x, weight, bias);
    fused_relu_scale_cuda(y, subtract_value, multiply_value);
    return y;
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>

__device__ __forceinline__ float act_op(float v, float sub, float mul) {
    float r = (v - sub) * mul;
    return r > 0.0f ? r : 0.0f;
}

__global__ void fused_pointwise_scalar_kernel(float* __restrict__ y,
                                              int64_t n,
                                              float sub,
                                              float mul) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = idx; i < n; i += stride) {
        y[i] = act_op(y[i], sub, mul);
    }
}

__global__ void fused_pointwise_vec4_kernel(float4* __restrict__ y4,
                                            int64_t n4,
                                            float sub,
                                            float mul) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = idx; i < n4; i += stride) {
        float4 v = y4[i];
        v.x = act_op(v.x, sub, mul);
        v.y = act_op(v.y, sub, mul);
        v.z = act_op(v.z, sub, mul);
        v.w = act_op(v.w, sub, mul);
        y4[i] = v;
    }
}

void fused_relu_scale_cuda(torch::Tensor y, double subtract_value, double multiply_value) {
    const int64_t n = y.numel();
    if (n == 0) return;

    const int threads = 256;
    const float sub = static_cast<float>(subtract_value);
    const float mul = static_cast<float>(multiply_value);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    uintptr_t addr = reinterpret_cast<uintptr_t>(y.data_ptr<float>());
    if ((addr % 16 == 0) && ((n & 3LL) == 0)) {
        int64_t n4 = n >> 2;
        int blocks = (int)((n4 + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        fused_pointwise_vec4_kernel<<<blocks, threads, 0, stream>>>(
            reinterpret_cast<float4*>(y.data_ptr<float>()), n4, sub, mul);
    } else {
        int blocks = (int)((n + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        fused_pointwise_scalar_kernel<<<blocks, threads, 0, stream>>>(
            y.data_ptr<float>(), n, sub, mul);
    }
}
'''

_ext = load_inline(
    name="linear_sub_mul_relu_fused_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_linear_act"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """Drop-in replacement for the reference model."""
    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 = subtract_value
        self.multiply_value = multiply_value

    def forward(self, x):
        return _ext.fused_linear_act(
            x,
            self.linear.weight,
            self.linear.bias,
            float(self.subtract_value),
            float(self.multiply_value),
        )