# NEGATIVE (caching across calls, by input pointer).
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
__global__ void mul_leaky_kernel(const float* __restrict__ x, float* __restrict__ y, float m, float s, long n) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)gridDim.x * blockDim.x;
    for (; i < n; i += stride) { float v = x[i] * m; y[i] = v > 0.f ? v : v * s; }
}
torch::Tensor mul_leaky(torch::Tensor x, double m, double s) {
    auto y = torch::empty_like(x);
    long n = x.numel();
    int threads = 256; int blocks = (int)std::min<long>((n + threads - 1) / threads, 65535L * 8);
    mul_leaky_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), (float)m, (float)s, n);
    return y;
}
"""
cpp_src = "torch::Tensor mul_leaky(torch::Tensor x, double m, double s);"

ext = load_inline(name="eval_neg_cache_ptr", cpp_sources=cpp_src, cuda_sources=cuda_src, functions=["mul_leaky"])

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, multiplier, negative_slope):
        super().__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.multiplier = multiplier
        self.negative_slope = negative_slope

        self.cache = {}

    def forward(self, x):
        k = x.data_ptr()
        if k not in self.cache:
            self.cache[k] = ext.mul_leaky(F.linear(x, self.gemm.weight, self.gemm.bias).contiguous(), self.multiplier, self.negative_slope)
        return self.cache[k]