import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication of a diagonal matrix with another matrix.
    C = diag(A) * B
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A, B):
        """
        Performs the matrix multiplication.

        Args:
            A (torch.Tensor): A 1D tensor representing the diagonal of the diagonal matrix. Shape: (N,).
            B (torch.Tensor): A 2D tensor representing the second matrix. Shape: (N, M).

        Returns:
            torch.Tensor: The result of the matrix multiplication. Shape: (N, M).
        """
        # Logically equivalent to torch.diag(A) @ B 
        # more efficient as no need to materialize a full N×N matrix
        return A.unsqueeze(1) * B

M = 4096
N = 4096

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

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

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

__global__ __launch_bounds__(256) void diag_mul_vec4_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int rows,
    int cols4)
{
    int r = blockIdx.x;
    if (r >= rows) return;
    float a = A[r];
    const float4* __restrict__ B4 = reinterpret_cast<const float4*>(B + ((long long)r * cols4 * 4));
    float4* __restrict__ C4 = reinterpret_cast<float4*>(C + ((long long)r * cols4 * 4));

    for (int c = threadIdx.x; c < cols4; c += blockDim.x) {
        float4 v = B4[c];
        v.x *= a;
        v.y *= a;
        v.z *= a;
        v.w *= a;
        C4[c] = v;
    }
}

__global__ __launch_bounds__(256) void diag_mul_scalar_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    long long total,
    int cols)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int r = (int)(idx / cols);
        C[idx] = A[r] * B[idx];
    }
}

torch::Tensor diag_mul_forward(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 && B.scalar_type() == torch::kFloat32, "float32 tensors expected");
    TORCH_CHECK(A.dim() == 1 && B.dim() == 2, "expected A[N], B[N,M]");
    TORCH_CHECK(A.size(0) == B.size(0), "shape mismatch");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(), "contiguous tensors expected");

    auto C = torch::empty_like(B);
    int rows = (int)B.size(0);
    int cols = (int)B.size(1);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((cols & 3) == 0) {
        int cols4 = cols >> 2;
        diag_mul_vec4_kernel<<<rows, 256, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), rows, cols4);
    } else {
        long long total = (long long)rows * cols;
        int blocks = (int)((total + 255) / 256);
        if (blocks > 65535) blocks = 65535;
        diag_mul_scalar_kernel<<<blocks, 256, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), total, cols);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
'''

_ext = load_inline(
    name='diag_mul_inline_ext_fast',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['diag_mul_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, A, B):
        return _ext.diag_mul_forward(A, B)