// Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//

#if !defined(__arm__) && !defined(__aarch64__)
#define SIMD_AVAILABLE
#endif

#include <OgnTutorialSIMDAddDatabase.h>

#ifdef SIMD_AVAILABLE
#include <immintrin.h>
#endif

// This node perform a sum using SIMD instruction set
class OgnTutorialSIMDAdd
{
public:
    static size_t computeVectorized(OgnTutorialSIMDAddDatabase& db, size_t count)
    {
        // Constants
        static size_t constexpr kSIMDSize = 4 * sizeof(float);
        static size_t constexpr kMask = kSIMDSize - 1;

        // Retrieve data
        auto opA = db.inputs.a.vectorized(count);
        auto opB = db.inputs.b.vectorized(count);
        auto res = db.outputs.result.vectorized(count);

        //Regular loop definition
        auto regularLoop = [&](size_t const begin, size_t const end) -> size_t const
        {
            for (size_t idx = begin; idx < end; ++idx)
                res[idx] = opA[idx] + opB[idx];
            return end;
        };

#ifdef SIMD_AVAILABLE

        //Alignment must be identical
        bool const correctlyAligned = ((size_t(opA.data()) & kMask) == (size_t(opB.data()) & kMask))
            && ((size_t(opA.data()) & kMask) == (size_t(res.data()) & kMask));

        if (!correctlyAligned)
        {
            regularLoop(0, count);
        }
        else
        {
            // Unaligned elements
            size_t const maskedAdress = (size_t(opA.data()) & kMask);
            size_t const unalignedCount = maskedAdress
                                            ? regularLoop(0, (kSIMDSize - maskedAdress) / sizeof(float))
                                            : 0;

            // Vectorized elements
            size_t const vectorizedCount = (count - unalignedCount) & (~kMask);
            size_t const vectorizedLoop = vectorizedCount / (kSIMDSize / sizeof(float));

            __m128* aSIMD   = (__m128*)(opA.data() + unalignedCount);
            __m128* bSIMD   = (__m128*)(opB.data() + unalignedCount);
            __m128* resSIMD = (__m128*)(res.data() + unalignedCount);
            for (size_t idx = 0; idx < vectorizedLoop; ++idx)
                resSIMD[idx] = _mm_add_ps(aSIMD[idx], bSIMD[idx]);

            // Remaining elements
            regularLoop(unalignedCount + vectorizedCount, count);
        }

#else

        regularLoop(0, count);

#endif

        return count;
    }
};

REGISTER_OGN_NODE()
