/* ************************************************************************
 * Copyright (C) 2025 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop-
 * ies of the Software, and to permit persons to whom the Software is furnished
 * to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM-
 * PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE-
 * CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *
 * ************************************************************************ */

#include <thread>

#ifdef _OPENMP
#include <omp.h>
#else
#ifndef __HIP_DEVICE_COMPILE__
#pragma GCC warning "_OPENMP not defined so client build can not utilize OPENMP."
#endif
#endif

#include "client_omp.hpp"
#include "rocblas_ostream.hpp"

// Conservative thread reducer for multi-threaded client scenarios
// Used in client_omp_manager constructor to avoid over-threading
static constexpr int c_thread_reducer = 4;

client_omp_manager::client_omp_manager(size_t std_thread_count)
    : m_original_omp_threads(1)
    , m_active(false)
{
#ifdef _OPENMP
    if(std_thread_count > 1)
    {
        // OPENMP behaviour not defined in std::threads so reduce potential for over threading
        const int processor_count     = std::thread::hardware_concurrency();
        m_original_omp_threads        = omp_get_max_threads();
        const int omp_current_threads = m_original_omp_threads;

        if(omp_current_threads * std_thread_count > processor_count - c_thread_reducer)
        {
            int omp_limit_threads = omp_current_threads / std_thread_count;
            omp_limit_threads     = std::max(1, omp_limit_threads);

            if(omp_limit_threads != m_original_omp_threads)
            {
                m_active = true;
                omp_set_num_threads(omp_limit_threads);
                static int once
                    = (rocblas_cout << "rocBLAS info: client (OPENMP) multi-thread reducing "
                                       "omp_set_num_threads from "
                                    << m_original_omp_threads << " to " << omp_limit_threads
                                    << " per thread." << std::endl,
                       1);
            }
        }
    }
#endif
}

client_omp_manager::~client_omp_manager()
{
#ifdef _OPENMP
    if(m_active)
    {
        omp_set_num_threads(m_original_omp_threads);
    }
#endif
}

void client_omp_manager::limit_by_processor_count()
{
    // Limit OMP usage to avoid performance degradation in reference library at high thread counts
    // See: rocBLAS Programmer's Guide - AOCL threading recommendations
#ifdef _OPENMP
    // Detect available cores using most conservative estimate
    // omp_get_max_threads() respects SLURM/Docker/OMP_NUM_THREADS limits
    // hardware_concurrency() reports physical cores
    // Use minimum to handle all environments correctly
    const int omp_threads       = omp_get_max_threads();
    const int hardware_threads  = std::thread::hardware_concurrency();
    const int available_threads = std::min(omp_threads, hardware_threads);

    if(available_threads <= 0)
    {
        return; // Sanity check - should not happen
    }

    // Determine optimal thread reduction based on system size
    // AOCL oversubscription is most severe on large core-count systems (observed 64x slowdown
    // on 24-core system). On smaller systems, use conservative reduction to preserve performance.
    int thread_reducer;
    if(available_threads <= 16)
    {
        thread_reducer = 2; // Small systems: minimal reduction
    }
    else
    {
        thread_reducer = 4; // Large systems: aggressive reduction to prevent AOCL degradation
    }

    int safe_thread_count = std::max(1, available_threads - thread_reducer);
    omp_set_num_threads(safe_thread_count);

    rocblas_cout << "rocBLAS info: Limiting OpenMP threads to " << safe_thread_count
                 << " (detected " << available_threads << " available, reduced by "
                 << thread_reducer << " to optimize AOCL performance)" << std::endl;
#endif
}