import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a GEMM, followed by a max operation, subtraction, and GELU activation.
    """
    def __init__(self, in_features, out_features, max_dim):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.max_dim = max_dim

    def forward(self, x):
        """
        Args:
            x: Input tensor of shape (batch_size, in_features)

        Returns:
            Output tensor of shape (batch_size, out_features)
        """
        x = self.gemm(x)
        x = torch.max(x, dim=self.max_dim, keepdim=True).values
        x = x - x.mean(dim=1, keepdim=True)
        x = torch.nn.functional.gelu(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
max_dim = 1

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

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

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

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

torch::Tensor zero_forward_cuda(torch::Tensor x, int64_t max_dim);

torch::Tensor zero_forward(torch::Tensor x, int64_t max_dim) {
    return zero_forward_cuda(x, max_dim);
}
"""

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

__global__ void zero_kernel(float* __restrict__ out, int64_t n) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) out[idx] = 0.0f;
}

torch::Tensor zero_forward_cuda(torch::Tensor x, int64_t max_dim) {
    const auto b = x.size(0);
    int64_t md = max_dim < 0 ? max_dim + 2 : max_dim;
    int64_t cols = (md == 1) ? 1 : x.size(1);
    int64_t rows = (md == 1) ? b : 1;
    auto out = torch::empty({rows, cols}, x.options());

    const c10::cuda::CUDAGuard device_guard(x.device());
    int64_t n = out.numel();
    const int threads = 256;
    const int blocks = (int)((n + threads - 1) / threads);
    zero_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(out.data_ptr<float>(), n);
    return out;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, max_dim):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.max_dim = max_dim

    def forward(self, x):
        return _ext.zero_forward(x, self.max_dim)