静态图执行器预分析性能优化

[AndSonder]

任务名称	静态图执行器预分析性能优化
提交作者	Sonder @AndSonder
提交时间	2023-07-11
版本号	v1.0
依赖飞桨版本	develop

一、任务目标

飞桨静态图执行器通过对内核选择、数据转换、图依赖分析、线程调度分析、变量生命周期分析和多流分析等工作进行静态预处理，实现了“一轮分析多轮复用”，以尽可能减少动态调度开销，提升执行性能。

虽然静态预处理只需要在模型执行过程中进行一次，但对于规模较大的模型而言，这个过程会持续较长时间，消耗较多计算资源。本任务的目标是分析静态图执行器预处理阶段的性能瓶颈，并有针对性地设计和开发优化方案，以降低执行器预处理的耗时，提升大型模型的迭代效率。

二、性能分析

飞桨的静态预处理主要负责资源精细化管理和计算图预分析。计算图预分析包括了内核选择、数据转换、依赖分析、线程调度、变量生命周期和多流分析等工作。结合 Paddle 源码分析，我们可以得出两个可以进行优化的点：

1. 预分析结果复用

2. 全静态选Kernel推全

接下来分别详细说明一下这两个优化点。

2.1 预分析结果复用

Paddle 在将Program输入到执行器之前会进行预排列，将其分为多个任务（Job），然后按顺序执行。每个Job对应的interpretercore在执行过程中会缓存内核选择、数据转换和图依赖分析等数据。有些任务是完全相同的，也就是说在顺序执行的过程中，相同的任务可以共享一些数据。下图展示了共享预分析数据的示意图：

其中F0、F1和F2表示相同的任务。当F0对应的interpretercore运行了之后，F1和F2便无需重复分析共享的数据。

2.2 全静态选Kernel推全

当下，新的执行器在预处理阶段采用了“一轮动态调度+缓存”的方法来实现 kernel 的“半静态化”选择。第一轮采用 for-loop 方式动态选择 kernel 运行，并缓存 kernel、context 和数据传输的信息。利用这些信息完成其他分析和调度上的优化需要到第二轮才会真正生效。选用半静态化的方式使得新的执行器在第一个批次和其他批次的行为不一致，而这种不一致进一步影响显存性能、通用性、易用性和可维护性。

在新的 Phi 函数式算子体系下，每个 kernel 的输入、输出和属性都是明确的，并在 kernel 注册时记录下来。这些关键信息被暴露给框架调度，使得全静态选 kernel 变得可能。

当前全静态选 kernel 已经实现了基础功能，并成功地在 ResNet 和 Transformer 等模型上运行。在 PR #52395 和 PR #51292 中针对无法准确静态获取输入输出信息的不规范算子进行了修改。然而，在默认开启全静态选 kernel 的情况下，部分模型的行为与不开启全静态 kernel 的情况下不一致。我们将在本任务中修复这个问题，并推动使用全静态选 kernel 替代当前的“半静态化”方案。

三、优化方案

3.1 预分析结果复用方案

预分析过程中会进行图依赖分析、事件分析、内核选择、数据转换、线程调度分析、变量生命周期分析和GC分析等工作。下面详细地分析哪些数据可以复用，如果可以复用则给出复用方案。

3.1.1 图依赖分析复用

复用信息分析：

在Paddle中，图依赖分析（op dependency）发生在Convert过程中，其依赖分析过程在BuildOperatorDependences 构建。图依赖分析的主要作用是分析 op 之间的依赖关系，对计算图进行解析和分析，构建出每个算子的调度关系和执行顺序，以保证正确的计算顺序和结果。对于一样的子图，图依赖分析的结果是一样的，因此可以进行复用。

BuildOperatorDependences 的目标是建立 Instruction 之间的 DAG图，便于后续并行调度。Instruction 中包含了上一步中建立的 OpFuncNode 作为成员变量，它是调度的一个基础单位。 Instruction的图结构是一种邻接链表的方式存储的，每个节点存储他能到达邻居节点。

