import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs min 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 min reduction over the specified dimension to the input tensor.

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

        Returns:
            torch.Tensor: Output tensor after min reduction over the specified dimension.
        """
        return torch.min(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 min_forward(torch::Tensor x, int64_t dim);
"""

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

__global__ __launch_bounds__(256)
void min_dim1_kernel(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;

    int base = (b * D1) * D2 + col;
    float m = CUDART_INF_F;

    // Hot path for the benchmark shape.  Unrolling removes a sizeable amount
    // of loop/control overhead while preserving the fully coalesced access
    // pattern across columns in each warp.
    if (D1 == 4096) {
        int idx = base;
        #pragma unroll 1
        for (int k = 0; k < 4096; k += 8) {
            float v0 = x[idx];              idx += D2;
            float v1 = x[idx];              idx += D2;
            float v2 = x[idx];              idx += D2;
            float v3 = x[idx];              idx += D2;
            float v4 = x[idx];              idx += D2;
            float v5 = x[idx];              idx += D2;
            float v6 = x[idx];              idx += D2;
            float v7 = x[idx];              idx += D2;
            m = fminf(m, v0);
            m = fminf(m, v1);
            m = fminf(m, v2);
            m = fminf(m, v3);
            m = fminf(m, v4);
            m = fminf(m, v5);
            m = fminf(m, v6);
            m = fminf(m, v7);
        }
    } else {
        int idx = base;
        for (int k = 0; k < D1; ++k) {
            m = fminf(m, x[idx]);
            idx += D2;
        }
    }
    out[b * D2 + col] = m;
}

__global__ __launch_bounds__(256)
void min_dim0_kernel(const float* __restrict__ x, float* __restrict__ out,
                     int B, int D1, int D2) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int d1 = blockIdx.y;
    if (col >= D2) return;
    int plane = D1 * D2;
    int idx = d1 * D2 + col;
    float m = CUDART_INF_F;
    for (int b = 0; b < B; ++b) {
        m = fminf(m, x[idx + b * plane]);
    }
    out[d1 * D2 + col] = m;
}

__global__ __launch_bounds__(256)
void min_dim2_kernel(const float* __restrict__ x, float* __restrict__ out,
                     int B, int D1, int D2) {
    extern __shared__ float smem[];
    int row = blockIdx.x;  // row over flattened B*D1
    int tid = threadIdx.x;
    int base = row * D2;
    float m = CUDART_INF_F;
    for (int j = tid; j < D2; j += blockDim.x) {
        m = fminf(m, x[base + j]);
    }
    smem[tid] = m;
    __syncthreads();

    for (int s = blockDim.x >> 1; s > 32; s >>= 1) {
        if (tid < s) smem[tid] = fminf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    if (tid < 32) {
        float v = smem[tid];
        v = fminf(v, smem[tid + 32]);
        unsigned mask = 0xffffffffu;
        v = fminf(v, __shfl_down_sync(mask, v, 16));
        v = fminf(v, __shfl_down_sync(mask, v, 8));
        v = fminf(v, __shfl_down_sync(mask, v, 4));
        v = fminf(v, __shfl_down_sync(mask, v, 2));
        v = fminf(v, __shfl_down_sync(mask, v, 1));
        if (tid == 0) out[row] = v;
    }
}

torch::Tensor min_forward(torch::Tensor x, int64_t dim) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(x.dim() == 3, "expected a 3D tensor");
    TORCH_CHECK(x.is_contiguous(), "expected contiguous input");

    if (dim < 0) dim += 3;
    TORCH_CHECK(dim >= 0 && dim < 3, "invalid reduction dim");

    const int B  = (int)x.size(0);
    const int D1 = (int)x.size(1);
    const int D2 = (int)x.size(2);
    const float* xp = x.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (dim == 1) {
        auto out = torch::empty({B, D2}, x.options());
        dim3 block(256);
        dim3 grid((D2 + block.x - 1) / block.x, B);
        min_dim1_kernel<<<grid, block, 0, stream>>>(xp, out.data_ptr<float>(), B, D1, D2);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return out;
    } else if (dim == 0) {
        auto out = torch::empty({D1, D2}, x.options());
        dim3 block(256);
        dim3 grid((D2 + block.x - 1) / block.x, D1);
        min_dim0_kernel<<<grid, block, 0, stream>>>(xp, out.data_ptr<float>(), B, D1, D2);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return out;
    } else { // dim == 2
        auto out = torch::empty({B, D1}, x.options());
        dim3 block(256);
        dim3 grid(B * D1);
        min_dim2_kernel<<<grid, block, 256 * sizeof(float), stream>>>(xp, out.data_ptr<float>(), B, D1, D2);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return out;
    }
}
"""

_min_reduce_ext = load_inline(
    name="min_reduce_dim1_ext_fast",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["min_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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 _min_reduce_ext.min_forward(x, self.dim)