1 /* 2 * Copyright (c) Meta Platforms, Inc. and affiliates. 3 * All rights reserved. 4 * 5 * This source code is licensed under the BSD-style license found in the 6 * LICENSE file in the root directory of this source tree. 7 */ 8 9 #pragma once 10 11 #include <cstdint> 12 #include <functional> 13 14 namespace executorch { 15 namespace extension { 16 17 /** 18 * A helper to run function in parallel. 19 * 20 * begin, end: describe the extent of the workitems via first and last workitem 21 * to be processed 22 * grain_size: number of workitems processed by user callback which is 23 * described below 24 * f: user function applied in parallel to the chunks, signature: 25 * void f(int64_t begin, int64_t end) 26 * Returns true if all work items are processed successfully, false otherwise 27 * 28 * Warning: parallel_for does NOT copy thread local states from the current 29 * thread to the worker threads. Users need to protect the access to captured 30 * data if they mutate them in f. 31 */ 32 bool parallel_for( 33 const int64_t begin, 34 const int64_t end, 35 const int64_t grain_size, 36 const std::function<void(int64_t, int64_t)>& f); 37 38 int64_t get_thread_num(); 39 40 void set_thread_num(int64_t thread_num); 41 42 } // namespace extension 43 } // namespace executorch 44 45 namespace torch { 46 namespace executor { 47 // TODO(T197294990): Remove these deprecated aliases once all users have moved 48 // to the new `::executorch` namespaces. 49 using ::executorch::extension::get_thread_num; 50 using ::executorch::extension::parallel_for; 51 using ::executorch::extension::set_thread_num; 52 } // namespace executor 53 } // namespace torch 54