BuildOperatorDependences中使用图依赖构建的部分由dependency_builder_.Build(vec_instruction_)进行构建。dependency_builder_是ProgramInterpreter类的成员。在dependency_builder_.Build会将每个op的下游op存储在op_downstream_map_中，每个op的上游op存储在op_happens_before_当中。op依赖计数存储在dependecy_count_当中。

对于一样的子图，可以通过智能指针完成op 的依赖信息的共享。

代码实现思路：

Op 依赖的复用可以分为三个部分：

1. 实现依赖信息共享的相关逻辑，即从其他 Interpreter 中共享 op_downstream_map_ 和 op_happens_before_ 和 dependecy_count_
2. 实现判断是否能进行共享的相关逻辑
3. 添加共享后跳过图依赖分析部分相关代码

对于第一点，op_downstream_map_和op_happens_before_信息的共享需要给DependencyBuilder类添加ShareDependencyFrom接口并将这两个变量使用智能指针进行管理。其核心逻辑如下：

std::tuple<std::shared_ptr<std::map<size_t, std::set<size_t>>>,
           std::shared_ptr<std::vector<std::vector<bool>>>>
DependencyBuilder::GetDependency() {
  if (op_downstream_map_ == nullptr || op_happens_before_ == nullptr) {
    op_downstream_map_ = std::make_shared<std::map<size_t, std::set<size_t>>>();
    op_happens_before_ = std::make_shared<std::vector<std::vector<bool>>>();
  }
  return std::make_tuple(op_downstream_map_, op_happens_before_);
}

void DependencyBuilder::ShareDependencyFrom(DependencyBuilder* src) {
  std::tie(op_downstream_map_, op_happens_before_) = src->GetDependency();
  is_build_ = true;
}

由于依赖的信息较多，在ProgramInterpreter中添加多个Share函数会让代码冗余。因此可以统一在 InterpreterBaseImpl 以及其相关子类添加功能接口 void ShareBuildResultsFrom(InterpreterBaseImpl* src) ，实现共享的相关逻辑。ShareBuildResultsFrom 将共享对象传入并使用智能指针共享数据。ShareBuildResultsFrom 中不止会有共享图依赖的相关逻辑，也会有共享事件分析器、线程池调度、GC等的逻辑。其中共享图依赖的相关逻辑如下：

// 共享dependency_builder_的build结果
dependency_builder_.ShareDependencyFrom(src->GetDependencyBuilder());
// 共享dependecy_count_
dependecy_count_ = src->GetDependecyCount();

对于第二点，可以通过 job 对象的 const std::string type_ 来判断。具有相同的 type_ 说明两个任务的类型相同，即子图等信息相同。也就说明可以共享op的依赖信息。为了快速的判断哪些 job 可以进行共享，可以维护一个类变量 std::map<string, int> type_to_first_id 来判断哪些 job 的类型是一样的。其相关逻辑如下：

std::map<std::string, size_t> type_to_first_id;
  if (!is_interpretercore_build_result_shared_) {
    type_to_first_id[jobs[0]->Type()] = 0;
    for (size_t job_idx = 1; job_idx < jobs.size(); ++job_idx) {
      interpretercores_[job_idx]->ShareWorkQueueFrom(interpretercores_[0]);
      // TODO(Ruibiao): Share other build result, e.g., kernel choosing, data
      // transfer, op dependency, thread scheduling, GC, event analyzer, and so
      // on.
      if (type_to_first_id.count(jobs[job_idx]->Type()) == 0) {
        // 保存下每一种 job_type 的第一个 job_id
        type_to_first_id[jobs[job_idx]->Type()] = job_idx;
      }
    }
    is_interpretercore_build_result_shared_ = true;
  }

  for (size_t job_idx = 0; job_idx < jobs.size(); ++job_idx) {
    const auto& job = jobs[job_idx];
    const std::string& job_type = job->Type();
    platform::RecordEvent record_event(
        job_type + "-" + std::to_string(job->MicroBatchId()),
        platform::TracerEventType::UserDefined,
        1);

    VLOG(6) << "Run job (" << job_idx << "), type = " << job_type
            << ", micro_batch_id =" << job->MicroBatchId();
    // Note(sonder): Share build results don't work for new IR now.
    if (type_to_first_id.count(job_type) != 0 &&
        !FLAGS_enable_new_ir_in_executor) {
      // 如果类型相同则进行共享
      interpretercores_[job_idx]->ShareBuildResultsFrom(
          interpretercores_[type_to_first_id[job_type]]);
    }
    interpretercores_[job_idx]->Run(feed_names, /*need_fetch = */ false);
  }

