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"""
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_mul_vec4_kernel(const float* __restrict__ A,
                                       float* __restrict__ C,
                                       float s,
                                       int64_t n4) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    const float4* __restrict__ Av = reinterpret_cast<const float4*>(A);
    float4* __restrict__ Cv = reinterpret_cast<float4*>(C);
    for (int64_t i = idx; i < n4; i += stride) {
        float4 v = Av[i];
        v.x *= s;
        v.y *= s;
        v.z *= s;
        v.w *= s;
        Cv[i] = v;
    }
}

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

torch::Tensor scalar_mul_forward(torch::Tensor A, double s) {
    auto C = torch::empty_like(A);
    const int64_t n = A.numel();
    if (n == 0) return C;
    const int threads = 256;
    int device = A.get_device();
    int sms = 0;
    cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device);
    // Enough blocks to saturate memory bandwidth without excessive launch overhead.
    int blocks = sms * 8;
    if (blocks < 1) blocks = 1;
    int64_t work = n;
    const float* Ap = A.data_ptr<float>();
    float* Cp = C.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    uintptr_t au = reinterpret_cast<uintptr_t>(Ap);
    uintptr_t cu = reinterpret_cast<uintptr_t>(Cp);
    if (((au | cu) & 15ULL) == 0ULL && (n % 4 == 0)) {
        work = n / 4;
        int64_t max_blocks = (work + threads - 1) / threads;
        if ((int64_t)blocks > max_blocks) blocks = (int)max_blocks;
        scalar_mul_vec4_kernel<<<blocks, threads, 0, stream>>>(Ap, Cp, (float)s, work);
    } else {
        int64_t max_blocks = (work + threads - 1) / threads;
        if ((int64_t)blocks > max_blocks) blocks = (int)max_blocks;
        scalar_mul_kernel<<<blocks, threads, 0, stream>>>(Ap, Cp, (float)s, work);
    }
    return C;
}
"""

_ext = load_inline(
    name="scalar_mul_ext_v1",
    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, float(s))