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

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void chebyshev_abs_square_kernel(const float* __restrict__ x, float* __restrict__ y, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float val = x[idx];
        float x2 = val * val;

        float p = fmaf(16.0f, x2, -20.0f);
        p = fmaf(p, x2, 5.0f);
        p = p * val;

        float abs_p = fabsf(p);
        y[idx] = abs_p * abs_p;
    }
}

torch::Tensor launch_chebyshev_abs_square(torch::Tensor x) {
    auto n = x.numel();
    auto y = torch::empty_like(x);

    const int threads = 256;
    const int blocks = (n + threads - 1) / threads;

    chebyshev_abs_square_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), n);
    return y;
}
"""

cpp_source = """
torch::Tensor launch_chebyshev_abs_square(torch::Tensor x);
"""

chebyshev_abs_square_module = load_inline(
    name='chebyshev_abs_square_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_chebyshev_abs_square'],
    verbose=False
)


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.op.launch_chebyshev_abs_square(x.contiguous())