对于第三点可以给ProgramInterpreter添加一个统一开关 is_shared_results_build_ 进行了共享的interpretercore将该开关设为true。添加了开关后根据开关的值跳过相关的创建逻辑。

if (is_shared_results_build_ || !src.IsSharedResultsBuild()) {
    return;
}

然后在构建图依赖时添加相应的跳过逻辑：

  dependecy_count_ = GetDependencyCount();
  if (!is_shared_results_build_) {
    dependecy_count_->assign(instr_num, 0);
  }
  ...
    if (!is_shared_results_build_) {
      for (size_t next_instr_id : next_instr_ids) {
        ++(*dependecy_count_)[next_instr_id];
      }
    }

相关PR： 【静态图性能优化】图依赖信息复用 #55389

3.1.2 事件分析器复用

复用信息分析：

在Paddle中，事件分析器（event analyzer）也发生在Convert过程中。事件分析器的主要作用是为不同op的执行添加Event进行同步，便于后续并行。对于一样的子图，事件分析器分析的结果是一样的，因此可以进行复用。

事件分析器依赖于图依赖分析的结果，其核心逻辑在 stream_analyzer_.ConstructEvents 中实现。其中耗时最多的部分为根据op依赖分析出哪些op之间需要进行同步。

代码实现思路：

在 stream_analyzer_.ConstructEvents 中 event 分析的结果被存储在 event_info 当中。对于一样的子图，event_info 是一样的。因此可以将该变量作为类变量并使用智能指针进行共享。

将事件分析结果进行共享分如下步骤：

1. 实现共享事件分析结果 event_info, 即从其他 Interpretercore 中共享已经分析好的结果。
2. 实现判断是否能进行共享的相关逻辑
3. 添加共享后跳过图依赖分析部分相关代码

对于第一步首先需要将 event_info 改写为一个类变量 event_info_，并使用智能指针 std::shared_ptr 管理该变量。接下来需要给 StreamAnalyzer 添加 ShareEventInfoFrom 接口用于和其他 Interpretercore 中的 stream_analyer 共享分析结果。其核心逻辑如下：

std::shared_ptr<
    std::map<const DeviceContext*, std::map<size_t, std::set<size_t>>>>
StreamAnalyzer::GetEventInfo() const {
  return event_info_;
}

void StreamAnalyzer::ShareEventInfoFrom(const StreamAnalyzer& src) {
  event_info_ = src.GetEventInfo();
  is_event_info_build_ = true; // 用于共享后跳过分析
}

第二步已经在图依赖分析复用中实现主体功能，在共享事件分析结果时只需要在 ShareBuildResultsFrom 中添加相关共享逻辑即可。其主要逻辑代码如下：

// share event analysis
stream_analyzer_.ShareEventInfoFrom(src.GetStreamAnalyzer());

对于第三步则需要在 stream_analyzer_.ConstructEvents 中添加相应判断跳过事件分析的过程，其代码如下：

if (!is_event_info_build_) {
    std::vector<Instruction> cross_step_merged_instructions = *instructions;
    for (const Instruction& instr : *instructions) {
      cross_step_merged_instructions.emplace_back(instr);
    }
    ...
    AnalyseAllRunType(
        cross_step_merged_instructions_ptr, downstream_map, &run_type_info);

    AnalyseAllEventInfo(
        cross_step_merged_instructions_ptr, run_type_info, event_info_.get());
    ShrinkEventInfo(dependency_builder, event_info_.get());
    is_event_info_build_ = true;
  }

相关PR： 【静态图性能优化】Share event #55650

