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

cpp_src = r"""
#include <torch/extension.h>
torch::Tensor diag_mul_cuda(torch::Tensor A, torch::Tensor B);
"""

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void row_scale_vec4_kernel(const float* __restrict__ A,
                                      const float4* __restrict__ B4,
                                      float4* __restrict__ C4,
                                      int M4) {
    int row = blockIdx.x;
    float a = A[row];
    int base = row * M4;
    for (int j = threadIdx.x; j < M4; j += blockDim.x) {
        float4 b = B4[base + j];
        b.x *= a;
        b.y *= a;
        b.z *= a;
        b.w *= a;
        C4[base + j] = b;
    }
}

__global__ void row_scale_scalar_kernel(const float* __restrict__ A,
                                        const float* __restrict__ B,
                                        float* __restrict__ C,
                                        int M) {
    int row = blockIdx.x;
    float a = A[row];
    int base = row * M;
    for (int col = threadIdx.x; col < M; col += blockDim.x) {
        C[base + col] = a * B[base + col];
    }
}

torch::Tensor diag_mul_cuda(torch::Tensor A, torch::Tensor B) {
    auto C = torch::empty_like(B);
    const int N = (int)B.size(0);
    const int M = (int)B.size(1);
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if ((M & 3) == 0) {
        int M4 = M >> 2;
        row_scale_vec4_kernel<<<N, threads, 0, stream>>>(
            A.data_ptr<float>(),
            reinterpret_cast<const float4*>(B.data_ptr<float>()),
            reinterpret_cast<float4*>(C.data_ptr<float>()),
            M4);
    } else {
        row_scale_scalar_kernel<<<N, threads, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M);
    }
    return C;
}
"""

_ext = load_inline(name="diag_mul_ext_final",
                   cpp_sources=[cpp_src],
                   cuda_sources=[cuda_src],
                   functions=["diag_mul_cuda"],
                   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, B):
        return _ext.diag_mul_cuda(A, B)