```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 vs. the previous version:
#   * 4x unrolled, grouped 128-bit (float4) loads: each thread issues four
#     independent __ldcs loads back-to-back before any store, giving deep
#     memory-level parallelism so far fewer resident threads are needed to
#     saturate DRAM.
#   * Tile-based mapping: one block covers a 256-thread x 4 tile of a row
#     (1024 float4s = exactly one 4096-wide row for the benchmark shape),
#     with the row index taken from blockIdx (no per-element division) and
#     the row scalar A[row] fetched once per block via __ldg.
#   * Streaming cache hints (__ldcs / __stcs): the single-touch data uses
#     evict-first policy so it does not thrash L2.
#   * Uniform fast path with no per-thread bounds checks when the tile is
#     full; a clean scalar tail handles partial tiles / odd shapes.
# ---------------------------------------------------------------------------

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;
constexpr int kUnroll  = 4;
constexpr int kTile    = kThreads * kUnroll;  // elements covered per block tile

// Vectorized (float4) kernel: C[row, :] = A[row] * B[row, :].
// Each block scales a kTile-float4 tile of one row; the four loads are
// issued together for maximal memory-level parallelism.
__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) {  // Mv = M / 4 (float4s per row)
    const int tx = threadIdx.x;
    const int64_t start = (int64_t)blockIdx.x * kTile;
    const bool full = (start + kTile) <= (int64_t)Mv;

    for (int row = blockIdx.y; row < N; row += gridDim.y) {
        const float a = __ldg(A + row);
        const int64_t row_off = (int64_t)row * Mv;
        const float4* __restrict__ Brow =
            reinterpret_cast<const float4*>(B) + row_off;
        float4* __restrict__ Crow =
            reinterpret_cast<float4*>(C) + row_off;

        if (full) {
            const int base = (int)start + tx;
            float4 v0 = __ldcs(Brow + base);
            float4 v1 = __ldcs(Brow + base + kThreads);
            float4 v2 = __ldcs(Brow + base + 2 * kThreads);
            float4 v3 = __ldcs(Brow + base + 3 * kThreads);
            v0.x *= a; v0.y *= a; v0.z *= a; v0.w *= a;
            v1.x *= a; v1.y *= a; v1.z *= a; v1.w *= a;
            v2.x *= a; v2.y *= a; v2.z *= a; v2.w *= a;
            v3.x *= a; v3.y *= a; v3.z *= a; v3.w *= a;
            __stcs(Crow + base, v0);
            __stcs(Crow + base + kThreads, v1);
            __stcs(Crow + base + 2 * kThreads, v2);
            __stcs(Crow + base + 3 * kThreads, v3);
        } else {
            for (int i = (int)start + tx; i < Mv; i += kThreads) {
                float4 v = __ldcs(Brow + i);
                v.x *= a; v.y *= a; v.z *= a; v.w *= a;
                __stcs(Crow + i, v);
            }
        }
    }
}

// Scalar fallback kernel (any M, any alignment), same tiling idea.
__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) {
    const int tx = threadIdx.x;
    const int64_t start = (int64_t)blockIdx.x * kTile;
    const bool full = (start + kTile) <= (int64_t)M;

    for (int row = blockIdx.y; row < N; row += gridDim.y) {
        const float a = __ldg(A + row);
        const int64_t row_off = (int64_t)row * M;
        const float* __restrict__ Brow = B + row_off;
        float* __restrict__ Crow = C + row_off;

        if (full) {
            const int base = (int)start + tx;
            float v0 = __ldcs(Brow + base);
            float v1 = __ldcs(Brow + base + kThreads);
            float v2 = __ldcs(Brow + base + 2 * kThreads);
            float v3 = __ldcs(Brow + base + 3 * kThreads);
            __stcs(Crow + base, v0 * a);
            __stcs(Crow + base + kThreads, v1 * a);
            __stcs(Crow + base + 2 * kThreads, v2 * a);
            __stcs(Crow + base + 3 * kThreads, v3 * a);
        } else {
            for (int i = (int)start + tx; i < M; i += kThreads) {
                __stcs(Crow + i, a * __ldcs(Brow + 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 unsigned blocks_y =
        (unsigned)std::min<int64_t>(N, 65535);
    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 unsigned blocks_x = (unsigned)((Mv + kTile - 1) / kTile);
        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 unsigned blocks_x = (unsigned)((M + kTile - 1) / kTile);
        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_v3",
    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
```