import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, scales the output, and then applies a minimum operation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height, width).
        """
        x = self.conv(x)
        x = x * self.scale_factor
        x = torch.min(x, dim=1, keepdim=True)[0]  # Minimum along channel dimension
        return x

batch_size = 64
in_channels = 64
out_channels = 128
height = width = 256
kernel_size = 3
scale_factor = 2.0

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scale_factor]

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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor bias_scale_min_forward(torch::Tensor x, torch::Tensor bias, double scale);
"""

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

__global__ void bias_min_c128_kernel(const float* __restrict__ x, const float* __restrict__ b,
                                     float* __restrict__ out, int total, int HW, float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int n = idx / HW;
    int s = idx - n * HW;
    int base = n * 128 * HW + s;
    if (scale >= 0.0f) {
        float mn = x[base] + b[0];
        #pragma unroll
        for (int c = 1; c < 128; ++c) {
            float v = x[base + c * HW] + b[c];
            mn = fminf(mn, v);
        }
        out[idx] = mn * scale;
    } else {
        float mx = x[base] + b[0];
        #pragma unroll
        for (int c = 1; c < 128; ++c) {
            float v = x[base + c * HW] + b[c];
            mx = fmaxf(mx, v);
        }
        out[idx] = mx * scale;
    }
}

__global__ void bias_min_kernel(const float* __restrict__ x, const float* __restrict__ b,
                                float* __restrict__ out, int total, int C, int HW, float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int n = idx / HW;
    int s = idx - n * HW;
    int base = n * C * HW + s;
    if (scale >= 0.0f) {
        float mn = x[base] + b[0];
        for (int c = 1; c < C; ++c) mn = fminf(mn, x[base + c * HW] + b[c]);
        out[idx] = mn * scale;
    } else {
        float mx = x[base] + b[0];
        for (int c = 1; c < C; ++c) mx = fmaxf(mx, x[base + c * HW] + b[c]);
        out[idx] = mx * scale;
    }
}

torch::Tensor bias_scale_min_forward(torch::Tensor x, torch::Tensor bias, double scale) {
    const int64_t N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    auto xc = x.contiguous();
    auto bc = bias.contiguous();
    auto out = torch::empty({N, 1, H, W}, xc.options());
    int total = (int)(N * H * W);
    int HW = (int)(H * W);
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if ((int)C == 128) {
        bias_min_c128_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>(), total, HW, (float)scale);
    } else {
        bias_min_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>(), total, (int)C, HW, (float)scale);
    }
    return out;
}
"""

_ext = load_inline(
    name="bias_min_ext_submit",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["bias_scale_min_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
    extra_cflags=["-O3"],
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor

    def forward(self, x):
        y = F.conv2d(x, self.conv.weight, None, self.conv.stride, self.conv.padding, self.conv.dilation, self.conv.groups)
        return _ext.bias_scale_min_forward(y, self.conv.bias, float(self.scale_factor))