import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    """
    Simple model that performs a HardTanh activation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies HardTanh activation to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of any shape.

        Returns:
            torch.Tensor: Output tensor with HardTanh applied, same shape as input.
        """
        return F.hardtanh(x, min_val=-1., max_val=1.)

batch_size = 4096
dim = 393216

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

def get_init_inputs():
    return []  # No special initialization inputs needed

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


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor hardtanh_forward_cuda(torch::Tensor x);

torch::Tensor hardtanh_forward(torch::Tensor x) {
    return hardtanh_forward_cuda(x);
}
"""

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

static __device__ __forceinline__ float ht_clamp(float v) {
    // Comparisons (rather than fmin/fmax) preserve NaNs like torch.clamp/hardtanh.
    if (v > 1.0f) return 1.0f;
    if (v < -1.0f) return -1.0f;
    return v;
}

__global__ __launch_bounds__(256, 4)
void hardtanh_scalar_kernel(const float* __restrict__ x, float* __restrict__ y, int64_t n) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < n; idx += stride) {
        y[idx] = ht_clamp(x[idx]);
    }
}

__global__ __launch_bounds__(256, 4)
void hardtanh_vec4_kernel(const float4* __restrict__ x4, float4* __restrict__ y4, int64_t n4) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < n4; idx += stride) {
        float4 v = x4[idx];
        v.x = ht_clamp(v.x);
        v.y = ht_clamp(v.y);
        v.z = ht_clamp(v.z);
        v.w = ht_clamp(v.w);
        y4[idx] = v;
    }
}

__global__ __launch_bounds__(256, 4)
void hardtanh_tail_kernel(const float* __restrict__ x, float* __restrict__ y, int64_t start, int64_t n) {
    int64_t i = start + (int64_t)threadIdx.x;
    if (i < n) y[i] = ht_clamp(x[i]);
}

torch::Tensor hardtanh_forward_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");

    // The benchmark input is contiguous; make other layouts correct as well.
    if (!x.is_contiguous()) {
        x = x.contiguous();
    }

    auto y = torch::empty_like(x);
    const int64_t n = x.numel();
    if (n == 0) return y;

    const float* xp = x.data_ptr<float>();
    float* yp = y.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int threads = 256;
    // Keep enough CTAs to saturate the device without creating an excessive launch grid.
    int device = x.get_device();
    int sms = 0;
    cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device);
    int max_blocks = sms > 0 ? sms * 8 : 1024;

    uintptr_t ax = reinterpret_cast<uintptr_t>(xp);
    uintptr_t ay = reinterpret_cast<uintptr_t>(yp);
    if (((ax | ay) & 15) == 0 && n >= 4) {
        int64_t n4 = n >> 2;
        int blocks = (int)((n4 + threads - 1) / threads);
        if (blocks > max_blocks) blocks = max_blocks;
        if (blocks < 1) blocks = 1;
        hardtanh_vec4_kernel<<<blocks, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(xp), reinterpret_cast<float4*>(yp), n4);
        int64_t start = n4 << 2;
        if (start < n) {
            hardtanh_tail_kernel<<<1, 256, 0, stream>>>(xp, yp, start, n);
        }
    } else {
        int blocks = (int)((n + threads - 1) / threads);
        if (blocks > max_blocks) blocks = max_blocks;
        if (blocks < 1) blocks = 1;
        hardtanh_scalar_kernel<<<blocks, threads, 0, stream>>>(xp, yp, n);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_ext = load_inline(
    name="hardtanh_inline_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["hardtanh_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.hardtanh_forward(x)