import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that finds the index of the minimum value along a specified dimension.
    """
    def __init__(self, dim: int):
        """
        Initializes the model with the dimension to perform argmin on.

        Args:
            dim (int): Dimension along which to find the minimum value.
        """
        super(Model, self).__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Finds the index of the minimum value along the specified dimension.

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

        Returns:
            torch.Tensor: Tensor containing the indices of the minimum values along the specified dimension.
        """
        return torch.argmin(x, dim=self.dim)

batch_size = 128
dim1 = 4096
dim2 = 4095
dim = 1

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

def get_init_inputs():
    return [dim]

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

_cpp_src = r"""
torch::Tensor argmin_dim1_cuda(torch::Tensor x);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstdint>

__global__ void argmin_dim1_kernel(const float* __restrict__ x, long long* __restrict__ out,
                                   int B, int D1, int D2) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int b = blockIdx.y;
    if (col >= D2 || b >= B) return;

    const float* base = x + ((long long)b * D1 * D2 + col);
    float minv = base[0];
    int mini = 0;
    // Strict less-than keeps the first occurrence, matching argmin tie behavior.
    #pragma unroll 4
    for (int i = 1; i < D1; ++i) {
        float v = base[(long long)i * D2];
        if (v < minv) {
            minv = v;
            mini = i;
        }
    }
    out[(long long)b * D2 + col] = (long long)mini;
}

torch::Tensor argmin_dim1_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 3, "x must be 3D");
    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);
    auto out = torch::empty({B, D2}, x.options().dtype(torch::kInt64));
    dim3 block(256);
    dim3 grid((D2 + block.x - 1) / block.x, B);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    argmin_dim1_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), (long long*)out.data_ptr<int64_t>(), B, D1, D2);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="argmin_dim1_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["argmin_dim1_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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.argmin_dim1_cuda(x)