/*++

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

Module Name:

    threading.cpp

Abstract:

    This module implements platform specific threading support.

--*/

#include "mlasi.h"

void
MlasExecuteThreaded(
    MLAS_THREADED_ROUTINE* ThreadedRoutine,
    void* Context,
    ptrdiff_t Iterations,
    MLAS_THREADPOOL* ThreadPool
    )
{
    //
    // Execute the routine directly if only one iteration is specified.
    //

    if (Iterations == 1) {
        ThreadedRoutine(Context, 0);
        return;
    }

#if defined(BUILD_MLAS_NO_ONNXRUNTIME)
    MLAS_UNREFERENCED_PARAMETER(ThreadPool);

    //
    // Fallback to OpenMP or a serialized implementation.
    //

    //
    // Execute the routine for the specified number of iterations.
    //
    if (ThreadPool == nullptr) {
        for (std::ptrdiff_t tid = 0; tid < Iterations; tid++) {
            ThreadedRoutine(Context, tid);
        }
    } else {
        ThreadPool->TrySimpleParallelFor(Iterations, [&](ptrdiff_t tid) {
            ThreadedRoutine(Context, tid);
        });
    }
#else
    //
    // Schedule the threaded iterations using the thread pool object.
    //

    MLAS_THREADPOOL::TrySimpleParallelFor(ThreadPool, Iterations, [&](ptrdiff_t tid) {
        ThreadedRoutine(Context, tid);
    });
#endif
}


void
MlasTrySimpleParallel(
    MLAS_THREADPOOL * ThreadPool,
    const std::ptrdiff_t Iterations,
    const std::function<void(std::ptrdiff_t tid)>& Work)
{
    //
    // Execute the routine directly if only one iteration is specified.
    //
    if (Iterations == 1) {
        Work(0);
        return;
    }

#if defined(BUILD_MLAS_NO_ONNXRUNTIME)
    MLAS_UNREFERENCED_PARAMETER(ThreadPool);

    //
    // Fallback to OpenMP or a serialized implementation.
    //

    //
    // Execute the routine for the specified number of iterations.
    //
    if (ThreadPool == nullptr) {
        for (std::ptrdiff_t i = 0; i < Iterations; i++) {
            Work(i);
        }
    } else {
        ThreadPool->TrySimpleParallelFor(Iterations, Work);
    }
#else
    //
    // Schedule the threaded iterations using the thread pool object.
    //

    MLAS_THREADPOOL::TrySimpleParallelFor(ThreadPool, Iterations, Work);
#endif
}