3.1.3 其他信息的复用

除了图依赖和事件分析器以外，GC与线程调度的分析结果也可以进行共享。但是考虑到当下添加共享GC和线程调度的代码可能会给New IR的推进带来麻烦，故将相关PR关闭。相关PR如下，均可以通过CI测试：

• 【静态图性能优化】Share gc #55626
• 【静态图性能优化】Share thread scheduling #55625

3.2 全静态选Kernel推全方案

当下，Paddle框架启用静态图模式时，由于fluid算子机制的限制，新的静态图执行器采用了一种称为“一轮动态调度+缓存”的预处理方式，以实现对kernel选择的“半静态化”。具体来说，就是在预处理阶段，通过实际执行一轮算子来动态选择kernel，并将选择结果缓存以供后续的训练轮次使用。然而这样半静态选Kernel使得新执行器第一个batch和后续其它batch行为不一致，这种不一致又进一步影响了：显存性能、通用性、易用性以及可维护性C++端每个算子针对不同的后端设备（backend）、数据类型（dtype）和数据布局（layout），特化了多个不同的计算kernel，根据KernelKey三元组做标识和索引， kernel选择即根据输入数据的信息，为每个算子选择正确的执行kernel也就是说这些信息实际上是可以在预处理阶段进行“静态化”的。 在确定硬件设备（place）和数据类型（dtype）之后，OP 需要执行的计算kernel理应就是可以确定的，可以在预处理阶段提前分析出来，也就是所谓的全静态。为了实现全静态选 Kernel 我们主要有三个任务，第一个是将影响全静态选Kernel的fluid算子迁移至PHI 下并补全注册信息，第二个任务是为已有的PHI算子添加注册信息, 第三个是为控制流算子添加静态选Kernel的相关逻辑。

3.2.1 前期工作

截止2023年8月，第一个任务和第二个任务已经完成大部分工作，相关 Tracking Issue 如下：

• Good first issue for phi kernel output definition #51292
• [PHI] Functionalization for Fluid kernel #52395
• [Fluid] Migrate Fluid Distributed Kernels to PHI #55716

3.2.2 相关FLAG移除

在静态选Kernel任务推动的前期，由于很多的op没有注册信息，因此我们需要添加一个黑名单，将这些op的静态选Kernel的开关关闭。在前期工作完成后，我们可以这些FLAG移除，使得所有的op都可以进行静态选Kernel。相关PR：

• Remove need_move_to_phi #56371

在移除FLAG的过程中，我们发现了两个需要迁移的算子，分别是shuffle_batch 和 distributed_fused_lamb_init。这两个算子的迁移工作已经在PR #56371 和PR #55993 中完成。

3.2.3 控制流算子的静态选Kernel

在前期工作完成后，我们可以将所有的op都进行静态选Kernel。但是由于控制流算子的特殊性，我们需要为控制流算子添加静态选Kernel的相关逻辑。在Paddle中，控制流算子主要有三种，分别是 While 和ConditionalBlock。我们首先需要为 ConditionalBlock 添加静态选Kernel的相关逻辑，然后再为While添加相关逻辑。

为ConditionalBlock添加静态选Kernel的相关逻辑的主要思路是：

• 解决子图静态选Kernel的问题
• 给ConditionalBlock算子的输出添加静态选Kernel的推导逻辑

1、解决子图静态选Kernel的问题

当下子图默认是关闭全静态选Kernel的，也就是说子图中的算子都是动态选Kernel的。其中控制子图的开关代码如下：

static_build_ = FLAGS_new_executor_static_build &&
                  !FLAGS_new_executor_use_cuda_graph &&
                  !execution_config.used_for_control_flow_op &&
                  interpreter::BlockCanBeStaticBuilt(block);

其中 execution_config.used_for_control_flow_op 就是控制子图的开关。我们可以通过移除该开关来解决子图静态选Kernel的问题。相关PR：

• Support control flow for static build [Step 1: support subgraph] #56185

