wehub-resource-sync
74 lines
2.3 KiB
C++
74 lines
2.3 KiB
C++
/*!
|
|
* \file support/threading_backend.h
|
|
* \brief Compatibility shim providing the threading helpers that used to live
|
|
* at <tvm/runtime/threading_backend.h>. The public TVM header was
|
|
* removed in a runtime refactor, but the underlying symbols
|
|
* (TVMBackendParallelLaunch, threading::MaxConcurrency,
|
|
* threading::SetMaxConcurrency) are still exported from libtvm_runtime.
|
|
* We re-declare the inline parallel-for template here so existing
|
|
* mlc-llm call sites keep compiling.
|
|
*/
|
|
#ifndef MLC_LLM_SUPPORT_THREADING_BACKEND_H_
|
|
#define MLC_LLM_SUPPORT_THREADING_BACKEND_H_
|
|
|
|
#include <tvm/runtime/base.h>
|
|
#include <tvm/runtime/c_backend_api.h>
|
|
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
|
|
namespace tvm {
|
|
namespace runtime {
|
|
namespace threading {
|
|
|
|
/*! \return the maximum number of effective workers for this system. */
|
|
int MaxConcurrency();
|
|
/*! \brief Setting the maximum number of available cores. */
|
|
void SetMaxConcurrency(int value);
|
|
|
|
} // namespace threading
|
|
|
|
namespace detail {
|
|
|
|
template <typename T>
|
|
struct ParallelForWithThreadingBackendLambdaInvoker {
|
|
static int TVMParallelLambdaInvoke(int task_id, TVMParallelGroupEnv* penv, void* cdata) {
|
|
int num_task = penv->num_task;
|
|
T* lambda_ptr = static_cast<T*>(cdata);
|
|
(*lambda_ptr)(task_id, num_task);
|
|
return 0;
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
inline void parallel_launch_with_threading_backend(T flambda) {
|
|
void* cdata = &flambda;
|
|
TVMBackendParallelLaunch(ParallelForWithThreadingBackendLambdaInvoker<T>::TVMParallelLambdaInvoke,
|
|
cdata, /*num_task=*/0);
|
|
}
|
|
|
|
} // namespace detail
|
|
|
|
template <typename T>
|
|
inline void parallel_for_with_threading_backend(T flambda, int64_t begin, int64_t end) {
|
|
if (end - begin == 1) {
|
|
flambda(begin);
|
|
return;
|
|
}
|
|
auto flaunch = [begin, end, flambda](int task_id, int num_task) {
|
|
int64_t total_len = end - begin;
|
|
int64_t step = (total_len + num_task - 1) / num_task;
|
|
int64_t local_begin = std::min(begin + step * task_id, end);
|
|
int64_t local_end = std::min(local_begin + step, end);
|
|
for (int64_t i = local_begin; i < local_end; ++i) {
|
|
flambda(i);
|
|
}
|
|
};
|
|
detail::parallel_launch_with_threading_backend(flaunch);
|
|
}
|
|
|
|
} // namespace runtime
|
|
} // namespace tvm
|
|
|
|
#endif // MLC_LLM_SUPPORT_THREADING_BACKEND_H_
|