import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix-scalar multiplication (C = A * s)
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, s: float) -> torch.Tensor:
        """
        Performs matrix-scalar multiplication.

        Args:
            A: Input matrix of shape (M, N)
            s: Scalar value

        Returns:
            C: Resulting matrix of shape (M, N)
        """
        return A * s

M = 16384 * 4
N = 4096 * 4

def get_inputs():
    A = torch.rand(M, N)
    s = 3.14
    return [A, s]

def get_init_inputs():
    return []  # No special initialization inputs needed

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor scale_forward(torch::Tensor A, double s);
'''

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

__global__ __launch_bounds__(256, 8)
void scale_float4x4_kernel(const float* __restrict__ in, float* __restrict__ out, float s, int n4) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    const float4* __restrict__ in4 = reinterpret_cast<const float4*>(in);
    float4* __restrict__ out4 = reinterpret_cast<float4*>(out);

    int groups = n4 >> 2;  // four float4 values per group = 16 floats
    for (int g = tid; g < groups; g += stride) {
        int i = g << 2;
        float4 v0 = in4[i];
        float4 v1 = in4[i + 1];
        float4 v2 = in4[i + 2];
        float4 v3 = in4[i + 3];
        v0.x *= s; v0.y *= s; v0.z *= s; v0.w *= s;
        v1.x *= s; v1.y *= s; v1.z *= s; v1.w *= s;
        v2.x *= s; v2.y *= s; v2.z *= s; v2.w *= s;
        v3.x *= s; v3.y *= s; v3.z *= s; v3.w *= s;
        out4[i] = v0;
        out4[i + 1] = v1;
        out4[i + 2] = v2;
        out4[i + 3] = v3;
    }
    int tail = groups << 2;
    for (int i = tail + tid; i < n4; i += stride) {
        float4 v = in4[i];
        v.x *= s; v.y *= s; v.z *= s; v.w *= s;
        out4[i] = v;
    }
}

__global__ __launch_bounds__(256, 8)
void scale_scalar_kernel(const float* __restrict__ in, float* __restrict__ out, float s, int64_t n) {
    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) {
        out[i] = in[i] * s;
    }
}

torch::Tensor scale_forward(torch::Tensor A, double s_double) {
    auto C = torch::empty_like(A);
    int64_t n = A.numel();
    if (n == 0) return C;

    constexpr int threads = 256;
    int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    float s = static_cast<float>(s_double);
    const float* in = A.data_ptr<float>();
    float* out = C.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    uintptr_t ai = reinterpret_cast<uintptr_t>(in);
    uintptr_t ci = reinterpret_cast<uintptr_t>(out);
    if (((ai | ci) & 15) == 0 && (n % 4 == 0) && (n / 4 <= INT_MAX)) {
        int n4 = static_cast<int>(n >> 2);
        int64_t work = (static_cast<int64_t>(n4) + 15) / 16;
        int blocks = static_cast<int>(work < static_cast<int64_t>(sms) * 8 ? work : static_cast<int64_t>(sms) * 8);
        if (blocks < 1) blocks = 1;
        scale_float4x4_kernel<<<blocks, threads, 0, stream>>>(in, out, s, n4);
    } else {
        int64_t need = (n + threads - 1) / threads;
        int blocks = static_cast<int>(need < static_cast<int64_t>(sms) * 8 ? need : static_cast<int64_t>(sms) * 8);
        if (blocks < 1) blocks = 1;
        scale_scalar_kernel<<<blocks, threads, 0, stream>>>(in, out, s, n);
    }
    return C;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
    def forward(self, A: torch.Tensor, s: float) -> torch.Tensor:
        return _ext.scale_forward(A, s)