# Copyright 2025 The OpenXLA Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== diff --git a/src/cpu/aarch64/acl_thread.cpp b/src/cpu/aarch64/acl_thread.cpp index 53175a05f9..89731cb356 100644 --- a/src/cpu/aarch64/acl_thread.cpp +++ b/src/cpu/aarch64/acl_thread.cpp @@ -1,5 +1,5 @@ /******************************************************************************* -* Copyright 2022-2024 Arm Ltd. and affiliates +* Copyright 2022-2025 Arm Ltd. and affiliates * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -83,17 +83,20 @@ void acl_set_threadpool_num_threads() { } // Swap BenchmarkScheduler for custom scheduler builds (i.e. ThreadPoolScheduler) void acl_set_tp_benchmark_scheduler() { - static std::once_flag flag_once; - // Create threadpool scheduler - std::unique_ptr threadpool_scheduler - = std::make_unique(); - arm_compute::IScheduler *_real_scheduler = nullptr; - _real_scheduler = threadpool_scheduler.release(); - // Create benchmark scheduler and set TP as real scheduler - std::shared_ptr benchmark_scheduler - = std::make_unique(*_real_scheduler); - std::call_once(flag_once, - [&]() { arm_compute::Scheduler::set(benchmark_scheduler); }); + static thread_local std::once_flag flag_once; + std::call_once(flag_once, [&]() { + // Create threadpool scheduler + std::unique_ptr threadpool_scheduler + = std::make_unique(); + arm_compute::IScheduler *_real_scheduler = nullptr; + _real_scheduler = threadpool_scheduler.release(); + + // Create benchmark scheduler and set TP as real scheduler + std::shared_ptr benchmark_scheduler + = std::make_unique(*_real_scheduler); + + arm_compute::Scheduler::set(benchmark_scheduler); + }); } #endif diff --git a/src/cpu/aarch64/acl_threadpool_scheduler.cpp b/src/cpu/aarch64/acl_threadpool_scheduler.cpp index 30910398d9..34cf44b7e2 100644 --- a/src/cpu/aarch64/acl_threadpool_scheduler.cpp +++ b/src/cpu/aarch64/acl_threadpool_scheduler.cpp @@ -1,5 +1,5 @@ /******************************************************************************* -* Copyright 2022-2024 Arm Ltd. and affiliates +* Copyright 2022-2025 Arm Ltd. and affiliates * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,24 +18,17 @@ #if DNNL_CPU_THREADING_RUNTIME == DNNL_RUNTIME_THREADPOOL -#include "cpu/aarch64/acl_thread.hpp" - #include "common/counting_barrier.hpp" #include "common/dnnl_thread.hpp" +#include "cpu/aarch64/acl_thread.hpp" #include "arm_compute/core/CPP/ICPPKernel.h" #include "arm_compute/core/Error.h" -#include "arm_compute/core/Helpers.h" -#include "arm_compute/core/Utils.h" #include "arm_compute/runtime/IScheduler.h" -// BARRIER #include #include -#include #include -#include -#include namespace dnnl { namespace impl { @@ -51,7 +44,7 @@ public: /// Function to check the next element in the range if there is one. bool get_next(unsigned int &next) { - next = atomic_fetch_add_explicit( + next = std::atomic_fetch_add_explicit( &_atomic_counter, 1u, std::memory_order_relaxed); return next < _end; } @@ -70,11 +63,8 @@ void process_workloads(std::vector &workloads, } while (feeder.get_next(workload_index)); } -ThreadpoolScheduler::ThreadpoolScheduler() { - using namespace dnnl::impl::threadpool_utils; - // Set number of threads to one when threadpool is not available. - _num_threads = get_active_threadpool() == nullptr ? 1 : num_threads_hint(); -} +ThreadpoolScheduler::ThreadpoolScheduler() + : _num_threads(dnnl_get_max_threads()) {} ThreadpoolScheduler::~ThreadpoolScheduler() = default; @@ -83,8 +73,8 @@ unsigned int ThreadpoolScheduler::num_threads() const { } void ThreadpoolScheduler::set_num_threads(unsigned int num_threads) { - arm_compute::lock_guard lock(this->_run_workloads_mutex); - _num_threads = num_threads == 0 ? num_threads_hint() : num_threads; + std::lock_guard lock(this->_mtx); + _num_threads = num_threads == 0 ? dnnl_get_max_threads() : num_threads; } void ThreadpoolScheduler::schedule(ICPPKernel *kernel, const Hints &hints) { @@ -104,7 +94,7 @@ void ThreadpoolScheduler::schedule_op(ICPPKernel *kernel, const Hints &hints, void ThreadpoolScheduler::run_workloads( std::vector &workloads) { - arm_compute::lock_guard lock(this->_run_workloads_mutex); + std::lock_guard lock(this->_mtx); const unsigned int num_threads = std::min(static_cast(_num_threads), diff --git a/src/cpu/aarch64/acl_threadpool_scheduler.hpp b/src/cpu/aarch64/acl_threadpool_scheduler.hpp index e9ba21c803..384dfec1b9 100644 --- a/src/cpu/aarch64/acl_threadpool_scheduler.hpp +++ b/src/cpu/aarch64/acl_threadpool_scheduler.hpp @@ -1,5 +1,5 @@ /******************************************************************************* -* Copyright 2022 Arm Ltd. and affiliates +* Copyright 2022, 2025 Arm Ltd. and affiliates * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -22,7 +22,8 @@ #if DNNL_CPU_THREADING_RUNTIME == DNNL_RUNTIME_THREADPOOL #include "arm_compute/runtime/IScheduler.h" -#include "support/Mutex.h" + +#include namespace dnnl { namespace impl { @@ -32,7 +33,7 @@ namespace aarch64 { class ThreadpoolScheduler final : public arm_compute::IScheduler { public: ThreadpoolScheduler(); - ~ThreadpoolScheduler(); + ~ThreadpoolScheduler() override; /// Sets the number of threads the scheduler will use to run the kernels. void set_num_threads(unsigned int num_threads) override; @@ -54,8 +55,8 @@ protected: void run_workloads(std::vector &workloads) override; private: - uint _num_threads {}; - arm_compute::Mutex _run_workloads_mutex {}; + unsigned int _num_threads {}; + std::mutex _mtx; }; } // namespace aarch64