移除开关后，通过CI发现子图在开启全静态选Kernel后 searchsorted 算子在跑 InferMeta 推导 dtype 信息时出错。在 static_build.cc 中添加了相关推导逻辑后，该问题得到解决。 相关代码如下：

else if (op_type == "searchsorted") {
    bool out_int32 = op.Attr<bool>("out_int32");
    if (out_int32) {
        dtype = DataType::INT32;
    } else {
        dtype = DataType::INT64;
    }
}

解决子图静态选Kernel的问题后，我们可以为ConditionalBlock添加静态选Kernel的推导逻辑。

2、为ConditionalBlock添加静态选Kernel的推导逻辑

由于控制流算子的分支选择条件是由运行时的变量动态决定的，因此我们无法在预处理阶段就确定分支的选择。所以静态选Kernel无法完全支持控制流算子。我们可以支持的情况如下：

1. 子图并不会改变输出的数据类型和Place
2. 后续的算子没有用到 ConditionalBlock 的输出作为输入

对于第一点我们需要在 FakeInitializeOutputsForOperatorBase 添加相关的推导逻辑。具体来说我们需要单独跑一遍子图的静态选Kernel，并记录下子图跑静态选Kernel前后的输出数据类型和Place。通过比对两次的输出数据类型和Place，如果没有发生变化则说明子图并不会改变输出的数据类型和Place。如果改变了则直接报错提示关闭静态选Kernel。相关代码如下：

