import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a transposed 3D convolution, clamps the output to a minimum value, 
    and then divides the result by a constant.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, min_value, divisor):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.min_value = min_value
        self.divisor = divisor

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.clamp(x, min=self.min_value)
        x = x / self.divisor
        return x

batch_size = 16
in_channels = 64
out_channels = 128
depth, height, width = 24, 48, 48
kernel_size = 3
stride = 2
padding = 1
min_value = -1.0
divisor = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, min_value, divisor]

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


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

torch::Tensor clamp_div_inplace(torch::Tensor x, double min_value, double divisor);
"""

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

__global__ void clamp_div_scalar_kernel(float* __restrict__ p, int64_t n, float minv, float inv_div) {
    int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = tid; i < n; i += stride) {
        float v = p[i];
        v = fmaxf(v, minv) * inv_div;
        p[i] = v;
    }
}

__global__ void clamp_div_float4_kernel(float4* __restrict__ p, int64_t n4, float minv, float inv_div) {
    int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = tid; i < n4; i += stride) {
        float4 v = p[i];
        v.x = fmaxf(v.x, minv) * inv_div;
        v.y = fmaxf(v.y, minv) * inv_div;
        v.z = fmaxf(v.z, minv) * inv_div;
        v.w = fmaxf(v.w, minv) * inv_div;
        p[i] = v;
    }
}

torch::Tensor clamp_div_inplace(torch::Tensor x, double min_value, double divisor) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    const int64_t n = x.numel();
    if (n == 0) return x;

    float* ptr = x.data_ptr<float>();
    const float minv = static_cast<float>(min_value);
    const float inv_div = 1.0f / static_cast<float>(divisor);
    constexpr int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
    if (((addr & 15) == 0) && ((n & 3) == 0)) {
        int64_t n4 = n >> 2;
        int blocks = (int)((n4 + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        clamp_div_float4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(ptr), n4, minv, inv_div);
    } else {
        int blocks = (int)((n + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        clamp_div_scalar_kernel<<<blocks, threads, 0, stream>>>(ptr, n, minv, inv_div);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return x;
}
"""

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


class ModelNew(nn.Module):
    """
    ConvTranspose3d followed by fused in-place clamp(min=...) and division.
    The ConvTranspose3d module is kept with the same name/layout as the reference
    so state_dict parameter copying is identical.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, min_value, divisor):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.min_value = min_value
        self.divisor = divisor

    def forward(self, x):
        y = self.conv_transpose(x)
        return _ext.clamp_div_inplace(y, float(self.min_value), float(self.divisor))