/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#pragma once

namespace __nvhpc_std {

namespace __bl = __blas_std;

template<class _Scalar,
         class _SizeType,
         class _ElementType,
         ::std::size_t ... _ext,
         class _Layout,
         class _Accessor>
void __scale_impl( __nvhpc_exec<__blas_exec_space> /* __exec */
                 , const _Scalar __alpha
                 , std::mdspan< _ElementType, std::extents<_SizeType, _ext ...>, _Layout, _Accessor > __x
                 )
{
#ifdef STDBLAS_VERBOSE
    __STDBLAS_BACKEND_MESSAGE( scale, BLAS );
#endif
    static_assert (__x.rank() <= 2);

    using __val_t = typename decltype(__x)::value_type;

    if ( __x.rank() == 1 )
    {
        __bl::__blas_scal( __x.extent(0), __alpha, __x.data_handle(), __x.stride(0) );
    }
    else if ( std::linalg::__is_contiguous( __x ) )
    {
        __bl::__blas_scal( __x.extent(0) * __x.extent(1), __alpha, __x.data_handle(), 1 );
    }
    else
    {
        __bl::__blas_scal_2d( std::linalg::__get_mem_row_count( __x )
                            , std::linalg::__get_mem_col_count( __x )
                            , __alpha
                            , __x.data_handle()
                            , std::linalg::__get_leading_dim( __x )
                            );
    }
}

}
