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
from torch.utils.cpp_extension import load_inline

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

// Convert PyTorch conv weight layout [OC,C,3,3] to [C*3*3,OC].
// The transformed matrix is tiny (128*64*9 floats) and stays resident in cache
// during the fused convolution/min kernel.
__global__ void transform_w_3x3_128_kernel(const float* __restrict__ w,
                                           float* __restrict__ wt) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = 128 * 64 * 9;
    if (idx >= total) return;
    int rem = idx;
    int kw = rem % 3; rem /= 3;
    int kh = rem % 3; rem /= 3;
    int c  = rem % 64; rem /= 64;
    int oc = rem;
    int r = (c * 9 + kh * 3 + kw) * 128 + oc;
    wt[r] = w[idx];
}

__device__ __forceinline__ float warp_min(float v) {
    unsigned mask = 0xffffffffu;
    v = fminf(v, __shfl_down_sync(mask, v, 16));
    v = fminf(v, __shfl_down_sync(mask, v, 8));
    v = fminf(v, __shfl_down_sync(mask, v, 4));
    v = fminf(v, __shfl_down_sync(mask, v, 2));
    v = fminf(v, __shfl_down_sync(mask, v, 1));
    return v;
}

// One warp computes one output pixel.  Each lane computes four output channels
// (lane, lane+32, lane+64, lane+96); the warp then reduces the 128 channel
// values to the required channel-minimum.  The convolution, scale and reduction
// are fully fused, avoiding the 64*128*254*254 intermediate tensor and the
// separate multiply/min kernels in the reference.
__global__ __launch_bounds__(256, 2)
void conv3x3_min_scale_64_128_kernel(const float* __restrict__ x,
                                      const float* __restrict__ wt,
                                      const float* __restrict__ bias,
                                      float* __restrict__ out,
                                      int N, int H, int W,
                                      float scale) {
    const int lane = threadIdx.x & 31;
    const int warp_in_block = threadIdx.x >> 5;
    const int warps_per_block = blockDim.x >> 5;
    const int OH = H - 2;
    const int OW = W - 2;
    const int64_t num_pix = (int64_t)N * OH * OW;
    int64_t pix = (int64_t)blockIdx.x * warps_per_block + warp_in_block;
    if (pix >= num_pix) return;

    int ow = (int)(pix % OW);
    int t = (int)(pix / OW);
    int oh = t % OH;
    int n  = t / OH;

    float a0 = __ldg(bias + lane);
    float a1 = __ldg(bias + lane + 32);
    float a2 = __ldg(bias + lane + 64);
    float a3 = __ldg(bias + lane + 96);

    // Base pointer for this image. Input is NCHW contiguous.
    const int img_stride = 64 * H * W;
    const float* xN = x + (int64_t)n * img_stride;

    #pragma unroll 1
    for (int c = 0; c < 64; ++c) {
        const float* xc = xN + (int64_t)c * H * W + oh * W + ow;
        int wr = c * 9 * 128 + lane;

        float xv;
        xv = __ldg(xc + 0);
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xv = __ldg(xc + 1); wr += 128;
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xv = __ldg(xc + 2); wr += 128;
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xc += W; wr += 128;
        xv = __ldg(xc + 0);
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xv = __ldg(xc + 1); wr += 128;
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xv = __ldg(xc + 2); wr += 128;
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xc += W; wr += 128;
        xv = __ldg(xc + 0);
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xv = __ldg(xc + 1); wr += 128;
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);

        xv = __ldg(xc + 2); wr += 128;
        a0 = fmaf(xv, __ldg(wt + wr +   0), a0);
        a1 = fmaf(xv, __ldg(wt + wr +  32), a1);
        a2 = fmaf(xv, __ldg(wt + wr +  64), a2);
        a3 = fmaf(xv, __ldg(wt + wr +  96), a3);
    }

    a0 *= scale; a1 *= scale; a2 *= scale; a3 *= scale;
    float m = fminf(fminf(a0, a1), fminf(a2, a3));
    m = warp_min(m);
    if (lane == 0) out[pix] = m;
}

torch::Tensor fused_conv_min_forward(torch::Tensor x, torch::Tensor weight,
                                     torch::Tensor bias, double scale) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "all tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 is supported");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32 && bias.scalar_type() == torch::kFloat32, "only float32 weights are supported");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), "tensors must be contiguous");
    TORCH_CHECK(x.dim() == 4 && weight.dim() == 4, "bad tensor rank");
    TORCH_CHECK(x.size(1) == 64 && weight.size(0) == 128 && weight.size(1) == 64 && weight.size(2) == 3 && weight.size(3) == 3,
                "this optimized kernel expects C=64, OC=128, K=3");

    const int N = (int)x.size(0);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    auto wt = torch::empty({64 * 9, 128}, x.options());
    auto out = torch::empty({N, 1, H - 2, W - 2}, x.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    transform_w_3x3_128_kernel<<<(128*64*9 + 255) / 256, 256, 0, stream>>>(
        weight.data_ptr<float>(), wt.data_ptr<float>());

    const int threads = 256;
    const int warps = threads / 32;
    int64_t num_pix = (int64_t)N * (H - 2) * (W - 2);
    int blocks = (int)((num_pix + warps - 1) / warps);
    conv3x3_min_scale_64_128_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), wt.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(),
        N, H, W, (float)scale);
    return out;
}
'''

_ext = load_inline(
    name='fused_conv_min_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['fused_conv_min_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

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 = float(scale_factor)

    def forward(self, x):
        return _ext.fused_conv_min_forward(x.contiguous(), self.conv.weight.contiguous(), self.conv.bias.contiguous(), self.scale_factor)