以下为鄙人读源码时做的一些粗略记录,如有不当望请多多指教(抱拳)
一、源码解读(主要为多线程执行过程)。从C++ 模型声明tflite::Interpreter类型做入口
1、tensorflow\lite\interpreter.cc
void Interpreter::SetNumThreads(int num_threads) {
for (auto& subgraph : subgraphs_) {
subgraph->context()->recommended_num_threads = num_threads;
}
for (int i = 0; i < kTfLiteMaxExternalContexts; ++i) {
auto* c = external_contexts_[i];
if (c && c->Refresh) {
c->Refresh(context_);
}
}
}
(a)、void SetNumThreads(int num_threads) ; 线程数设置 最新github上改了函数名字,改成void cpus(int num_threads);
(b)、顺着函数里面subgraph找到, tensorflow\lite\core\subgraph.h 里面的class Subgraph{} 以及里面的成员TfLiteContext* context()。
(c)、TfLiteContext 在tensorflow\lite\c\c_api_internal.h 里面定义typedef struct TfLiteContext {} 以及成员变量recommended_num_threads 。
2、根据recommended_num_threads 找到tensorflow\lite\kernels\cpu_backend_context.cc
CpuBackendContext* CpuBackendContext::GetFromContext(TfLiteContext* context) {
auto* external_context = static_cast(
context->GetExternalContext(context, kTfLiteCpuBackendContext));
if (external_context == nullptr) {
TF_LITE_FATAL(
"ExternalCpuBackendContext isn't properly initialized during TFLite "
"interpreter initialization.");
}
auto* cpu_backend_context = static_cast(
external_context->internal_backend_context());
if (cpu_backend_context == nullptr) {
// We do the lazy initialization here for the TfLiteInternalBackendContext
// that's wrapped inside ExternalCpuBackendContext.
cpu_backend_context = new CpuBackendContext();
if (context->recommended_num_threads != -1) {
cpu_backend_context->SetMaxNumThreads(context->recommended_num_threads);
}
external_context->set_internal_backend_context(
std::unique_ptr(cpu_backend_context));
}
return cpu_backend_context;
}
(a)、函数里面重要的代码。CPU设置的最大线程数,cpu_backend_context->SetMaxNumThreads(context->recommended_num_threads) 。
(b)、cpu_backend_context.cc 文件里面还有另外一个函数,CpuBackendContext::CpuBackendContext() 里面设置SetMaxNumThreads(1); 说明,lite里面默认线程数为1.
(c)、此文件里面还有最后一个设置函数线程的函数,void CpuBackendContext::SetMaxNumThreads(int max_num_threads) {….},里面有个ruy_context_->max_num_threads = max_num_threads,顺着这个找下去,能找到ruy::Context
3、根据2的ruy:Context 定位到tensorflow\lite\experimental\ruy\context.h。此文件较多注释,可详细看注释。
/ The state private to each Ruy thread.
struct PerThreadState {
// Each thread may be running on a different microarchitecture. For example,
// some threads may be on big cores, while others are on little cores. Thus,
// it's best for the tuning to be per-thread.
TuningResolver tuning_resolver;
// Each thread has its own local allocator.
Allocator allocator;
};
……….
……….
……….
void EnsureNPerThreadStates(int thread_count) {
while (per_thread_states.size() < static_cast(thread_count)) {
per_thread_states.emplace_back(new PerThreadState);
}
}
Tuning GetMainThreadTuning() {
EnsureNPerThreadStates(1);
TuningResolver* tuning_resolver = &per_thread_states[0]->tuning_resolver;
tuning_resolver->SetTuning(explicit_tuning);
return tuning_resolver->Resolve();
}
(a)、里面定义了关于线程的两个重要函数void EnsureNPerThreadStates(int thread_count) {…}和 Tuning GetMainThreadTuning() { …} 。
(b)、EnsureNPerThreadStates () 函数里面有个PerThreadState 结构体,(通过注释可明显知道,此处为线程分配的地方)
(c)、Struct Context final{…}里面有个 int max_num_threads = 1; 再次确定CPU线程分配默认为1。
4、查找3里面的线程分配结构体里面的TuningResolver ,找到定义处: tensorflow\lite\experimental\ruy\tune.h
enum class Tuning {
// kAuto means please use auto-detection. It's the default in the
// user-visible parts (see Context). It's meant to be resolved to an
// actual tuning at some point by means of TuningResolver.
kAuto,
// Target an out-order CPU. Example: ARM Cortex-A75.
kOutOfOrder,
// Target an in-order CPU. Example: ARM Cortex-A55.
kInOrder
};
// Why a TuningResolver class?
//
// Ideally, this Library would offer a single function,
// Tuning GetCurrentCPUTuning();
//
// However, determining information about the current CPU is not necessarily,
// cheap, so we currently cache that and only invalidate/reevaluate after
// a fixed amount of time. This need to store state is why this library
// has to expose a class, TuningResolver, not just a function.
class TuningResolver {
public:
TuningResolver();
// Allows the user to specify an explicit Tuning value, bypassing auto
// detection; or to specify Tuning::kAuto, reverting to auto detection.
void SetTuning(Tuning tuning) { unresolved_tuning_ = tuning; }
(a)、此文件有较多注释,可详细看注释,从注释里面可知道SetTuning 是设置CPU类型的地方,默认为KAuto,自动识别硬件类型。
5、根据SetTuning 函数调用情况,定位到tensorflow\lite\experimental\ruy\trmul.cc 文件
a) SetTuning所在函数存在一个 params->RunKernel(tuning, origin, rounded_dims); params 被TrMulParams定义。
6、找到TrMulParams定义文件, tensorflow\lite\experimental\ruy\trmul_params.h
a) 此文件定义了一个TrMulParams结构体,结构体里面定义了一个RunKernel()函数
7、根据RunKernel()函数,找到tensorflow\lite\experimental\ruy\kernel_common.h。
void RunKernelTyped(Tuning tuning, const PackedMatrix& lhs,
const PackedMatrix& rhs, const Spec& spec,
int start_row, int start_col, int end_row, int end_col,
Matrix* dst) {
using Kernel = Kernel;
Kernel kernel(tuning);
using LhsLayout = typename Kernel::LhsLayout;
using RhsLayout = typename Kernel::RhsLayout;
// end_row and end_col may be larger than dst dimensions.
// that is because kernels write directly to the destination matrix, whose
// dimensions may not be a multiple of the kernel dimensions, and we try to
// keep this annoyance localized as an implementation detail in kernels,
// by allowing to pass rounded-up values down as far as possible.
// These assertions encode the contract.
RUY_DCHECK_LE(0, start_row);
RUY_DCHECK_LE(start_row, end_row);
RUY_DCHECK_LT(end_row, dst->layout.rows + LhsLayout::kCols);
RUY_DCHECK_EQ((end_row - start_row) % LhsLayout::kCols, 0);
RUY_DCHECK_LE(0, start_col);
RUY_DCHECK_LE(start_col, end_col);
RUY_DCHECK_LT(end_col, dst->layout.cols + RhsLayout::kCols);
RUY_DCHECK_EQ((end_col - start_col) % RhsLayout::kCols, 0);
#if RUY_OPT_ENABLED(RUY_OPT_FAT_KERNEL)
kernel.Run(lhs, rhs, spec, start_row, start_col, end_row, end_col, dst);
#else
for (int col = start_col; col < end_col; col += RhsLayout::kCols) {
int block_end_col = std::min(col + RhsLayout::kCols, end_col);
for (int row = start_row; row < end_row; row += LhsLayout::kCols) {
int block_end_row = std::min(row + LhsLayout::kCols, end_row);
kernel.Run(lhs, rhs, spec, row, col, block_end_row, block_end_col, dst);
}
}
#endif
}
(a)、此文件里面有个RunKernel()重写函数,与6里面的TrMulParams里面的RunKernel()函数关系在tensorflow\lite\experimental\ruy\dispatch.h文件里面有说明。
(b)、RunKernel()函数里面有个RunKernelTyped ()函数,此函数里面执行了最终的kernel.Run()函数,但此函数里面并为提及相关CPU信息,是一些矩阵操作。唯一涉及到的就是Turnning tuning。但默认tuning是KAuto即自动检测核类型。
二、多线程执行示意图
三、模型执行过程