import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Max reduction over a specific dimension.
    """
    def __init__(self, dim: int):
        """
        Initializes the model with the dimension to reduce over.

        Args:
            dim (int): The dimension to reduce over.
        """
        super(Model, self).__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Max reduction over the specified dimension to the input tensor.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after Max reduction over the specified dimension.
        """
        return torch.max(x, dim=self.dim)[0]

batch_size = 128
dim1 = 4096
dim2 = 4095

def get_inputs():
    x = torch.rand(batch_size, dim1, dim2)
    return [x]

def get_init_inputs():
    return [1] # Example, change to desired dimension

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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor max_reduce_forward(torch::Tensor x, int64_t dim);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__global__ void max_dim1_4096_4095(const float* __restrict__ x, float* __restrict__ out) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int b = blockIdx.y;
    if (col >= 4095) return;
    const float* p = x + ((long long)b * 4096 * 4095 + col);
    float m = -FLT_MAX;
    #pragma unroll 4
    for (int i = 0; i < 4096; ++i) {
        float v = *p;
        m = fmaxf(m, v);
        p += 4095;
    }
    out[(long long)b * 4095 + col] = m;
}

__global__ void max_dim1_generic(const float* __restrict__ x, float* __restrict__ out, int B, int D1, int D2) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int b = blockIdx.y;
    if (col >= D2) return;
    const float* p = x + ((long long)b * D1 * D2 + col);
    float m = -FLT_MAX;
    for (int i = 0; i < D1; ++i) {
        m = fmaxf(m, *p);
        p += D2;
    }
    out[(long long)b * D2 + col] = m;
}

__global__ void max_dim2_kernel(const float* __restrict__ x, float* __restrict__ out, int B, int D1, int D2) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * D1;
    if (idx >= total) return;
    const float* p = x + (long long)idx * D2;
    float m = -FLT_MAX;
    for (int j = 0; j < D2; ++j) m = fmaxf(m, p[j]);
    out[idx] = m;
}

__global__ void max_dim0_kernel(const float* __restrict__ x, float* __restrict__ out, int B, int D1, int D2) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = D1 * D2;
    if (idx >= total) return;
    int i = idx / D2;
    int j = idx - i * D2;
    float m = -FLT_MAX;
    const float* p = x + (long long)i * D2 + j;
    long long stride = (long long)D1 * D2;
    for (int b = 0; b < B; ++b) m = fmaxf(m, p[(long long)b * stride]);
    out[idx] = m;
}

torch::Tensor max_reduce_forward(torch::Tensor x, int64_t dim) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, "float32 only");
    TORCH_CHECK(x.dim() == 3, "expected 3D tensor");
    int B = (int)x.size(0), D1 = (int)x.size(1), D2 = (int)x.size(2);
    if (dim < 0) dim += 3;
    auto out_opts = x.options();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (dim == 1) {
        auto out = torch::empty({B, D2}, out_opts);
        int threads = 512;
        dim3 grid((D2 + threads - 1) / threads, B);
        if (D1 == 4096 && D2 == 4095)
            max_dim1_4096_4095<<<grid, threads, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>());
        else
            max_dim1_generic<<<grid, threads, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>(), B, D1, D2);
        return out;
    } else if (dim == 2) {
        auto out = torch::empty({B, D1}, out_opts);
        int threads = 256;
        max_dim2_kernel<<<(B * D1 + threads - 1) / threads, threads, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>(), B, D1, D2);
        return out;
    } else {
        TORCH_CHECK(dim == 0, "dim out of range");
        auto out = torch::empty({D1, D2}, out_opts);
        int threads = 256;
        max_dim0_kernel<<<(D1 * D2 + threads - 1) / threads, threads, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>(), B, D1, D2);
        return out;
    }
}
'''

_ext = load_inline(name='max_reduce_final_ext', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['max_reduce_forward'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, dim: int):
        super(ModelNew, self).__init__()
        self.dim = dim
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.max_reduce_forward(x, self.dim)