```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. The fused
# kernel computes the result directly in O(N * M), which is purely
# memory-bound (read B once, write C once).
#
# This version improves on the previous single-configuration kernel by
# providing a family of vectorized (float4) kernel configurations that trade
# off block size, per-thread unroll (memory-level parallelism) and cache
# eviction policy (streaming / write-through / plain). At run time, on the
# first forward call for a given problem shape, a lightweight autotuner
# benchmarks every variant on the live tensors with CUDA events and caches
# the fastest configuration, so the kernel adapts to whatever GPU it runs on.
# ---------------------------------------------------------------------------

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

namespace {

// Cache-policy modes for the load/store path.
constexpr int kPlain  = 0;  // __ldg load, plain store
constexpr int kStream = 1;  // __ldcs load, __stcs store (evict-first streaming)
constexpr int kWt     = 2;  // __ldcs load, __stwt store (write-through)

template <int Mode>
__device__ __forceinline__ float4 load4(const float4* __restrict__ p) {
    return (Mode == kPlain) ? __ldg(p) : __ldcs(p);
}

template <int Mode>
__device__ __forceinline__ void store4(float4* __restrict__ p, const float4 v) {
    if (Mode == kPlain) {
        *p = v;
    } else if (Mode == kStream) {
        __stcs(p, v);
    } else {
        __stwt(p, v);
    }
}

__device__ __forceinline__ float4 scale4(float4 v, const float a) {
    v.x *= a; v.y *= a; v.z *= a; v.w *= a;
    return v;
}

// Vectorized (float4) kernel: C[row, :] = A[row] * B[row, :].
// Each block covers a (kThreads * kUnroll)-float4 tile of one row; all
// kUnroll loads are issued back-to-back for deep memory-level parallelism
// before any store, so DRAM can be saturated with relatively few threads.
template <int kThreads, int kUnroll, int kMode>
__global__ void __launch_bounds__(kThreads)
diag_scale_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)
    constexpr int kTile = kThreads * kUnroll;
    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 float4* __restrict__ Brow =
            reinterpret_cast<const float4*>(B) + (int64_t)row * Mv;
        float4* __restrict__ Crow =
            reinterpret_cast<float4*>(C) + (int64_t)row * Mv;

        if (full) {
            const int base = (int)start + tx;
            float4 v[kUnroll];
            #pragma unroll
            for (int u = 0; u < kUnroll; ++u) {
                v[u] = load4<kMode>(Brow + base + u * kThreads);
            }
            #pragma unroll
            for (int u = 0; u < kUnroll; ++u) {
                store4<kMode>(Crow + base + u * kThreads, scale4(v[u], a));
            }
        } else {
            for (int i = (int)start + tx; i < Mv; i += kThreads) {
                float4 v = load4<kMode>(Brow + i);
                store4<kMode>(Crow + i, scale4(v, a));
            }
        }
    }
}

// Scalar fallback kernel (any M, any alignment), same tiling idea.
template <int kThreads, int kUnroll>
__global__ void __launch_bounds__(kThreads)
diag_scale_scalar_kernel(const float* __restrict__ A,
                         const float* __restrict__ B,
                         float* __restrict__ C,
                         const int N,
                         const int M) {
    constexpr int kTile = kThreads * kUnroll;
    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 float* __restrict__ Brow = B + (int64_t)row * M;
        float* __restrict__ Crow = C + (int64_t)row * M;

        if (full) {
            const int base = (int)start + tx;
            float v[kUnroll];
            #pragma unroll
            for (int u = 0; u < kUnroll; ++u) {
                v[u] = __ldcs(Brow + base + u * kThreads);
            }
            #pragma unroll
            for (int u = 0; u < kUnroll; ++u) {
                __stcs(Crow + base + u * kThreads, v[u] * 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, int64_t variant) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda() &&
                A.dim() == 1 && B.dim() == 2 &&
                A.size(0) == B.size(0) &&
                A.device() == B.device() &&
                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 by = (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) {
        constexpr int T = 256, U = 4;
        const unsigned bx = (unsigned)((M + T * U - 1) / (T * U));
        dim3 grid(bx, by);
        diag_scale_scalar_kernel<T, U><<<grid, T, 0, stream>>>(
            A_ptr, B_ptr, C_ptr, (int)N, (int)M);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return C;
    }

    const int Mv = (int)(M >> 2);
    auto launch4 = [&](auto kern, int threads, int unroll) {
        const unsigned bx =
            (unsigned)((Mv + threads * unroll - 1) / (threads * unroll));
        dim3 grid(bx, by);
        kern<<<grid, threads, 0, stream>>>(A_ptr, B_ptr, C_ptr, (int)N, Mv);
    };

    switch (variant) {
        case 1:  launch4(diag_scale_vec4_kernel<256,  8, kStream>,  256,  8); break;
        case 2:  launch4(diag_scale_vec4_kernel<128,  8, kStream>,  128,  8); break;
        case 3:  launch4(diag_scale_vec4_kernel<512,  8, kStream>,  512,  8); break;
        case 4:  launch4(diag_scale_vec4_kernel<256,  4, kPlain >,  256,  4); break;
        case 5:  launch4(diag_scale_vec4_kernel<256,  8, kPlain >,  256,  8); break;
        case 6:  launch4(diag_scale_vec4_kernel<256,  8, kWt    >,  256,  8); break;
        case 7:  launch4(diag_scale_vec4_kernel<128,  4, kStream>,  128,  4); break;
        case 8:  launch4(diag_scale_vec4_kernel< 64, 16, kStream>,   64, 16); break;
        case 9:  launch4(diag_scale_vec4_kernel<512,  4, kStream>,  512,  4); break;
        case 10: launch4(diag_scale_vec4_kernel<1024, 1, kStream>, 1024,  1); break;
        case 11: launch4(diag_scale_vec4_kernel<256,  2, kStream>,  256,  2); break;
        case 0:
        default: launch4(diag_scale_vec4_kernel<256,  4, kStream>,  256,  4); break;
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

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

# Compile the inline CUDA kernels.
diag_matmul = load_inline(
    name="diag_matmul_fused_v4",
    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
_NUM_VARIANTS = 12
_TUNE_REPS = 5


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.

    Several kernel configurations (block size, per-thread unroll, cache
    eviction policy) are compiled, and the fastest one for the actual
    problem shape / GPU is selected by a one-time CUDA-event autotune on the
    first forward call; the choice is cached per (device, N, M).
    """

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

    def _pick_variant(self, A, B):
        """One-time micro-benchmark of all kernel variants on the live data."""
        try:
            # Warm every variant once (allocator warm-up, fault in code path).
            for v in range(_NUM_VARIANTS):
                self._op(A, B, v)
            torch.cuda.synchronize(A.device)

            best_variant = 0
            best_time = None
            for v in range(_NUM_VARIANTS):
                try:
                    start = torch.cuda.Event(enable_timing=True)
                    end = torch.cuda.Event(enable_timing=True)
                    start.record()
                    for _ in range(_TUNE_REPS):
                        self._op(A, B, v)
                    end.record()
                    torch.cuda.synchronize(A.device)
                    t = start.elapsed_time(end)
                except Exception:
                    continue
                if best_time is None or t < best_time:
                    best_time = t
                    best_variant = v
            return best_variant
        except Exception:
            # Autotuning unsupported in this context (e.g. stream capture);
            # variant 0 is a solid default on all architectures.
            return 0

    def forward(self, A, B):
        if (A.is_cuda and B.is_cuda and A.dim() == 1 and B.dim() == 2 and
                A.dtype == torch.float32 and B.dtype == torch.float32):
            key = (A.get_device(), B.shape[0], B.shape[1])
            variant = self._variant_cache.get(key)
            if variant is None:
                variant = self._pick_variant(A, B)
                self._variant_cache[key] = variant
            try:
                return self._op(A, B, variant)
            except Exception:
                pass
        # Generic fallback (identical semantics to torch.diag(A) @ B).
        return A.unsqueeze(-1) * B
```