```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

# Custom CUDA kernel for matrix-scalar multiplication (C = A * s).
#
# This operator is purely memory-bandwidth bound (read 4 GiB, write 4 GiB),
# so this version is tuned to match/beat the fastest known streaming kernels:
#   * 128-bit vectorized loads/stores (float4), plain ld/st (no cache hints
#     that can perturb the L2 write stream)
#   * 32-bit indexing on the hot path (numel < 2^31) to cut register pressure
#     and reach full occupancy (2048 threads/SM)
#   * a dedicated "uniform" kernel used when the element count is an exact
#     multiple of the block tile: zero bounds checks, 4 independent 16-byte
#     loads in flight per thread for deep memory-level parallelism
#   * minimal host-side work (single branch, no TensorIterator dispatch)
scalar_mul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <climits>

namespace {

constexpr int kThreads = 256;
constexpr int kUnroll  = 4;
constexpr int kTile    = kThreads * kUnroll;   // float4 elements per block

// Hot path: n4 is an exact multiple of kTile -> no bounds checks at all,
// 4 outstanding 16B loads per thread, 32-bit index math.
__global__ void __launch_bounds__(kThreads)
scalar_mul_vec4_uniform_kernel(const float4* __restrict__ A,
                               float4* __restrict__ C,
                               const float s) {
    const int base = (int)blockIdx.x * kTile + (int)threadIdx.x;

    float4 v[kUnroll];
    #pragma unroll
    for (int i = 0; i < kUnroll; ++i) {
        v[i] = A[base + i * kThreads];
    }
    #pragma unroll
    for (int i = 0; i < kUnroll; ++i) {
        v[i].x *= s;
        v[i].y *= s;
        v[i].z *= s;
        v[i].w *= s;
    }
    #pragma unroll
    for (int i = 0; i < kUnroll; ++i) {
        C[base + i * kThreads] = v[i];
    }
}

// Generic vectorized path: one float4 per thread with a single bounds check.
__global__ void __launch_bounds__(kThreads)
scalar_mul_vec4_kernel(const float4* __restrict__ A,
                       float4* __restrict__ C,
                       const float s,
                       const int n4) {
    const int idx = (int)blockIdx.x * kThreads + (int)threadIdx.x;
    if (idx < n4) {
        float4 v = A[idx];
        v.x *= s;
        v.y *= s;
        v.z *= s;
        v.w *= s;
        C[idx] = v;
    }
}

// Scalar grid-stride fallback (unaligned pointers / odd sizes / huge tensors).
__global__ void scalar_mul_fallback_kernel(const float* __restrict__ A,
                                           float* __restrict__ C,
                                           const float s,
                                           const long n) {
    const long stride = (long)gridDim.x * blockDim.x;
    for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < n; idx += stride) {
        C[idx] = A[idx] * s;
    }
}

}  // namespace

torch::Tensor scalar_mul_cuda(torch::Tensor A, double s) {
    if (!A.is_cuda() || A.scalar_type() != at::kFloat) {
        return at::mul(A, at::Scalar(s));
    }

    auto Ac = A.is_contiguous() ? A : A.contiguous();
    auto C  = at::empty(Ac.sizes(), Ac.options());

    const long n = (long)Ac.numel();
    if (n == 0) {
        return C;
    }

    const float* a_ptr = Ac.data_ptr<float>();
    float*       c_ptr = C.data_ptr<float>();
    const float  sf    = static_cast<float>(s);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const bool aligned =
        ((((uintptr_t)a_ptr) | ((uintptr_t)c_ptr)) & 0xF) == 0;
    const long n4 = n >> 2;

    if (aligned && (n & 3) == 0 && n4 <= (long)INT_MAX) {
        if ((n4 % kTile) == 0) {
            const int blocks = (int)(n4 / kTile);
            scalar_mul_vec4_uniform_kernel<<<blocks, kThreads, 0, stream>>>(
                reinterpret_cast<const float4*>(a_ptr),
                reinterpret_cast<float4*>(c_ptr),
                sf);
        } else {
            const int blocks = (int)((n4 + kThreads - 1) / kThreads);
            scalar_mul_vec4_kernel<<<blocks, kThreads, 0, stream>>>(
                reinterpret_cast<const float4*>(a_ptr),
                reinterpret_cast<float4*>(c_ptr),
                sf, (int)n4);
        }
    } else {
        long blocks = (n + kThreads - 1) / kThreads;
        const long cap = 1L << 20;
        if (blocks > cap) blocks = cap;
        scalar_mul_fallback_kernel<<<(unsigned int)blocks, kThreads, 0, stream>>>(
            a_ptr, c_ptr, sf, n);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

scalar_mul_cpp_source = "torch::Tensor scalar_mul_cuda(torch::Tensor A, double s);"

# Compile the inline CUDA extension
scalar_mul_ext = load_inline(
    name="scalar_mul_ext_v4",
    cpp_sources=scalar_mul_cpp_source,
    cuda_sources=scalar_mul_cuda_source,
    functions=["scalar_mul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs a matrix-scalar multiplication (C = A * s)
    using a custom, vectorized, bandwidth-maximizing CUDA kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._scalar_mul = scalar_mul_ext.scalar_mul_cuda

    def forward(self, A: torch.Tensor, s: float) -> torch.Tensor:
        return self._scalar_mul(A, s)
```