/*
 * 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

#include <algorithm>

namespace __nvhpc_std {

namespace __bl = __blas_std;

/* z = x + y */
template< class _ElementType_x
        , class _SizeType_x
        , ::std::size_t ... _ext_x
        , class _Layout_x
        , class _Accessor_x
        , class _ElementType_y
        , class _SizeType_y
        , ::std::size_t ... _ext_y
        , class _Layout_y
        , class _Accessor_y
        , class _ElementType_z
        , class _SizeType_z
        , ::std::size_t ... _ext_z
        , class _Layout_z
        , class _Accessor_z
        >
void __add_impl( __nvhpc_exec< __blas_exec_space > __exec
               , std::mdspan< _ElementType_x, std::extents<_SizeType_x, _ext_x ...>, _Layout_x, _Accessor_x > __x
               , std::mdspan< _ElementType_y, std::extents<_SizeType_y, _ext_y ...>, _Layout_y, _Accessor_y > __y
               , std::mdspan< _ElementType_z, std::extents<_SizeType_z, _ext_z ...>, _Layout_z, _Accessor_z > __z
               )
{
#ifdef STDBLAS_VERBOSE
    __STDBLAS_BACKEND_MESSAGE( add, BLAS );
#endif

    if constexpr ( __x.rank() == 1 )
    {
        int const __n = __z.extent(0);
        std::for_each( std::execution::par,
                       __stdblas_counting_iterator(0), __stdblas_counting_iterator(__n),
                       [=]( int __i ) {
                           __z[__i] = __x[__i] + __y[__i];
                       } );
    }
    else
    {
        int const __rows = __z.extent(0);
        int const __cols = __z.extent(1);

        int const __total = __rows * __cols;
        std::for_each( std::execution::par,
                       __stdblas_counting_iterator(0), __stdblas_counting_iterator(__total),
                       [=]( int __flat ) {
                           int __i = __flat % __rows;
                           int __j = __flat / __rows;
                           __z[__i, __j] = __x[__i, __j] + __y[__i, __j];
                       } );
    }
}

}
