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 scalar_mul_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__ void scalar_mul4_kernel(const float4* __restrict__ in, float4* __restrict__ out, unsigned int n4, float s) {
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int stride = blockDim.x * gridDim.x;
    for (unsigned int i = idx; i < n4; i += stride) {
        float4 v = in[i];
        v.x *= s;
        v.y *= s;
        v.z *= s;
        v.w *= s;
        out[i] = v;
    }
}

__global__ void scalar_mul_kernel(const float* __restrict__ in, float* __restrict__ out, unsigned int n, float s) {
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int stride = blockDim.x * gridDim.x;
    for (unsigned int i = idx; i < n; i += stride) {
        out[i] = in[i] * s;
    }
}

torch::Tensor scalar_mul_forward(torch::Tensor A, double s) {
    auto C = torch::empty_like(A);
    unsigned long long nn = (unsigned long long)A.numel();
    if (nn == 0) return C;
    const int threads = 1024;
    float sf = (float)s;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* in = A.data_ptr<float>();
    float* out = C.data_ptr<float>();
    uintptr_t ai = (uintptr_t)in;
    uintptr_t ao = (uintptr_t)out;
    if (((ai | ao) & 15) == 0 && ((nn & 3ull) == 0)) {
        unsigned int n4 = (unsigned int)(nn >> 2);
        int blocks = (int)((n4 + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        scalar_mul4_kernel<<<blocks, threads, 0, stream>>>((const float4*)in, (float4*)out, n4, sf);
    } else {
        unsigned int n = (unsigned int)nn;
        int blocks = (int)((n + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        scalar_mul_kernel<<<blocks, threads, 0, stream>>>(in, out, n, sf);
    }
    return C;
}
"""

_ext = load_inline(
    name="scalar_mul_ext_final",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["scalar_mul_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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.scalar_mul_forward(A, s)