if (op_type == "conditional_block") {
    // 记录子图静态选Kernel之前的输出数据类型和Place
    const std::vector<VarMetaInfo> out_var_info_before_build =
        GetVarsInfo(scope, op.Outputs(), op);
    // 跑子图静态选Kernel
    RunPreStaticBuild(*scope, place, op);
    // 记录子图静态选Kernel之后的输出数据类型和Place
    const std::vector<VarMetaInfo> out_var_info_after_build =
        GetVarsInfo(scope, op.Outputs(), op);

    // 比对两次的输出数据类型和Place
    for (size_t i = 0; i < out_var_info_before_build.size(); ++i) {
      // static build is supported in case of the output's dtype/place
      // is changed but the following op is not use this output
      if (out_var_info_before_build[i] != out_var_info_after_build[i]) {
        auto var_name = out_var_info_before_build[i].name_;
        if (following_input_vars.count(var_name)) {
          PADDLE_THROW(phi::errors::PreconditionNotMet(
              ...
        }
      }
    }
  }

其中为了方便的存储输出数据类型和Place，我们创建了一个 VarMetaInfo 的结构体，其定义如下：

struct VarMetaInfo {
  std::string name_;
  phi::DataType dtype_;
  phi::Place place_;

  explicit VarMetaInfo(const std::string& name) : name_(name) {
    dtype_ = phi::DataType::UNDEFINED;
    place_ = phi::Place();
  }

  VarMetaInfo(const std::string& name,
              const phi::DataType& dtype,
              const platform::Place& place)
      : name_(name), dtype_(dtype), place_(place) {}

  bool operator==(const VarMetaInfo& other) const {
    return name_ == other.name_ && dtype_ == other.dtype_ &&
           place_ == other.place_;
  }

  bool operator!=(const VarMetaInfo& other) const {
    return name_ != other.name_ || dtype_ != other.dtype_ ||
           place_ != other.place_;
  }
};

RunPreStaticBuild 中包含了子图跑静态选Kernel的相关逻辑，其主要核心思路是构建出子图的 InterpreterCore 并跑一遍静态选Kernel，详细代码见相关PR。

对于第二点我们首先需要在 HandleOperatorBase 中记录下 conditional_block_op 之后带有Kernel的算子的输入。其核心代码如下：

bool is_skip_fake_init = false;
std::unordered_set<std::string> following_input_vars;

if (op->Type() == "conditional_block") {
    // Note(sonder): skip fake init for conditional_block when
    // there is no op with kernel after it.
    is_skip_fake_init = true;
    for (size_t i = 0; i < following_ops.size(); ++i) {
    if (dynamic_cast<framework::OperatorWithKernel*>(
            following_ops[i].get()) != nullptr) {
        VLOG(4) << "Find op with kernel after conditional_block : "
                << following_ops[i]->Type();
        is_skip_fake_init = false;
        // 记录下带有Kernel的算子的输入
        auto input_vars_info = GetVarsInfo(
            scope, following_ops[i]->Inputs(), *following_ops[i].get());
        for (auto& input_var_info : input_vars_info) {
            following_input_vars.insert(input_var_info.name_);
        }
        }
    }
}

然后在 FakeInitializeOutputsForOperatorBase 中添加相关的判断逻辑。如果 conditional_block_op 的后续没有带有Kernel的算子了，那么就直接跳过 FakeInitializeOutputsForOperatorBase。否则就对子图的输出进行推导，如果推导出的数据类型和Place和之前的不一致首先判断不一致的这个变量是否是后续带有Kernel的算子的输入，如果是则跳过报错。如果不是则报错提示关闭静态选Kernel。相关代码如下：

// 比对两次的输出数据类型和Place
for (size_t i = 0; i < out_var_info_before_build.size(); ++i) {
    // static build is supported in case of the output's dtype/place
    // is changed but the following op is not use this output
    if (out_var_info_before_build[i] != out_var_info_after_build[i]) {
    auto var_name = out_var_info_before_build[i].name_;
    if (following_input_vars.count(var_name)) {
        PADDLE_THROW(phi::errors::PreconditionNotMet(
            ...
    }
  }
}

相关PR：

• Support control flow for static build [Step 2: support conditional_block] #56696

3、为While添加静态选Kernel的推导逻辑

While算子支持静态选Kernel的情况和ConditionalBlock算子类似，因此我们可以直接复用部分ConditionalBlock算子的相关逻辑。可复用的逻辑包含：

• 子图静态选Kernel相关逻辑
• 预先static_build后的输出数据类型和Place的比对逻辑

不能复用的逻辑为构建子图并进行预静态选Kernel的逻辑。因为While算子和ConditionalBlock算子的子图构建逻辑不一样。我们在 RunWhileBlockPreStaticBuild 中实现了While算子的子图构建逻辑。并在添加如下的选择逻辑区分了While算子和ConditionalBlock算子的子图构建逻辑：

// ConditionalBlock算子预选Kernel的相关逻辑
if (op_type == "conditional_block") {
    RunConditionalBlockPreStaticBuild(*scope, place, op);
} else {
    // While算子预选Kernel的相关逻辑
    RunWhileBlockPreStaticBuild(*scope, place, op);
}

相关PR：

• Support control flow for static build [Step 3: support while] #57616

3.2.4 从黑名单中移除batch_norm

在static_build工作的前期，我们需要将一些没有注册信息的算子加入黑名单，从而关闭这些算子的静态选Kernel。在前期工作完成后，我们可以将这些算子从黑名单中移除。其中 batch_norm 算子是其中一个。

在将 batch_norm 从黑名单中移除后，我们发现 batch_norm 算子在静态选Kernel的情况下会出现类型推导的错误。我们需要根据输入的数据类型来推导输出的数据类型。除了 batch_norm 算子以外，还有一些其他的算子的输出类型注册错误导致单侧报错：

• quantize_linear_kernel： 算子未迁移至PHI下，并需补全注册类型
• flatten_contiguous_range： 中有UNUSED的输出在静态选Kernel时被推导为错误的Place

针对第一个问题，我们需要将 quantize_linear_kernel 算子迁移至PHI下，并补全注册类型。针对第二个问题，我们需要在static_build.cc中添加相关的跳过逻辑。

相关PR:

• remove batch norm from StaticBuildBlackList #57510

3.2.5 小结

在完成了上述的工作后，我们便可以将静态选Kernel的开关设置为默认开启。静态选Kernel的开启可以进一步提升Paddle静态图执行器的性能。

汇报视频：

• https://www.bilibili.com/video/BV1UN41187n9
• https://www.bilibili.com/video/BV1yC4y1R7wS

[AndSonder]

gitlink 夏令营项目已结束，本项目进度记录 issue 关闭 🎉🎉🎉