```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# ---------------------------------------------------------------------------
# Fused custom CUDA kernel for C = diag(A) @ B.
#
# Key insight: (diag(A) @ B)[i, j] = A[i] * B[i, j], i.e. the operation is a
# per-row scaling of B. The naive implementation materializes an (N, N)
# diagonal matrix (O(N^2) memory) and runs an O(N^2 * M) GEMM. This kernel
# computes the result directly in O(N * M), which is purely memory-bound.
#
# Performance features:
#   * 128-bit vectorized (float4) loads/stores when M % 4 == 0.
#   * Cache-streaming hints (__ldcs / __stcs) so the one-touch data does not
#     thrash L2, improving effective DRAM bandwidth.
#   * Persistent-style launch: grid sized to fill the machine exactly once
#     (SM_count * 8 blocks of 256 threads), minimizing tail effects while
#     keeping full occupancy for latency hiding.
#   * Minimal host-side overhead (no dispatcher-heavy calls on the hot path).
# ---------------------------------------------------------------------------

diag_matmul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <climits>
#include <cstdint>

namespace {

constexpr int kThreads = 256;

// Vectorized (float4) streaming kernel: C[row, :] = A[row] * B[row, :].
__global__ void __launch_bounds__(kThreads)
diag_matmul_vec4_kernel(const float* __restrict__ A,
                        const float* __restrict__ B,
                        float* __restrict__ C,
                        const int N,
                        const int Mv) {
    for (int row = blockIdx.y; row < N; row += gridDim.y) {
        const float a = __ldg(A + row);
        const int64_t off = (int64_t)row * Mv;
        const float4* __restrict__ Brow =
            reinterpret_cast<const float4*>(B) + off;
        float4* __restrict__ Crow = reinterpret_cast<float4*>(C) + off;
        for (int i = blockIdx.x * kThreads + threadIdx.x;
             i < Mv;
             i += gridDim.x * kThreads) {
            float4 b = __ldcs(Brow + i);
            b.x *= a;
            b.y *= a;
            b.z *= a;
            b.w *= a;
            __stcs(Crow + i, b);
        }
    }
}

// Scalar fallback kernel (any M, any alignment).
__global__ void __launch_bounds__(kThreads)
diag_matmul_scalar_kernel(const float* __restrict__ A,
                          const float* __restrict__ B,
                          float* __restrict__ C,
                          const int N,
                          const int M) {
    for (int row = blockIdx.y; row < N; row += gridDim.y) {
        const float a = __ldg(A + row);
        const int64_t off = (int64_t)row * M;
        for (int i = blockIdx.x * kThreads + threadIdx.x;
             i < M;
             i += gridDim.x * kThreads) {
            __stcs(C + off + i, a * __ldcs(B + off + i));
        }
    }
}

}  // namespace

torch::Tensor diag_matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda() &&
                A.dim() == 1 && B.dim() == 2 &&
                A.size(0) == B.size(0) &&
                A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "diag_matmul_cuda expects float32 CUDA tensors A(N,) and B(N, M)");

    const c10::cuda::OptionalCUDAGuard guard(B.device());

    // Cheap flag checks first; only pay for a dispatch copy when needed.
    const at::Tensor A_c = A.is_contiguous() ? A : A.contiguous();
    const at::Tensor B_c = B.is_contiguous() ? B : B.contiguous();
    at::Tensor C = at::empty({B_c.size(0), B_c.size(1)}, B_c.options());

    const int64_t N = B_c.size(0);
    const int64_t M = B_c.size(1);
    if (N == 0 || M == 0) {
        return C;
    }
    TORCH_CHECK(N <= INT_MAX && M <= INT_MAX, "Tensor dimensions too large");

    const float* A_ptr = A_c.data_ptr<float>();
    const float* B_ptr = B_c.data_ptr<float>();
    float* C_ptr = C.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int max_blocks = std::max(1, sms * 8);

    const bool vec_ok = ((M & 3) == 0) &&
        (((reinterpret_cast<uintptr_t>(B_ptr) |
           reinterpret_cast<uintptr_t>(C_ptr)) & 0xF) == 0);

    if (vec_ok) {
        const int Mv = (int)(M >> 2);
        const int pieces = (Mv + kThreads - 1) / kThreads;
        const int blocks_x = std::min(pieces, max_blocks);
        const int blocks_y =
            std::min((int)N, std::max(1, max_blocks / blocks_x));
        dim3 grid(blocks_x, blocks_y);
        diag_matmul_vec4_kernel<<<grid, kThreads, 0, stream>>>(
            A_ptr, B_ptr, C_ptr, (int)N, Mv);
    } else {
        const int pieces = (int)((M + kThreads - 1) / kThreads);
        const int blocks_x = std::min(pieces, max_blocks);
        const int blocks_y =
            std::min((int)N, std::max(1, max_blocks / blocks_x));
        dim3 grid(blocks_x, blocks_y);
        diag_matmul_scalar_kernel<<<grid, kThreads, 0, stream>>>(
            A_ptr, B_ptr, C_ptr, (int)N, (int)M);
    }
    return C;
}
"""

diag_matmul_cpp_source = (
    "torch::Tensor diag_matmul_cuda(torch::Tensor A, torch::Tensor B);"
)

# Compile the inline CUDA kernel.
diag_matmul = load_inline(
    name="diag_matmul_fused_v2",
    cpp_sources=diag_matmul_cpp_source,
    cuda_sources=diag_matmul_cuda_source,
    functions=["diag_matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)

_diag_matmul_cuda = diag_matmul.diag_matmul_cuda


class ModelNew(nn.Module):
    """
    Optimized version of Model that computes C = diag(A) @ B.

    Instead of building an (N, N) diagonal matrix and performing an
    O(N^2 * M) matrix multiplication, this module uses a fused custom CUDA
    kernel that directly evaluates C[i, j] = A[i] * B[i, j] in O(N * M),
    which is purely memory-bound and avoids materializing the diagonal
    matrix entirely.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _diag_matmul_cuda

    def forward(self, A, B):
        # Fast path: the custom kernel validates its own inputs; if anything
        # is unsupported (wrong dtype/device/shape) it raises and we fall
        # back to the mathematically identical broadcast multiply.
        try:
            return self._op(A, B)
        except Exception:
            # Generic fallback (identical semantics to torch.diag(A) @ B).
            return A.unsqueeze(-1) * B
```