import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Argmax over a specified dimension.
    """
    def __init__(self, dim: int):
        """
        Initializes the model with the dimension to perform argmax.

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

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

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

        Returns:
            torch.Tensor: Output tensor with argmax applied, with the specified dimension removed.
        """
        return torch.argmax(x, dim=self.dim)

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]

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

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

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cstdint>
#define NEG_INF_F (-3.4028234663852886e38f)

__global__ void argmax_dim1_main4_kernel(const float* __restrict__ x, int64_t* __restrict__ out,
                                         int B, int D1, int D2, int groups4) {
    int b = blockIdx.y;
    int group = blockIdx.x * blockDim.x + threadIdx.x;
    if (group >= groups4) return;
    int c = group << 2;
    const float* p = x + (long long)b * D1 * D2 + c;
    float m0 = p[0], m1 = p[1], m2 = p[2], m3 = p[3];
    int i0 = 0, i1 = 0, i2 = 0, i3 = 0;
    #pragma unroll 4
    for (int k = 1; k < D1; ++k) {
        p += D2;
        float v0 = p[0];
        float v1 = p[1];
        float v2 = p[2];
        float v3 = p[3];
        if (v0 > m0) { m0 = v0; i0 = k; }
        if (v1 > m1) { m1 = v1; i1 = k; }
        if (v2 > m2) { m2 = v2; i2 = k; }
        if (v3 > m3) { m3 = v3; i3 = k; }
    }
    int64_t* o = out + (long long)b * D2 + c;
    o[0] = (int64_t)i0; o[1] = (int64_t)i1; o[2] = (int64_t)i2; o[3] = (int64_t)i3;
}

__global__ void argmax_dim1_tail_kernel(const float* __restrict__ x, int64_t* __restrict__ out,
                                        int B, int D1, int D2, int startc) {
    int b = blockIdx.x;
    int c = startc + threadIdx.x;
    if (c >= D2) return;
    const float* p = x + (long long)b * D1 * D2 + c;
    float m = p[0]; int mi = 0;
    #pragma unroll 4
    for (int k = 1; k < D1; ++k) {
        p += D2;
        float v = p[0];
        if (v > m) { m = v; mi = k; }
    }
    out[(long long)b * D2 + c] = (int64_t)mi;
}

__global__ void argmax_dim2_float_kernel(const float* __restrict__ x, int64_t* __restrict__ out,
                                         int B, int D1, int D2) {
    int row = blockIdx.x;
    const float* p = x + (long long)row * D2;
    float best = NEG_INF_F; int besti = 0;
    for (int c = threadIdx.x; c < D2; c += blockDim.x) { float v = p[c]; if (v > best) { best = v; besti = c; } }
    __shared__ float svals[256]; __shared__ int sidx[256];
    svals[threadIdx.x] = best; sidx[threadIdx.x] = besti; __syncthreads();
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            float v = svals[threadIdx.x + stride]; int idx = sidx[threadIdx.x + stride];
            if (v > svals[threadIdx.x] || (v == svals[threadIdx.x] && idx < sidx[threadIdx.x])) { svals[threadIdx.x] = v; sidx[threadIdx.x] = idx; }
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) out[row] = (int64_t)sidx[0];
}

__global__ void argmax_dim0_float_kernel(const float* __restrict__ x, int64_t* __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 d1 = idx / D2; int d2 = idx - d1 * D2;
    float best = x[(long long)d1 * D2 + d2]; int besti = 0;
    for (int b = 1; b < B; ++b) { float v = x[((long long)b * D1 + d1) * D2 + d2]; if (v > best) { best = v; besti = b; } }
    out[idx] = (int64_t)besti;
}

torch::Tensor argmax_forward(torch::Tensor x, int64_t dim) {
    dim = (dim % 3 + 3) % 3;
    int B = (int)x.size(0), D1 = (int)x.size(1), D2 = (int)x.size(2);
    auto opts = torch::TensorOptions().device(x.device()).dtype(torch::kInt64);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (dim == 1) {
        torch::Tensor out = torch::empty({B, D2}, opts);
        int groups4 = D2 >> 2;
        const int threads = 128;
        if (groups4 > 0) {
            dim3 grid((groups4 + threads - 1) / threads, B);
            argmax_dim1_main4_kernel<<<grid, threads, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<int64_t>(), B, D1, D2, groups4);
        }
        int startc = groups4 << 2;
        if (startc < D2) {
            argmax_dim1_tail_kernel<<<B, 32, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<int64_t>(), B, D1, D2, startc);
        }
        return out;
    } else if (dim == 2) {
        torch::Tensor out = torch::empty({B, D1}, opts);
        argmax_dim2_float_kernel<<<B * D1, 256, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<int64_t>(), B, D1, D2);
        return out;
    } else {
        torch::Tensor out = torch::empty({D1, D2}, opts);
        int total = D1 * D2;
        argmax_dim0_float_kernel<<<(total + 255) / 256, 256, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<int64_t>(), B, D1, D2);
        return out;
    }
}
'''

_ext = load_inline(name='argmax_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=['argmax_forward'], verbose=False,
                   extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])

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.argmax_forward(x.contiguous(), self.dim)