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

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

// Optimized for the benchmark shape [128,4096,4095], reducing dimension 1.
// One block handles a contiguous tile of columns for one batch item.  For each
// row of the reduced dimension, lanes in a warp read contiguous columns, so all
// global loads are coalesced although each output element is a strided reduce.
__global__ void max_dim1_kernel(const float* __restrict__ x,
                                float* __restrict__ y,
                                int B, int D1, int D2) {
    int b = blockIdx.y;
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= D2 || b >= B) return;

    const float* p = x + ((long long)b * D1 * D2 + j);
    float m = -FLT_MAX;

    // D1 is 4096 in the target; unrolling improves ILP a bit while preserving
    // coalescing across the block for every row.
    int i = 0;
    #pragma unroll 4
    for (; i + 3 < D1; i += 4) {
        float v0 = p[(long long)(i + 0) * D2];
        float v1 = p[(long long)(i + 1) * D2];
        float v2 = p[(long long)(i + 2) * D2];
        float v3 = p[(long long)(i + 3) * D2];
        m = fmaxf(m, v0);
        m = fmaxf(m, v1);
        m = fmaxf(m, v2);
        m = fmaxf(m, v3);
    }
    for (; i < D1; ++i) {
        m = fmaxf(m, p[(long long)i * D2]);
    }
    y[(long long)b * D2 + j] = m;
}

// Contiguous last-dimension reduction.  Included for completeness; not expected
// for the provided initializer, but avoids relying on PyTorch fallbacks.
__global__ void max_dim2_kernel(const float* __restrict__ x,
                                float* __restrict__ y,
                                int B, int D1, int D2) {
    int row = blockIdx.x; // row in flattened [B,D1]
    if (row >= B * D1) return;
    const float* p = x + (long long)row * D2;
    float m = -FLT_MAX;
    for (int k = threadIdx.x; k < D2; k += blockDim.x) m = fmaxf(m, p[k]);

    // block reduction in shared memory
    __shared__ float smem[256];
    smem[threadIdx.x] = m;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) smem[threadIdx.x] = fmaxf(smem[threadIdx.x], smem[threadIdx.x + s]);
        __syncthreads();
    }
    if (threadIdx.x < 32) {
        float v = smem[threadIdx.x];
        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 16));
        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 8));
        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 4));
        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 2));
        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 1));
        if (threadIdx.x == 0) y[row] = v;
    }
}

// Reduction over batch dimension for 3-D tensors.
__global__ void max_dim0_kernel(const float* __restrict__ x,
                                float* __restrict__ y,
                                int B, int D1, int D2) {
    long long n = (long long)D1 * D2;
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    float m = -FLT_MAX;
    for (int b = 0; b < B; ++b) {
        m = fmaxf(m, x[(long long)b * n + idx]);
    }
    y[idx] = m;
}

torch::Tensor max_forward(torch::Tensor x, int64_t dim) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(x.dim() == 3, "only 3D tensors supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    int B = (int)x.size(0);
    int D1 = (int)x.size(1);
    int D2 = (int)x.size(2);
    if (dim < 0) dim += 3;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (dim == 1) {
        auto y = torch::empty({B, D2}, x.options());
        int threads = 256;
        dim3 block(threads);
        dim3 grid((D2 + threads - 1) / threads, B);
        max_dim1_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), B, D1, D2);
        return y;
    } else if (dim == 2) {
        auto y = torch::empty({B, D1}, x.options());
        int threads = 256;
        max_dim2_kernel<<<B * D1, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), B, D1, D2);
        return y;
    } else { // dim == 0
        TORCH_CHECK(dim == 0, "dim out of range");
        auto y = torch::empty({D1, D2}, x.options());
        int threads = 256;
        long long n = (long long)D1 * D2;
        int grid = (int)((n + threads - 1) / threads);
        max_dim0_kernel<<<grid, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), B, D1, D2);
        return y;
    }
}
'''

_ext = load_inline(
    name='max_reduce_dim1_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['max_forward'],
    extra_cuda_cflags=['-O3'],
    extra_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_forward(x, self.dim)